mirror of
https://github.com/lionsoul2014/ip2region.git
synced 2025-12-08 19:25:22 +00:00
121 lines
4.0 KiB
Python
121 lines
4.0 KiB
Python
# Copyright 2022 The Ip2Region Authors. All rights reserved.
|
|
# Use of this source code is governed by a Apache2.0-style
|
|
# license that can be found in the LICENSE file.
|
|
#
|
|
# Created by luckydog on 2022/7/1.
|
|
# Copyright © 2022年 luckydog. All rights reserved.
|
|
#
|
|
# from ast import main
|
|
import io
|
|
|
|
from xdbSearcher import XdbSearcher
|
|
import argparse
|
|
import time
|
|
import sys
|
|
|
|
|
|
def printHelp():
|
|
print("python bench_test.py [command options]")
|
|
print("options: ")
|
|
print(" --db string ip2region binary xdb file path")
|
|
print(" --src string source ip text file path")
|
|
print(" --cache-policy string cache policy: file/vectorIndex/content")
|
|
|
|
def trim(string):
|
|
if string[:1] != ' ' and string[-1:] != ' ':
|
|
return string
|
|
elif string[:1] == ' ':
|
|
return trim(string[1:])
|
|
else:
|
|
return trim(string[:-1])
|
|
|
|
def start_bench(dbFile="", srcFile="", cachePolicy="vectorIndex"):
|
|
if cachePolicy == "file":
|
|
try:
|
|
searcher = XdbSearcher(dbfile=dbFile)
|
|
except Exception as err:
|
|
print(err)
|
|
return
|
|
elif cachePolicy == "vectorIndex":
|
|
try:
|
|
vi = XdbSearcher.loadVectorIndexFromFile(dbfile=dbFile)
|
|
if vi is None:
|
|
print(f"failed to load vector index from {dbFile}\n")
|
|
searcher = XdbSearcher(dbfile=dbFile, vectorIndex=vi)
|
|
|
|
except Exception as err:
|
|
print(err)
|
|
return
|
|
else:
|
|
try:
|
|
cb = XdbSearcher.loadContentFromFile(dbfile=dbFile)
|
|
if cb is None:
|
|
print(f"failed to load xdb content from {dbFile}\n")
|
|
searcher = XdbSearcher(contentBuff=cb)
|
|
|
|
except Exception as err:
|
|
print(err)
|
|
return
|
|
# do the bench test
|
|
|
|
try:
|
|
count = 0
|
|
costs = 0
|
|
sTime = time.time()
|
|
f = io.open(srcFile, "rb")
|
|
while True:
|
|
line = trim(f.readline(1024)).decode("utf-8").replace("\n", "")
|
|
if len(line) < 1:
|
|
break
|
|
|
|
ps = line.split("|",2)
|
|
if len(ps) != 3:
|
|
print(f"invalid ip segment line :{line}")
|
|
return
|
|
sip = XdbSearcher.ip2long(None, ps[0])
|
|
eip = XdbSearcher.ip2long(None, ps[1])
|
|
|
|
if sip > eip:
|
|
print(f"start ip({ps[0]}) should not be greater than end ip({ps[1]})")
|
|
return
|
|
|
|
mip = (sip + eip) >> 1
|
|
|
|
for ip in [sip, (sip + mip) >> 1, mip, (mip + eip) >> 1, eip]:
|
|
try:
|
|
cTime = time.time()
|
|
region = searcher.search(ip)
|
|
costs = costs + (time.time() - cTime)
|
|
except Exception as error:
|
|
print(f"failed to search ip :{ip}")
|
|
return
|
|
|
|
if region is None:
|
|
print(f"failed to search ip :{ip}")
|
|
return
|
|
if region != ps[2]:
|
|
print(f"failed search({ip}) with ({region} != {ps[2]})")
|
|
return
|
|
count = count + 1
|
|
|
|
# close the searcher at last
|
|
f.close()
|
|
searcher.close()
|
|
print(f"Bench finished, {{cachePolicy: {cachePolicy}, total: {count}, took: {round(time.time() - sTime, 2)} s, cost: {round(costs/count*1000, 4)} ms/op}}")
|
|
except Exception as err:
|
|
print(f"failed to open source text file :{err}")
|
|
return
|
|
|
|
|
|
if __name__ == '__main__':
|
|
if len(sys.argv) < 2:
|
|
printHelp()
|
|
exit(0)
|
|
parse = argparse.ArgumentParser()
|
|
parse.add_argument("--db", help="ip2region binary xdb file path")
|
|
parse.add_argument("--src", help="source ip text file path")
|
|
parse.add_argument("--cache-policy", choices=["file", "vectorIndex", "content"],
|
|
help="cache policy: file/vectorIndex/content")
|
|
args = parse.parse_args()
|
|
start_bench(dbFile=args.db, srcFile=args.src, cachePolicy=args.cache_policy)
|