优化IP名单同步速度

This commit is contained in:
GoEdgeLab
2023-04-19 12:01:02 +08:00
parent 30c50f356c
commit 50838e8688
20 changed files with 334 additions and 56 deletions

View File

@@ -1,6 +1,5 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
//go:build linux
// +build linux
package nftables

View File

@@ -1,4 +1,5 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
//go:build linux
package nftables

View File

@@ -1,6 +1,5 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
//go:build linux
// +build linux
package nftables_test
@@ -11,7 +10,10 @@ import (
)
func getIPv4Chain(t *testing.T) *nftables.Chain {
var conn = nftables.NewConn()
conn, err := nftables.NewConn()
if err != nil {
t.Fatal(err)
}
table, err := conn.GetTable("test_ipv4", nftables.TableFamilyIPv4)
if err != nil {
if err == nftables.ErrTableNotFound {

View File

@@ -15,10 +15,14 @@ type Conn struct {
rawConn *nft.Conn
}
func NewConn() *Conn {
return &Conn{
rawConn: &nft.Conn{},
func NewConn() (*Conn, error) {
conn, err := nft.New()
if err != nil {
return nil, err
}
return &Conn{
rawConn: conn,
}, nil
}
func (this *Conn) Raw() *nft.Conn {

View File

@@ -4,7 +4,10 @@
package nftables
import "errors"
import (
"errors"
"strings"
)
var ErrTableNotFound = errors.New("table not found")
var ErrChainNotFound = errors.New("chain not found")
@@ -15,5 +18,5 @@ func IsNotFound(err error) bool {
if err == nil {
return false
}
return err == ErrTableNotFound || err == ErrChainNotFound || err == ErrSetNotFound || err == ErrRuleNotFound
return err == ErrTableNotFound || err == ErrChainNotFound || err == ErrSetNotFound || err == ErrRuleNotFound || strings.Contains(err.Error(), "no such file or directory")
}

View File

@@ -0,0 +1,62 @@
// Copyright 2023 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
package nftables
import (
"sync"
"time"
)
type Expiration struct {
m map[string]time.Time // key => expires time
lastGCAt int64
locker sync.RWMutex
}
func NewExpiration() *Expiration {
return &Expiration{
m: map[string]time.Time{},
}
}
func (this *Expiration) AddUnsafe(key []byte, expires time.Time) {
this.m[string(key)] = expires
}
func (this *Expiration) Add(key []byte, expires time.Time) {
this.locker.Lock()
this.m[string(key)] = expires
this.gc()
this.locker.Unlock()
}
func (this *Expiration) Remove(key []byte) {
this.locker.Lock()
delete(this.m, string(key))
this.locker.Unlock()
}
func (this *Expiration) Contains(key []byte) bool {
this.locker.RLock()
_, ok := this.m[string(key)]
this.locker.RUnlock()
return ok
}
func (this *Expiration) gc() {
// we won't gc too frequently
var currentTime = time.Now().Unix()
if this.lastGCAt >= currentTime {
return
}
this.lastGCAt = currentTime
var now = time.Now().Add(-10 * time.Second) // gc elements expired before 10 seconds ago
for key, expires := range this.m {
if expires.Year() >= 2000 && now.After(expires) {
delete(this.m, key)
}
}
}

View File

@@ -0,0 +1,51 @@
// Copyright 2023 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
package nftables_test
import (
"github.com/TeaOSLab/EdgeNode/internal/firewalls/nftables"
"github.com/iwind/TeaGo/rands"
"github.com/iwind/TeaGo/types"
"net"
"testing"
"time"
)
func TestExpiration_Add(t *testing.T) {
var expiration = nftables.NewExpiration()
{
expiration.Add([]byte{'a', 'b', 'c'}, time.Now())
t.Log(expiration.Contains([]byte{'a', 'b', 'c'}))
}
{
expiration.Add([]byte{'a', 'b', 'c'}, time.Now().Add(-1*time.Second))
t.Log(expiration.Contains([]byte{'a', 'b', 'c'}))
}
{
expiration.Add([]byte{'a', 'b', 'c'}, time.Now().Add(-10*time.Second))
t.Log(expiration.Contains([]byte{'a', 'b', 'c'}))
}
{
expiration.Add([]byte{'a', 'b', 'c'}, time.Now().Add(1*time.Second))
expiration.Remove([]byte{'a', 'b', 'c'})
t.Log(expiration.Contains([]byte{'a', 'b', 'c'}))
}
{
expiration.Add(net.ParseIP("10.254.0.75").To4(), time.Now())
t.Log(expiration.Contains(net.ParseIP("10.254.0.75").To4()))
}
}
func BenchmarkNewExpiration(b *testing.B) {
var expiration = nftables.NewExpiration()
for i := 0; i < 10_000; i++ {
expiration.Add([]byte(types.String(types.String(rands.Int(0, 255))+"."+types.String(rands.Int(0, 255))+"."+types.String(rands.Int(0, 255))+"."+types.String(rands.Int(0, 255)))), time.Now().Add(3600*time.Second))
}
b.ResetTimer()
b.RunParallel(func(pb *testing.PB) {
for pb.Next() {
expiration.Add([]byte(types.String(types.String(rands.Int(0, 255))+"."+types.String(rands.Int(0, 255))+"."+types.String(rands.Int(0, 255))+"."+types.String(rands.Int(0, 255)))), time.Now().Add(3600*time.Second))
}
})
}

View File

@@ -1,4 +1,5 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
//go:build linux
package nftables

View File

@@ -1,4 +1,5 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
//go:build linux
package nftables

View File

@@ -1,4 +1,5 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
//go:build linux
package nftables

View File

@@ -1,6 +1,5 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
//go:build linux
// +build linux
package nftables
@@ -35,17 +34,25 @@ type Set struct {
conn *Conn
rawSet *nft.Set
batch *SetBatch
expiration *Expiration
}
func NewSet(conn *Conn, rawSet *nft.Set) *Set {
return &Set{
conn: conn,
rawSet: rawSet,
var set = &Set{
conn: conn,
rawSet: rawSet,
expiration: nil,
batch: &SetBatch{
conn: conn,
rawSet: rawSet,
},
}
// retrieve set elements to improve "delete" speed
set.initElements()
return set
}
func (this *Set) Raw() *nft.Set {
@@ -57,11 +64,21 @@ func (this *Set) Name() string {
}
func (this *Set) AddElement(key []byte, options *ElementOptions, overwrite bool) error {
// check if already exists
if this.expiration != nil && !overwrite && this.expiration.Contains(key) {
return nil
}
var expiresTime = time.Time{}
var rawElement = nft.SetElement{
Key: key,
}
if options != nil {
rawElement.Timeout = options.Timeout
if options.Timeout > 0 {
expiresTime = time.UnixMilli(time.Now().UnixMilli() + options.Timeout.Milliseconds())
}
}
err := this.conn.Raw().SetAddElements(this.rawSet, []nft.SetElement{
rawElement,
@@ -71,9 +88,19 @@ func (this *Set) AddElement(key []byte, options *ElementOptions, overwrite bool)
}
err = this.conn.Commit()
if err != nil {
if err == nil {
if this.expiration != nil {
this.expiration.Add(key, expiresTime)
}
} else {
var isFileExistsErr = strings.Contains(err.Error(), "file exists")
if !overwrite && isFileExistsErr {
// ignore file exists error
return nil
}
// retry if exists
if overwrite && strings.Contains(err.Error(), "file exists") {
if overwrite && isFileExistsErr {
deleteErr := this.conn.Raw().SetDeleteElements(this.rawSet, []nft.SetElement{
{
Key: key,
@@ -85,6 +112,11 @@ func (this *Set) AddElement(key []byte, options *ElementOptions, overwrite bool)
})
if err == nil {
err = this.conn.Commit()
if err == nil {
if this.expiration != nil {
this.expiration.Add(key, expiresTime)
}
}
}
}
}
@@ -107,6 +139,11 @@ func (this *Set) AddIPElement(ip string, options *ElementOptions, overwrite bool
}
func (this *Set) DeleteElement(key []byte) error {
// if set element does not exist, we return immediately
if this.expiration != nil && !this.expiration.Contains(key) {
return nil
}
err := this.conn.Raw().SetDeleteElements(this.rawSet, []nft.SetElement{
{
Key: key,
@@ -116,9 +153,17 @@ func (this *Set) DeleteElement(key []byte) error {
return err
}
err = this.conn.Commit()
if err != nil {
if err == nil {
if this.expiration != nil {
this.expiration.Remove(key)
}
} else {
if strings.Contains(err.Error(), "no such file or directory") {
err = nil
if this.expiration != nil {
this.expiration.Remove(key)
}
}
}
return err

View File

@@ -1,4 +1,5 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
//go:build linux
package nftables

View File

@@ -0,0 +1,8 @@
// Copyright 2023 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
//go:build linux && !plus
package nftables
func (this *Set) initElements() {
// NOT IMPLEMENTED
}

View File

@@ -34,7 +34,7 @@ func getIPv4Set(t *testing.T) *nftables.Set {
func TestSet_AddElement(t *testing.T) {
var set = getIPv4Set(t)
err := set.AddElement(net.ParseIP("192.168.2.31").To4(), &nftables.ElementOptions{Timeout: 86400 * time.Second})
err := set.AddElement(net.ParseIP("192.168.2.31").To4(), &nftables.ElementOptions{Timeout: 86400 * time.Second}, false)
if err != nil {
t.Fatal(err)
}

View File

@@ -1,6 +1,5 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
//go:build linux
// +build linux
package nftables

View File

@@ -1,6 +1,5 @@
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
//go:build linux
// +build linux
package nftables_test
@@ -10,7 +9,10 @@ import (
)
func getIPv4Table(t *testing.T) *nftables.Table {
var conn = nftables.NewConn()
conn, err := nftables.NewConn()
if err != nil {
t.Fatal(err)
}
table, err := conn.GetTable("test_ipv4", nftables.TableFamilyIPv4)
if err != nil {
if err == nftables.ErrTableNotFound {