mirror of
https://github.com/TeaOSLab/EdgeNode.git
synced 2026-01-05 23:35:48 +08:00
大幅优化IP名单查询速度
This commit is contained in:
GoEdgeLab
@@ -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]
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user