使用 Redis 有序集合實現 IP 歸屬地查詢

yongxinz發表於2019-04-15
Redis

工作中經常遇到一類需求，根據 IP 地址段來查詢 IP 對應的歸屬地資訊。如果把查詢過程放到關係型資料庫中，會帶來很大的 IO 消耗，速度也不能滿足，顯然是不合適的。

那有哪些更好的辦法呢？為此做了一些嘗試，下面來詳細說明。

構建索引檔案

在 GitHub 上看到一個 ip2region 專案，作者通過生成一個包含有二級索引的檔案來實現快速查詢，查詢速度足夠快，毫秒級別。但如果想更新地址段或歸屬地資訊，每次都要重新生成檔案，並不是很方便。

不過還是推薦大家看看這個專案，其中建索引的思想還是很值得學習的。作者的開源專案中只有查詢的相關程式碼，並沒有生成索引檔案的程式碼，我依照原理圖寫了一段生成索引檔案的程式碼，如下：

# -*- coding:utf-8 -*-


import time
import socket
import struct

IP_REGION_FILE = './data/ip_to_region.db'

SUPER_BLOCK_LENGTH = 8
INDEX_BLOCK_LENGTH = 12
HEADER_INDEX_LENGTH = 8192


def generate_db_file():
    pointer = SUPER_BLOCK_LENGTH + HEADER_INDEX_LENGTH

    region, index = '', ''

    # 檔案格式
    # 1.0.0.0|1.0.0.255|澳大利亞|0|0|0|0
    # 1.0.1.0|1.0.3.255|中國|0|福建省|福州市|電信
    with open('./ip.merge.txt', 'r') as f:
        for line in f.readlines():
            item = line.strip().split('|')
            print item[0], item[1], item[2], item[3], item[4], item[5], item[6]
            start_ip = struct.pack('I', struct.unpack('!L', socket.inet_aton(item[0]))[0])
            end_ip = struct.pack('I', struct.unpack('!L', socket.inet_aton(item[1]))[0])
            region_item = '|'.join([item[2], item[3], item[4], item[5], item[6]])
            region += region_item

            ptr = struct.pack('I', int(bin(len(region_item))[2:].zfill(8) + bin(pointer)[2:].zfill(24), 2))
            index += start_ip + end_ip + ptr
            pointer += len(region_item)

    index_start_ptr = pointer
    index_end_ptr = pointer + len(index) - 12
    super_block = struct.pack('I', index_start_ptr) + struct.pack('I', index_end_ptr)

    n = 0
    header_index = ''
    for index_block in range(pointer, index_end_ptr, 8184):
        header_index_block_ip = index[n * 8184:n * 8184 + 4]
        header_index_block_ptr = index_block
        header_index += header_index_block_ip + struct.pack('I', header_index_block_ptr)

        n += 1

    header_index += index[len(index) - 12: len(index) - 8] + struct.pack('I', index_end_ptr)

    with open(IP_REGION_FILE, 'wb') as f:
        f.write(super_block)
        f.write(header_index)
        f.seek(SUPER_BLOCK_LENGTH + HEADER_INDEX_LENGTH, 0)
        f.write(region)
        f.write(index)


if __name__ == '__main__':
    start_time = time.time()
    generate_db_file()

    print 'cost time: ', time.time() - start_time
複製程式碼

使用 Redis 快取

目前有兩種方式對 IP 以及歸屬地資訊進行快取：

第一種是將起始 IP，結束 IP 以及中間所有 IP 轉換成整型，然後以字串方式，用轉換後的 IP 作為 key，歸屬地資訊作為 value 存入 Redis；

第二種是採用有序集合和雜湊方式，首先將起始 IP 和結束 IP 新增到有序集合 ip2cityid，城市 ID 作為成員，轉換後的 IP 作為分值，然後再將城市 ID 和歸屬地資訊新增到雜湊 cityid2city，城市 ID 作為 key，歸屬地資訊作為 value。

第一種方式就不多做介紹了，簡單粗暴，非常不推薦。查詢速度當然很快，毫秒級別，但缺點也十分明顯，我用 1000 條資料做了測試，快取時間長，大概 20 分鐘，佔用空間大，將近 1G。

下面介紹第二種方式，直接看程式碼：

# generate_to_redis.py
# -*- coding:utf-8 -*-

import time
import json
from redis import Redis


def ip_to_num(x):
    return sum([256 ** j * int(i) for j, i in enumerate(x.split('.')[::-1])])


# 連線 Redis
conn = Redis(host='127.0.0.1', port=6379, db=10)

start_time = time.time()

# 檔案格式
# 1.0.0.0|1.0.0.255|澳大利亞|0|0|0|0
# 1.0.1.0|1.0.3.255|中國|0|福建省|福州市|電信
with open('./ip.merge.txt', 'r') as f:
    i = 1
    for line in f.readlines():
        item = line.strip().split('|')
        # 將起始 IP 和結束 IP 新增到有序集合 ip2cityid
        # 成員分別是城市 ID 和 ID + #, 分值是根據 IP 計算的整數值
        conn.zadd('ip2cityid', str(i), ip_to_num(item[0]), str(i) + '#', ip_to_num(item[1]) + 1)
        # 將城市資訊新增到雜湊 cityid2city，key 是城市 ID，值是城市資訊的 json 序列
        conn.hset('cityid2city', str(i), json.dumps([item[2], item[3], item[4], item[5]]))

        i += 1

end_time = time.time()

print 'start_time: ' + str(start_time) + ', end_time: ' + str(end_time) + ', cost time: ' + str(end_time - start_time)
複製程式碼
# test.py
# -*- coding:utf-8 -*-

import sys
import time
import json
import socket
import struct
from redis import Redis

# 連線 Redis
conn = Redis(host='127.0.0.1', port=6379, db=10)

# 將 IP 轉換成整數
ip = struct.unpack("!L", socket.inet_aton(sys.argv[1]))[0]

start_time = time.time()
# 將有序集合從大到小排序，取小於輸入 IP 值的第一條資料
cityid = conn.zrevrangebyscore('ip2cityid', ip, 0, start=0, num=1)
# 如果返回 cityid 是空，或者匹配到了 # 號，說明沒有找到對應地址段
if not cityid or cityid[0].endswith('#'):
    print 'no city info...'
else:
    # 根據城市 ID 到雜湊表取出城市資訊
    ret = json.loads(conn.hget('cityid2city', cityid[0]))
    print ret[0], ret[1], ret[2]

end_time = time.time()
print 'start_time: ' + str(start_time) + ', end_time: ' + str(end_time) + ', cost time: ' + str(end_time - start_time)
複製程式碼
# python generate_to_redis.py 
start_time: 1554300310.31, end_time: 1554300425.65, cost time: 115.333260059
複製程式碼
# python test_2.py 1.0.16.0
日本 0 0
start_time: 1555081532.44, end_time: 1555081532.45, cost time: 0.000912189483643
複製程式碼

測試資料大概 50 萬條，快取所用時間不到 2 分鐘，佔用記憶體 182M，查詢速度毫秒級別。顯而易見，這種方式更值得嘗試。

zrevrangebyscore 方法的時間複雜度是 O(log(N)+M)， N 為有序集的基數， M 為結果集的基數。可見當 N 的值越大，查詢效率越慢，具體在多大的資料量還可以高效查詢，這個有待驗證。不過這個問題我覺得並不用擔心，遇到了再說吧。

以上。

GitHub 求 star： 我的技術部落格

相關文章