优化通过IP查询区域性能

internal/iplibrary/ip2Region.go

@@ -0,0 +1,130 @@
+// 源码改自：https://github.com/lionsoul2014/ip2region/blob/master/binding/golang/ip2region/ip2Region.go
+
+package iplibrary
+
+import (
+	"errors"
+	"io/ioutil"
+	"strconv"
+	"strings"
+)
+
+const (
+	IndexBlockLength = 12
+)
+
+var err error
+
+type IP2Region struct {
+	headerSip []int64
+	headerPtr []int64
+	headerLen int64
+
+	// super block index info
+	firstIndexPtr int64
+	lastIndexPtr  int64
+	totalBlocks   int64
+
+	dbData []byte
+}
+
+type IpInfo struct {
+	CityId   int64
+	Country  string
+	Region   string
+	Province string
+	City     string
+	ISP      string
+}
+
+func (ip IpInfo) String() string {
+	return strconv.FormatInt(ip.CityId, 10) + "|" + ip.Country + "|" + ip.Region + "|" + ip.Province + "|" + ip.City + "|" + ip.ISP
+}
+
+func getIpInfo(cityId int64, line []byte) *IpInfo {
+	lineSlice := strings.Split(string(line), "|")
+	ipInfo := &IpInfo{}
+	length := len(lineSlice)
+	ipInfo.CityId = cityId
+	if length < 5 {
+		for i := 0; i <= 5-length; i++ {
+			lineSlice = append(lineSlice, "")
+		}
+	}
+
+	ipInfo.Country = lineSlice[0]
+	ipInfo.Region = lineSlice[1]
+	ipInfo.Province = lineSlice[2]
+	ipInfo.City = lineSlice[3]
+	ipInfo.ISP = lineSlice[4]
+	return ipInfo
+}
+
+func NewIP2Region(path string) (*IP2Region, error) {
+	var region = &IP2Region{}
+	region.dbData, err = ioutil.ReadFile(path)
+
+	if err != nil {
+		return nil, err
+	}
+
+	region.firstIndexPtr = region.ipLongAtOffset(0)
+	region.lastIndexPtr = region.ipLongAtOffset(4)
+	region.totalBlocks = (region.lastIndexPtr-region.firstIndexPtr)/IndexBlockLength + 1
+	return region, nil
+}
+
+func (this *IP2Region) MemorySearch(ipStr string) (ipInfo *IpInfo, err error) {
+	ip, err := ip2long(ipStr)
+	if err != nil {
+		return nil, err
+	}
+
+	h := this.totalBlocks
+	var dataPtr, l int64
+	for l <= h {
+		m := (l + h) >> 1
+		p := this.firstIndexPtr + m*IndexBlockLength
+		sip := this.ipLongAtOffset(p)
+		if ip < sip {
+			h = m - 1
+		} else {
+			eip := this.ipLongAtOffset(p + 4)
+			if ip > eip {
+				l = m + 1
+			} else {
+				dataPtr = this.ipLongAtOffset(p + 8)
+				break
+			}
+		}
+	}
+	if dataPtr == 0 {
+		return nil, errors.New("not found")
+	}
+
+	dataLen := (dataPtr >> 24) & 0xFF
+	dataPtr = dataPtr & 0x00FFFFFF
+	return getIpInfo(this.ipLongAtOffset(dataPtr), this.dbData[(dataPtr)+4:dataPtr+dataLen]), nil
+}
+
+func (this *IP2Region) ipLongAtOffset(offset int64) int64 {
+	return int64(this.dbData[offset]) |
+		int64(this.dbData[offset+1])<<8 |
+		int64(this.dbData[offset+2])<<16 |
+		int64(this.dbData[offset+3])<<24
+}
+
+func ip2long(IpStr string) (int64, error) {
+	bits := strings.Split(IpStr, ".")
+	if len(bits) != 4 {
+		return 0, errors.New("ip format error")
+	}
+
+	var sum int64
+	for i, n := range bits {
+		bit, _ := strconv.ParseInt(n, 10, 64)
+		sum += bit << uint(24-8*i)
+	}
+
+	return sum, nil
+}

internal/iplibrary/library_ip2region.go

