大幅优化IP名单查询速度

This commit is contained in:
GoEdgeLab
2021-10-04 17:42:38 +08:00
parent 19b373c18c
commit 74e1cc392d
8 changed files with 305 additions and 162 deletions

View File

@@ -10,7 +10,7 @@ const (
IPItemTypeAll IPItemType = "all" // 所有IP
)
// IP条目
// IPItem IP条目
type IPItem struct {
Type string `json:"type"`
Id int64 `json:"id"`
@@ -20,7 +20,7 @@ type IPItem struct {
EventLevel string `json:"eventLevel"`
}
// 检查是否包含某个IP
// Contains 检查是否包含某个IP
func (this *IPItem) Contains(ip uint64) bool {
switch this.Type {
case IPItemTypeIPv4:

View File

@@ -3,6 +3,7 @@ package iplibrary
import (
"github.com/TeaOSLab/EdgeNode/internal/utils"
"github.com/iwind/TeaGo/assert"
"runtime"
"testing"
"time"
)
@@ -72,3 +73,36 @@ func TestIPItem_Contains(t *testing.T) {
a.IsTrue(item.Contains(utils.IP2Long("192.168.1.1")))
}
}
func TestIPItem_Memory(t *testing.T) {
var list = NewIPList()
for i := 0; i < 2_000_000; i ++ {
list.Add(&IPItem{
Type: "ip",
Id: int64(i),
IPFrom: utils.IP2Long("192.168.1.1"),
IPTo: 0,
ExpiredAt: time.Now().Unix(),
EventLevel: "",
})
}
t.Log("waiting")
time.Sleep(10 * time.Second)
}
func BenchmarkIPItem_Contains(b *testing.B) {
runtime.GOMAXPROCS(1)
item := &IPItem{
IPFrom: utils.IP2Long("192.168.1.1"),
IPTo: utils.IP2Long("192.168.1.101"),
ExpiredAt: 0,
}
ip := utils.IP2Long("192.168.1.1")
for i := 0; i < b.N; i++ {
for j := 0; j < 10_000; j++ {
item.Contains(ip)
}
}
}

View File

@@ -3,24 +3,26 @@ package iplibrary
import (
"github.com/TeaOSLab/EdgeNode/internal/utils"
"github.com/TeaOSLab/EdgeNode/internal/utils/expires"
"sort"
"sync"
)
// IPList IP名单
// TODO IP名单可以分片关闭这样让每一片的数据量减少查询更快
type IPList struct {
itemsMap map[int64]*IPItem // id => item
ipMap map[uint64][]int64 // ip => itemIds
expireList *expires.List
itemsMap map[int64]*IPItem // id => item
sortedItems []*IPItem
allItemsMap map[int64]*IPItem // id => item
isAll bool
expireList *expires.List
locker sync.RWMutex
}
func NewIPList() *IPList {
list := &IPList{
itemsMap: map[int64]*IPItem{},
ipMap: map[uint64][]int64{},
itemsMap: map[int64]*IPItem{},
allItemsMap: map[int64]*IPItem{},
}
expireList := expires.NewList()
@@ -34,14 +36,94 @@ func NewIPList() *IPList {
}
func (this *IPList) Add(item *IPItem) {
this.addItem(item, true)
}
// AddDelay 延迟添加需要手工调用Sort()函数
func (this *IPList) AddDelay(item *IPItem) {
this.addItem(item, false)
}
func (this *IPList) Sort() {
this.locker.Lock()
this.sortItems()
this.locker.Unlock()
}
func (this *IPList) Delete(itemId int64) {
this.locker.Lock()
this.deleteItem(itemId)
this.locker.Unlock()
}
// Contains 判断是否包含某个IP
func (this *IPList) Contains(ip uint64) bool {
this.locker.RLock()
if len(this.allItemsMap) > 0 {
this.locker.RUnlock()
return true
}
var item = this.lookupIP(ip)
this.locker.RUnlock()
return item != nil
}
// ContainsIPStrings 是否包含一组IP中的任意一个并返回匹配的第一个Item
func (this *IPList) ContainsIPStrings(ipStrings []string) (item *IPItem, found bool) {
if len(ipStrings) == 0 {
return
}
this.locker.RLock()
if len(this.allItemsMap) > 0 {
for _, allItem := range this.allItemsMap {
item = allItem
break
}
if item != nil {
this.locker.RUnlock()
found = true
return
}
}
for _, ipString := range ipStrings {
if len(ipString) == 0 {
continue
}
item = this.lookupIP(utils.IP2Long(ipString))
if item != nil {
this.locker.RUnlock()
found = true
return
}
}
this.locker.RUnlock()
return
}
func (this *IPList) addItem(item *IPItem, sortable bool) {
if item == nil {
return
}
if item.ExpiredAt > 0 && item.ExpiredAt < utils.UnixTime() {
return
}
if item.IPFrom == 0 && item.IPTo == 0 {
if item.Type != "all" {
if item.Type != IPItemTypeAll {
return
}
} else if item.IPTo > 0 {
if item.IPFrom > item.IPTo {
item.IPFrom, item.IPTo = item.IPTo, item.IPFrom
} else if item.IPFrom == 0 {
item.IPFrom = item.IPTo
item.IPTo = 0
}
}
this.locker.Lock()
@@ -56,157 +138,86 @@ func (this *IPList) Add(item *IPItem) {
// 展开
if item.IPFrom > 0 {
if item.IPTo == 0 {
this.addIP(item.IPFrom, item.Id)
} else {
if item.IPFrom > item.IPTo {
item.IPTo, item.IPFrom = item.IPFrom, item.IPTo
}
for i := item.IPFrom; i <= item.IPTo; i++ {
// 最多不能超过65535防止整个系统内存爆掉
if i >= item.IPFrom+65535 {
break
}
this.addIP(i, item.Id)
}
}
} else if item.IPTo > 0 {
this.addIP(item.IPTo, item.Id)
this.sortedItems = append(this.sortedItems, item)
} else {
this.addIP(0, item.Id)
// 更新isAll
this.isAll = true
this.allItemsMap[item.Id] = item
}
if item.ExpiredAt > 0 {
this.expireList.Add(item.Id, item.ExpiredAt)
}
if sortable {
this.sortItems()
}
this.locker.Unlock()
}
func (this *IPList) Delete(itemId int64) {
this.locker.Lock()
defer this.locker.Unlock()
this.deleteItem(itemId)
// 更新isAll
this.isAll = len(this.ipMap[0]) > 0
// 对列表进行排序
func (this *IPList) sortItems() {
sort.Slice(this.sortedItems, func(i, j int) bool {
var item1 = this.sortedItems[i]
var item2 = this.sortedItems[j]
if item1.IPFrom == item2.IPFrom {
return item1.IPTo < item2.IPTo
}
return item1.IPFrom < item2.IPFrom
})
}
// Contains 判断是否包含某个IP
func (this *IPList) Contains(ip uint64) bool {
this.locker.RLock()
if this.isAll {
this.locker.RUnlock()
return true
}
_, ok := this.ipMap[ip]
this.locker.RUnlock()
return ok
}
// ContainsIPStrings 是否包含一组IP
func (this *IPList) ContainsIPStrings(ipStrings []string) (found bool, item *IPItem) {
if len(ipStrings) == 0 {
return
}
this.locker.RLock()
if this.isAll {
itemIds := this.ipMap[0]
if len(itemIds) > 0 {
itemId := itemIds[0]
item = this.itemsMap[itemId]
}
this.locker.RUnlock()
found = true
return
}
for _, ipString := range ipStrings {
if len(ipString) == 0 {
continue
}
itemIds, ok := this.ipMap[utils.IP2Long(ipString)]
if ok {
if len(itemIds) > 0 {
itemId := itemIds[0]
item = this.itemsMap[itemId]
// 不加锁的情况下查找Item
func (this *IPList) lookupIP(ip uint64) *IPItem {
var count = len(this.sortedItems)
var resultIndex = -1
sort.Search(count, func(i int) bool {
var item = this.sortedItems[i]
if item.IPFrom < ip {
if item.IPTo >= ip {
resultIndex = i
}
this.locker.RUnlock()
found = true
return
return false
} else if item.IPFrom == ip {
resultIndex = i
return false
}
return true
})
if resultIndex < 0 || resultIndex >= count {
return nil
}
this.locker.RUnlock()
return
return this.sortedItems[resultIndex]
}
// 在不加锁的情况下删除某个Item
// 将会被别的方法引用,切记不能加锁
func (this *IPList) deleteItem(itemId int64) {
item, ok := this.itemsMap[itemId]
_, ok := this.itemsMap[itemId]
if !ok {
return
}
delete(this.itemsMap, itemId)
// 展开
if item.IPFrom > 0 {
if item.IPTo == 0 {
this.deleteIP(item.IPFrom, item.Id)
} else {
if item.IPFrom > item.IPTo {
item.IPTo, item.IPFrom = item.IPFrom, item.IPTo
}
for i := item.IPFrom; i <= item.IPTo; i++ {
// 最多不能超过65535防止整个系统内存爆掉
if i >= item.IPFrom+65535 {
break
}
this.deleteIP(i, item.Id)
}
}
} else if item.IPTo > 0 {
this.deleteIP(item.IPTo, item.Id)
} else {
this.deleteIP(0, item.Id)
}
}
// 添加单个IP
func (this *IPList) addIP(ip uint64, itemId int64) {
itemIds, ok := this.ipMap[ip]
// 是否为All Item
_, ok = this.allItemsMap[itemId]
if ok {
itemIds = append(itemIds, itemId)
} else {
itemIds = []int64{itemId}
}
this.ipMap[ip] = itemIds
}
// 删除单个IP
func (this *IPList) deleteIP(ip uint64, itemId int64) {
itemIds, ok := this.ipMap[ip]
if !ok {
delete(this.allItemsMap, itemId)
return
}
newItemIds := []int64{}
for _, oldItemId := range itemIds {
if oldItemId == itemId {
continue
// 删除排序中的Item
var index = -1
for itemIndex, item := range this.sortedItems {
if item.Id == itemId {
index = itemIndex
break
}
newItemIds = append(newItemIds, oldItemId)
}
if len(newItemIds) > 0 {
this.ipMap[ip] = newItemIds
} else {
delete(this.ipMap, ip)
if index >= 0 {
copy(this.sortedItems[index:], this.sortedItems[index+1:])
this.sortedItems = this.sortedItems[:len(this.sortedItems)-1]
}
}

View File

@@ -4,6 +4,7 @@ import (
"github.com/TeaOSLab/EdgeNode/internal/utils"
"github.com/iwind/TeaGo/assert"
"github.com/iwind/TeaGo/logs"
"github.com/iwind/TeaGo/rands"
"runtime"
"strconv"
"testing"
@@ -16,7 +17,7 @@ func TestIPList_Add_Empty(t *testing.T) {
Id: 1,
})
logs.PrintAsJSON(ipList.itemsMap, t)
logs.PrintAsJSON(ipList.ipMap, t)
logs.PrintAsJSON(ipList.allItemsMap, t)
}
func TestIPList_Add_One(t *testing.T) {
@@ -31,15 +32,30 @@ func TestIPList_Add_One(t *testing.T) {
})
ipList.Add(&IPItem{
Id: 3,
IPFrom: utils.IP2Long("2001:db8:0:1::101"),
IPFrom: utils.IP2Long("192.168.0.2"),
})
ipList.Add(&IPItem{
Id: 4,
IPFrom: utils.IP2Long("192.168.0.2"),
IPTo: utils.IP2Long("192.168.0.1"),
})
ipList.Add(&IPItem{
Id: 5,
IPFrom: utils.IP2Long("2001:db8:0:1::101"),
})
ipList.Add(&IPItem{
Id: 6,
IPFrom: 0,
Type: "all",
})
t.Log("===items===")
logs.PrintAsJSON(ipList.itemsMap, t)
logs.PrintAsJSON(ipList.ipMap, t) // ip => items
t.Log("===sorted items===")
logs.PrintAsJSON(ipList.sortedItems, t)
t.Log("===all items===")
logs.PrintAsJSON(ipList.allItemsMap, t) // ip => items
}
func TestIPList_Update(t *testing.T) {
@@ -50,14 +66,31 @@ func TestIPList_Update(t *testing.T) {
})
/**ipList.Add(&IPItem{
Id: 2,
IPFrom: IP2Long("192.168.1.1"),
IPFrom: utils.IP2Long("192.168.1.1"),
})**/
ipList.Add(&IPItem{
Id: 1,
IPTo: utils.IP2Long("192.168.1.2"),
})
logs.PrintAsJSON(ipList.itemsMap, t)
logs.PrintAsJSON(ipList.ipMap, t)
logs.PrintAsJSON(ipList.sortedItems, t)
}
func TestIPList_Update_AllItems(t *testing.T) {
ipList := NewIPList()
ipList.Add(&IPItem{
Id: 1,
Type: IPItemTypeAll,
IPFrom: 0,
})
ipList.Add(&IPItem{
Id: 1,
IPTo: 0,
})
t.Log("===items map===")
logs.PrintAsJSON(ipList.itemsMap, t)
t.Log("===all items map===")
logs.PrintAsJSON(ipList.allItemsMap, t)
}
func TestIPList_Add_Range(t *testing.T) {
@@ -71,9 +104,9 @@ func TestIPList_Add_Range(t *testing.T) {
Id: 2,
IPTo: utils.IP2Long("192.168.1.2"),
})
t.Log(len(ipList.ipMap), "ips")
t.Log(len(ipList.itemsMap), "ips")
logs.PrintAsJSON(ipList.itemsMap, t)
logs.PrintAsJSON(ipList.ipMap, t)
logs.PrintAsJSON(ipList.allItemsMap, t)
}
func TestIPList_Add_Overflow(t *testing.T) {
@@ -85,8 +118,8 @@ func TestIPList_Add_Overflow(t *testing.T) {
IPFrom: utils.IP2Long("192.168.1.1"),
IPTo: utils.IP2Long("192.169.255.1"),
})
t.Log(len(ipList.ipMap), "ips")
a.IsTrue(len(ipList.ipMap) <= 65535)
t.Log(len(ipList.itemsMap), "ips")
a.IsTrue(len(ipList.itemsMap) <= 65535)
}
func TestNewIPList_Memory(t *testing.T) {
@@ -104,20 +137,50 @@ func TestNewIPList_Memory(t *testing.T) {
}
func TestIPList_Contains(t *testing.T) {
var a = assert.NewAssertion(t)
list := NewIPList()
for i := 0; i < 255; i++ {
list.Add(&IPItem{
list.AddDelay(&IPItem{
Id: int64(i),
IPFrom: utils.IP2Long(strconv.Itoa(i) + ".168.0.1"),
IPTo: utils.IP2Long(strconv.Itoa(i) + ".168.255.1"),
ExpiredAt: 0,
})
}
t.Log(len(list.ipMap), "ip")
for i := 0; i < 255; i++ {
list.AddDelay(&IPItem{
Id: int64(1000 + i),
IPFrom: utils.IP2Long("192.167.2." + strconv.Itoa(i)),
})
}
list.Sort()
t.Log(len(list.itemsMap), "ip")
before := time.Now()
t.Log(list.Contains(utils.IP2Long("192.168.1.100")))
t.Log(list.Contains(utils.IP2Long("192.168.2.100")))
a.IsTrue(list.Contains(utils.IP2Long("192.168.1.100")))
a.IsTrue(list.Contains(utils.IP2Long("192.168.2.100")))
a.IsFalse(list.Contains(utils.IP2Long("192.169.3.100")))
a.IsFalse(list.Contains(utils.IP2Long("192.167.3.100")))
a.IsTrue(list.Contains(utils.IP2Long("192.167.2.100")))
t.Log(time.Since(before).Seconds()*1000, "ms")
}
func TestIPList_Contains_Many(t *testing.T) {
list := NewIPList()
for i := 0; i < 1_000_000; i++ {
list.AddDelay(&IPItem{
Id: int64(i),
IPFrom: utils.IP2Long(strconv.Itoa(rands.Int(0, 255)) + "." + strconv.Itoa(rands.Int(0, 255)) + "." + strconv.Itoa(rands.Int(0, 255)) + "." + strconv.Itoa(rands.Int(0, 255))),
IPTo: utils.IP2Long(strconv.Itoa(rands.Int(0, 255)) + "." + strconv.Itoa(rands.Int(0, 255)) + "." + strconv.Itoa(rands.Int(0, 255)) + "." + strconv.Itoa(rands.Int(0, 255))),
ExpiredAt: 0,
})
}
list.Sort()
t.Log(len(list.itemsMap), "ip")
before := time.Now()
_ = list.Contains(utils.IP2Long("192.168.1.100"))
t.Log(time.Since(before).Seconds()*1000, "ms")
}
@@ -146,6 +209,32 @@ func TestIPList_ContainsAll(t *testing.T) {
}
func TestIPList_ContainsIPStrings(t *testing.T) {
var a = assert.NewAssertion(t)
list := NewIPList()
for i := 0; i < 255; i++ {
list.Add(&IPItem{
Id: int64(i),
IPFrom: utils.IP2Long(strconv.Itoa(i) + ".168.0.1"),
IPTo: utils.IP2Long(strconv.Itoa(i) + ".168.255.1"),
ExpiredAt: 0,
})
}
t.Log(len(list.itemsMap), "ip")
{
item, ok := list.ContainsIPStrings([]string{"192.168.1.100"})
t.Log("item:", item)
a.IsTrue(ok)
}
{
item, ok := list.ContainsIPStrings([]string{"192.167.1.100"})
t.Log("item:", item)
a.IsFalse(ok)
}
}
func TestIPList_Delete(t *testing.T) {
list := NewIPList()
list.Add(&IPItem{
@@ -160,13 +249,13 @@ func TestIPList_Delete(t *testing.T) {
})
t.Log("===BEFORE===")
logs.PrintAsJSON(list.itemsMap, t)
logs.PrintAsJSON(list.ipMap, t)
logs.PrintAsJSON(list.allItemsMap, t)
list.Delete(1)
t.Log("===AFTER===")
logs.PrintAsJSON(list.itemsMap, t)
logs.PrintAsJSON(list.ipMap, t)
logs.PrintAsJSON(list.allItemsMap, t)
}
func TestGC(t *testing.T) {
@@ -184,27 +273,27 @@ func TestGC(t *testing.T) {
ExpiredAt: 0,
})
logs.PrintAsJSON(list.itemsMap, t)
logs.PrintAsJSON(list.ipMap, t)
logs.PrintAsJSON(list.allItemsMap, t)
time.Sleep(2 * time.Second)
t.Log("===AFTER GC===")
logs.PrintAsJSON(list.itemsMap, t)
logs.PrintAsJSON(list.ipMap, t)
logs.PrintAsJSON(list.sortedItems, t)
}
func BenchmarkIPList_Contains(b *testing.B) {
runtime.GOMAXPROCS(1)
list := NewIPList()
for i := 192; i < 194; i++ {
for i := 1; i < 194; i++ {
list.Add(&IPItem{
Id: int64(1),
IPFrom: utils.IP2Long(strconv.Itoa(i) + ".1.0.1"),
IPTo: utils.IP2Long(strconv.Itoa(i) + ".2.0.1"),
Id: int64(i),
IPFrom: utils.IP2Long(strconv.Itoa(i%255) + "." + strconv.Itoa(i%255) + ".0.1"),
IPTo: utils.IP2Long(strconv.Itoa(i%255) + "." + strconv.Itoa(i%255) + ".0.1"),
ExpiredAt: time.Now().Unix() + 60,
})
}
b.Log(len(list.ipMap), "ip")
b.Log(len(list.itemsMap), "ip")
for i := 0; i < b.N; i++ {
_ = list.Contains(utils.IP2Long("192.168.1.100"))
}

View File

@@ -5,9 +5,9 @@ import (
"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
"github.com/TeaOSLab/EdgeNode/internal/errors"
"github.com/TeaOSLab/EdgeNode/internal/events"
"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
"github.com/iwind/TeaGo/Tea"
"github.com/iwind/TeaGo/files"
"github.com/iwind/TeaGo/logs"
"github.com/iwind/TeaGo/types"
"regexp"
"strings"
@@ -21,7 +21,7 @@ func init() {
// 初始化
library, err := SharedManager.Load()
if err != nil {
logs.Println("[IP_LIBRARY]" + err.Error())
remotelogs.Error("IP_LIBRARY", err.Error())
return
}
SharedLibrary = library

View File

@@ -112,12 +112,16 @@ func (this *IPListManager) fetch() (hasNext bool, err error) {
return false, nil
}
this.locker.Lock()
var changedLists = map[*IPList]bool{}
for _, item := range items {
list, ok := this.listMap[item.ListId]
if !ok {
list = NewIPList()
this.listMap[item.ListId] = list
}
changedLists[list] = true
if item.IsDeleted {
list.Delete(item.Id)
@@ -127,7 +131,7 @@ func (this *IPListManager) fetch() (hasNext bool, err error) {
continue
}
list.Add(&IPItem{
list.AddDelay(&IPItem{
Id: item.Id,
Type: item.Type,
IPFrom: utils.IP2Long(item.IpFrom),
@@ -140,6 +144,11 @@ func (this *IPListManager) fetch() (hasNext bool, err error) {
SharedActionManager.DeleteItem(item.ListType, item)
SharedActionManager.AddItem(item.ListType, item)
}
for changedList := range changedLists {
changedList.Sort()
}
this.locker.Unlock()
this.version = items[len(items)-1].Version

View File

@@ -7,9 +7,9 @@ import (
"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
"github.com/TeaOSLab/EdgeNode/internal/errors"
"github.com/TeaOSLab/EdgeNode/internal/events"
"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
"github.com/TeaOSLab/EdgeNode/internal/rpc"
"github.com/iwind/TeaGo/Tea"
"github.com/iwind/TeaGo/logs"
"os"
"time"
)
@@ -21,16 +21,16 @@ func init() {
})
}
// IP库更新程序
// Updater IP库更新程序
type Updater struct {
}
// 获取新对象
// NewUpdater 获取新对象
func NewUpdater() *Updater {
return &Updater{}
}
// 开始更新
// Start 开始更新
func (this *Updater) Start() {
// 这里不需要太频繁检查更新因为通常不需要更新IP库
ticker := time.NewTicker(1 * time.Hour)
@@ -38,7 +38,7 @@ func (this *Updater) Start() {
for range ticker.C {
err := this.loop()
if err != nil {
logs.Println("[IP_LIBRARY]" + err.Error())
remotelogs.Error("IP_LIBRARY", err.Error())
}
}
}()

View File

@@ -67,7 +67,7 @@ func (this *HTTPRequest) checkWAFRequest(firewallPolicy *firewallconfigs.HTTPFir
if ref.IsOn && ref.ListId > 0 {
list := iplibrary.SharedIPListManager.FindList(ref.ListId)
if list != nil {
found, _ := list.ContainsIPStrings(remoteAddrs)
_, found := list.ContainsIPStrings(remoteAddrs)
if found {
breakChecking = true
return
@@ -81,7 +81,7 @@ func (this *HTTPRequest) checkWAFRequest(firewallPolicy *firewallconfigs.HTTPFir
if ref.IsOn && ref.ListId > 0 {
list := iplibrary.SharedIPListManager.FindList(ref.ListId)
if list != nil {
found, item := list.ContainsIPStrings(remoteAddrs)
item, found := list.ContainsIPStrings(remoteAddrs)
if found {
// 触发事件
if item != nil && len(item.EventLevel) > 0 {