@@ -4,17 +4,16 @@ import (
 	"fmt"
 	"github.com/TeaOSLab/EdgeNode/internal/errors"
 	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
-	"github.com/lionsoul2014/ip2region/binding/golang/ip2region"
 	"net"
 	"strings"
 )
 
 type IP2RegionLibrary struct {
-	db *ip2region.Ip2Region
+	db *IP2Region
 }
 
 func (this *IP2RegionLibrary) Load(dbPath string) error {
-	db, err := ip2region.New(dbPath)
+	db, err := NewIP2Region(dbPath)
 	if err != nil {
 		return err
 	}
@@ -49,6 +48,10 @@ func (this *IP2RegionLibrary) Lookup(ip string) (*Result, error) {
 		return nil, err
 	}
 
+	if info == nil {
+		return nil, nil
+	}
+
 	if info.Country == "0" {
 		info.Country = ""
 	}
@@ -76,7 +79,5 @@ func (this *IP2RegionLibrary) Lookup(ip string) (*Result, error) {
 }
 
 func (this *IP2RegionLibrary) Close() {
-	if this.db != nil {
-		this.db.Close()
-	}
+
 }

internal/iplibrary/library_ip2region_test.go

@@ -3,25 +3,83 @@ package iplibrary
 import (
 	"github.com/iwind/TeaGo/Tea"
 	_ "github.com/iwind/TeaGo/bootstrap"
-	"github.com/iwind/TeaGo/logs"
 	"github.com/iwind/TeaGo/rands"
 	"runtime"
 	"strconv"
+	"sync"
 	"testing"
 	"time"
 )
 
+func TestIP2RegionLibrary_Lookup_MemoryUsage(t *testing.T) {
+	var mem = &runtime.MemStats{}
+	runtime.ReadMemStats(mem)
+
+	library := &IP2RegionLibrary{}
+	err := library.Load(Tea.Root + "/resources/ipdata/ip2region/ip2region.db")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var mem2 = &runtime.MemStats{}
+	runtime.ReadMemStats(mem2)
+	t.Log((mem2.HeapInuse-mem.HeapInuse)/1024/1024, "MB")
+}
+
+func TestIP2RegionLibrary_Lookup_Single(t *testing.T) {
+	library := &IP2RegionLibrary{}
+	err := library.Load(Tea.Root + "/resources/ipdata/ip2region/ip2region.db")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	for _, ip := range []string{"8.8.9.9"} {
+		result, err := library.Lookup(ip)
+		if err != nil {
+			t.Fatal(err)
+		}
+		t.Log("IP:", ip, "result:", result)
+	}
+}
+
 func TestIP2RegionLibrary_Lookup(t *testing.T) {
 	library := &IP2RegionLibrary{}
 	err := library.Load(Tea.Root + "/resources/ipdata/ip2region/ip2region.db")
 	if err != nil {
 		t.Fatal(err)
 	}
-	result, err := library.Lookup("114.240.223.47")
+
+	for _, ip := range []string{"", "a", "1.1.1", "192.168.1.100", "114.240.223.47", "8.8.9.9", "::1"} {
+		result, err := library.Lookup(ip)
+		if err != nil {
+			t.Fatal(err)
+		}
+		t.Log("IP:", ip, "result:", result)
+	}
+}
+
+func TestIP2RegionLibrary_Lookup_Concurrent(t *testing.T) {
+	library := &IP2RegionLibrary{}
+	err := library.Load(Tea.Root + "/resources/ipdata/ip2region/ip2region.db")
 	if err != nil {
 		t.Fatal(err)
 	}
-	logs.PrintAsJSON(result, t)
+
+	var count = 4000
+	var wg = sync.WaitGroup{}
+	wg.Add(count)
+	for i := 0; i < count; i++ {
+		go func() {
+			defer wg.Done()
+
+			for i := 0; i < 100; i++ {
+				_, _ = library.Lookup(strconv.Itoa(rands.Int(0, 254)) + "." + strconv.Itoa(rands.Int(0, 254)) + "." + strconv.Itoa(rands.Int(0, 254)) + "." + strconv.Itoa(rands.Int(0, 254)))
+			}
+		}()
+	}
+
+	wg.Done()
+	t.Log("ok")
 }
 
 func TestIP2RegionLibrary_Memory(t *testing.T) {
@@ -43,13 +101,13 @@ func TestIP2RegionLibrary_Memory(t *testing.T) {
 func BenchmarkIP2RegionLibrary_Lookup(b *testing.B) {
 	runtime.GOMAXPROCS(1)
 
-	library := &IP2RegionLibrary{}
+	var library = &IP2RegionLibrary{}
 	err := library.Load(Tea.Root + "/resources/ipdata/ip2region/ip2region.db")
 	if err != nil {
 		b.Fatal(err)
 	}
 
 	for i := 0; i < b.N; i++ {
-		_, _ = library.Lookup(strconv.Itoa(rands.Int(0, 254)) + "." + strconv.Itoa(rands.Int(0, 254)) + "." + strconv.Itoa(rands.Int(0, 254)) + "." + strconv.Itoa(rands.Int(0, 254)))
+		_, _ = library.Lookup("8.8.8.8")
 	}
 }