使用KV数据库来管理IP名单

internal/utils/byte/utils.go

@@ -0,0 +1,28 @@
+// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package byteutils
+
+// Copy bytes
+func Copy(b []byte) []byte {
+	var l = len(b)
+	if l == 0 {
+		return []byte{}
+	}
+	var d = make([]byte, l)
+	copy(d, b)
+	return d
+}
+
+// Append bytes
+func Append(b []byte, b2 ...byte) []byte {
+	return append(Copy(b), b2...)
+}
+
+// Contact bytes
+func Contact(b []byte, b2 ...[]byte) []byte {
+	b = Copy(b)
+	for _, b3 := range b2 {
+		b = append(b, b3...)
+	}
+	return b
+}

internal/utils/byte/utils_test.go

@@ -0,0 +1,56 @@
+// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package byteutils_test
+
+import (
+	"bytes"
+	byteutils "github.com/TeaOSLab/EdgeNode/internal/utils/byte"
+	"github.com/iwind/TeaGo/assert"
+	"testing"
+)
+
+func TestCopy(t *testing.T) {
+	var a = assert.NewAssertion(t)
+
+	var prefix []byte
+	prefix = append(prefix, 1, 2, 3)
+	t.Log(prefix, byteutils.Copy(prefix))
+	a.IsTrue(bytes.Equal(byteutils.Copy(prefix), []byte{1, 2, 3}))
+}
+
+func TestAppend(t *testing.T) {
+	var as = assert.NewAssertion(t)
+
+	var prefix []byte
+	prefix = append(prefix, 1, 2, 3)
+
+	// [1 2 3 4 5 6] [1 2 3 7]
+	var a = byteutils.Append(prefix, 4, 5, 6)
+	var b = byteutils.Append(prefix, 7)
+	t.Log(a, b)
+
+	as.IsTrue(bytes.Equal(a, []byte{1, 2, 3, 4, 5, 6}))
+	as.IsTrue(bytes.Equal(b, []byte{1, 2, 3, 7}))
+}
+
+func TestConcat(t *testing.T) {
+	var a = assert.NewAssertion(t)
+
+	var prefix []byte
+	prefix = append(prefix, 1, 2, 3)
+
+	var b = byteutils.Contact(prefix, []byte{4, 5, 6}, []byte{7})
+	t.Log(b)
+
+	a.IsTrue(bytes.Equal(b, []byte{1, 2, 3, 4, 5, 6, 7}))
+}
+
+func TestAppend_Raw(t *testing.T) {
+	var prefix []byte
+	prefix = append(prefix, 1, 2, 3)
+
+	// [1 2 3 7 5 6] [1 2 3 7]
+	var a = append(prefix, 4, 5, 6)
+	var b = append(prefix, 7)
+	t.Log(a, b)
+}

internal/utils/kvstore/errors.go

@@ -9,10 +9,16 @@ import (
 
 var ErrTableNotFound = errors.New("table not found")
 var ErrKeyTooLong = errors.New("too long key")
+var ErrSkip= errors.New("skip") // skip count in iterator
 
-func IsKeyNotFound(err error) bool {
-	if err == nil {
-		return false
-	}
-	return errors.Is(err, pebble.ErrNotFound)
+func IsNotFound(err error) bool {
+	return err != nil && errors.Is(err, pebble.ErrNotFound)
+}
+
+func IsSkipError(err error) bool {
+	return err != nil && errors.Is(err, ErrSkip)
+}
+
+func Skip() (bool, error) {
+	return true, ErrSkip
 }

internal/utils/kvstore/options.go

@@ -7,3 +7,7 @@ import "github.com/cockroachdb/pebble"
 var DefaultWriteOptions = &pebble.WriteOptions{
 	Sync: false,
 }
+
+var DefaultWriteSyncOptions = &pebble.WriteOptions{
+	Sync: true,
+}

internal/utils/kvstore/query.go

@@ -6,6 +6,7 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
+	byteutils "github.com/TeaOSLab/EdgeNode/internal/utils/byte"
 )
 
 type DataType = int
@@ -222,11 +223,11 @@ func (this *Query[T]) iterateKeys(fn IteratorFunc[T]) error {
 	var prefix []byte
 	switch this.dataType {
 	case DataTypeKey:
-		prefix = append(this.table.Namespace(), KeyPrefix...)
+		prefix = byteutils.Append(this.table.Namespace(), []byte(KeyPrefix)...)
 	case DataTypeField:
-		prefix = append(this.table.Namespace(), FieldPrefix...)
+		prefix = byteutils.Append(this.table.Namespace(), []byte(FieldPrefix)...)
 	default:
-		prefix = append(this.table.Namespace(), KeyPrefix...)
+		prefix = byteutils.Append(this.table.Namespace(), []byte(KeyPrefix)...)
 	}
 
 	var prefixLen = len(prefix)
@@ -238,21 +239,21 @@ func (this *Query[T]) iterateKeys(fn IteratorFunc[T]) error {
 	var offsetKey []byte
 	if this.reverse {
 		if len(this.offsetKey) > 0 {
-			offsetKey = append(prefix, this.offsetKey...)
+			offsetKey = byteutils.Append(prefix, []byte(this.offsetKey)...)
 		} else {
-			offsetKey = append(prefix, 0xFF)
+			offsetKey = byteutils.Append(prefix, 0xFF)
 		}
 
 		opt.LowerBound = prefix
 		opt.UpperBound = offsetKey
 	} else {
 		if len(this.offsetKey) > 0 {
-			offsetKey = append(prefix, this.offsetKey...)
+			offsetKey = byteutils.Append(prefix, []byte(this.offsetKey)...)
 		} else {
 			offsetKey = prefix
 		}
 		opt.LowerBound = offsetKey
-		opt.UpperBound = append(offsetKey, 0xFF)
+		opt.UpperBound = byteutils.Append(prefix, 0xFF)
 	}
 
 	var hasOffsetKey = len(this.offsetKey) > 0
@@ -267,7 +268,7 @@ func (this *Query[T]) iterateKeys(fn IteratorFunc[T]) error {
 
 	var count int
 
-	var itemFn = func() (goNext bool, err error) {
+	var itemFn = func() (goNextItem bool, err error) {
 		var keyBytes = it.Key()
 
 		// skip first offset key
@@ -297,7 +298,11 @@ func (this *Query[T]) iterateKeys(fn IteratorFunc[T]) error {
 			Value: value,
 		})
 		if callbackErr != nil {
-			return false, callbackErr
+			if IsSkipError(callbackErr) {
+				return true, nil
+			} else {
+				return false, callbackErr
+			}
 		}
 		if !goNext {
 			return false, nil
@@ -361,9 +366,9 @@ func (this *Query[T]) iterateFields(fn IteratorFunc[T]) error {
 		if len(this.fieldOffsetKey) > 0 {
 			offsetKey = this.fieldOffsetKey
 		} else if len(this.offsetKey) > 0 {
-			offsetKey = append(prefix, this.offsetKey...)
+			offsetKey = byteutils.Append(prefix, []byte(this.offsetKey)...)
 		} else {
-			offsetKey = append(prefix, 0xFF)
+			offsetKey = byteutils.Append(prefix, 0xFF)
 		}
 		opt.LowerBound = prefix
 		opt.UpperBound = offsetKey
@@ -371,14 +376,14 @@ func (this *Query[T]) iterateFields(fn IteratorFunc[T]) error {
 		if len(this.fieldOffsetKey) > 0 {
 			offsetKey = this.fieldOffsetKey
 		} else if len(this.offsetKey) > 0 {
-			offsetKey = append(prefix, this.offsetKey...)
+			offsetKey = byteutils.Append(prefix, []byte(this.offsetKey)...)
 			offsetKey = append(offsetKey, 0)
 		} else {
 			offsetKey = prefix
 		}
 
 		opt.LowerBound = offsetKey
-		opt.UpperBound = append(prefix, 0xFF)
+		opt.UpperBound = byteutils.Append(prefix, 0xFF)
 	}
 
 	it, itErr := this.tx.NewIterator(opt)
@@ -391,7 +396,7 @@ func (this *Query[T]) iterateFields(fn IteratorFunc[T]) error {
 
 	var count int
 
-	var itemFn = func() (goNext bool, err error) {
+	var itemFn = func() (goNextItem bool, err error) {
 		var fieldKeyBytes = it.Key()
 
 		fieldValueBytes, keyBytes, decodeKeyErr := this.table.DecodeFieldKey(this.fieldName, fieldKeyBytes)
@@ -423,7 +428,7 @@ func (this *Query[T]) iterateFields(fn IteratorFunc[T]) error {
 		if !this.keysOnly {
 			value, getErr := this.table.getWithKeyBytes(this.tx, this.table.FullKeyBytes(keyBytes))
 			if getErr != nil {
-				if IsKeyNotFound(getErr) {
+				if IsNotFound(getErr) {
 					return true, nil
 				}
 				return false, getErr
@@ -432,11 +437,15 @@ func (this *Query[T]) iterateFields(fn IteratorFunc[T]) error {
 			resultItem.Value = value
 		}
 
-		goNext, err = fn(this.tx, resultItem)
+		goNextItem, err = fn(this.tx, resultItem)
 		if err != nil {
-			return
+			if IsSkipError(err) {
+				return true, nil
+			} else {
+				return false, err
+			}
 		}
-		if !goNext {
+		if !goNextItem {
 			return false, nil
 		}
 

internal/utils/kvstore/query_test.go

@@ -138,6 +138,26 @@ func TestQuery_FindAll_Offset(t *testing.T) {
 	}
 }
 
+func TestQuery_FindAll_Skip(t *testing.T) {
+	var table = testOpenStoreTable[*testCachedItem](t, "cache_items", &testCacheItemEncoder[*testCachedItem]{})
+
+	{
+		err := table.Query().
+			Offset("a3").
+			Limit(10).
+			FindAll(func(tx *kvstore.Tx[*testCachedItem], item kvstore.Item[*testCachedItem]) (goNext bool, err error) {
+				if item.Key == "a30" || item.Key == "a3000005" {
+					return kvstore.Skip()
+				}
+				t.Log("key:", item.Key, "value:", item.Value)
+				return true, nil
+			})
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
 func TestQuery_FindAll_Count(t *testing.T) {
 	var table = testOpenStoreTable[*testCachedItem](t, "cache_items", &testCacheItemEncoder[*testCachedItem]{})
 

internal/utils/kvstore/store.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"github.com/TeaOSLab/EdgeNode/internal/events"
+	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
 	memutils "github.com/TeaOSLab/EdgeNode/internal/utils/mem"
 	"github.com/cockroachdb/pebble"
 	"github.com/iwind/TeaGo/Tea"
@@ -85,6 +86,31 @@ func OpenStoreDir(dir string, storeName string) (*Store, error) {
 	return store, nil
 }
 
+var storeOnce = &sync.Once{}
+var defaultSore *Store
+
+func DefaultStore() (*Store, error) {
+	if defaultSore != nil {
+		return defaultSore, nil
+	}
+
+	storeOnce.Do(func() {
+		store, err := NewStore("default")
+		if err != nil {
+			remotelogs.Error("KV", "create default store failed: "+err.Error())
+			return
+		}
+		err = store.Open()
+		if err != nil {
+			remotelogs.Error("KV", "open default store failed: "+err.Error())
+			return
+		}
+		defaultSore = store
+	})
+
+	return defaultSore, nil
+}
+
 func (this *Store) Open() error {
 	var opt = &pebble.Options{
 		Logger: NewLogger(),
@@ -144,6 +170,10 @@ func (this *Store) RawDB() *pebble.DB {
 	return this.rawDB
 }
 
+func (this *Store) Flush() error {
+	return this.rawDB.Flush()
+}
+
 func (this *Store) Close() error {
 	if this.isClosed {
 		return nil

internal/utils/kvstore/store_test.go

@@ -6,7 +6,9 @@ import (
 	"github.com/TeaOSLab/EdgeNode/internal/utils/kvstore"
 	"github.com/cockroachdb/pebble"
 	"github.com/iwind/TeaGo/Tea"
+	"github.com/iwind/TeaGo/assert"
 	_ "github.com/iwind/TeaGo/bootstrap"
+	"sync"
 	"testing"
 	"time"
 )
@@ -19,6 +21,44 @@ func TestMain(m *testing.M) {
 	}
 }
 
+func TestStore_Default(t *testing.T) {
+	var a = assert.NewAssertion(t)
+
+	store, err := kvstore.DefaultStore()
+	if err != nil {
+		t.Fatal(err)
+	}
+	a.IsTrue(store != nil)
+}
+
+func TestStore_Default_Concurrent(t *testing.T) {
+	var lastStore *kvstore.Store
+
+	const threads = 32
+
+	var wg = &sync.WaitGroup{}
+	wg.Add(threads)
+	for i := 0; i < threads; i++ {
+		go func() {
+			defer wg.Done()
+
+			store, err := kvstore.DefaultStore()
+			if err != nil {
+				t.Log("ERROR", err)
+				t.Fail()
+			}
+
+			if lastStore != nil && lastStore != store {
+				t.Log("ERROR", "should be single instance")
+				t.Fail()
+			}
+
+			lastStore = store
+		}()
+	}
+	wg.Wait()
+}
+
 func TestStore_Open(t *testing.T) {
 	store, err := kvstore.OpenStore("test")
 	if err != nil {
@@ -29,6 +69,7 @@ func TestStore_Open(t *testing.T) {
 	}()
 
 	t.Log("opened")
+	_ = store
 }
 
 func TestStore_RawDB(t *testing.T) {

internal/utils/kvstore/table.go

@@ -64,6 +64,10 @@ func (this *Table[T]) DB() *DB {
 	return this.db
 }
 
+func (this *Table[T]) Encoder() ValueEncoder[T] {
+	return this.encoder
+}
+
 func (this *Table[T]) Set(key string, value T) error {
 	if len(key) > KeyMaxLength {
 		return ErrKeyTooLong
@@ -75,7 +79,22 @@ func (this *Table[T]) Set(key string, value T) error {
 	}
 
 	return this.WriteTx(func(tx *Tx[T]) error {
-		return this.set(tx, key, valueBytes, value, false)
+		return this.set(tx, key, valueBytes, value, false, false)
+	})
+}
+
+func (this *Table[T]) SetSync(key string, value T) error {
+	if len(key) > KeyMaxLength {
+		return ErrKeyTooLong
+	}
+
+	valueBytes, err := this.encoder.Encode(value)
+	if err != nil {
+		return err
+	}
+
+	return this.WriteTxSync(func(tx *Tx[T]) error {
+		return this.set(tx, key, valueBytes, value, false, true)
 	})
 }
 
@@ -90,7 +109,7 @@ func (this *Table[T]) Insert(key string, value T) error {
 	}
 
 	return this.WriteTx(func(tx *Tx[T]) error {
-		return this.set(tx, key, valueBytes, value, true)
+		return this.set(tx, key, valueBytes, value, true, false)
 	})
 }
 
@@ -111,7 +130,7 @@ func (this *Table[T]) ComposeFieldKey(keyBytes []byte, fieldName string, fieldVa
 func (this *Table[T]) Exist(key string) (found bool, err error) {
 	_, closer, err := this.db.store.rawDB.Get(this.FullKey(key))
 	if err != nil {
-		if IsKeyNotFound(err) {
+		if IsNotFound(err) {
 			return false, nil
 		}
 		return false, err
@@ -173,6 +192,20 @@ func (this *Table[T]) WriteTx(fn func(tx *Tx[T]) error) error {
 	return tx.Commit()
 }
 
+func (this *Table[T]) WriteTxSync(fn func(tx *Tx[T]) error) error {
+	var tx = NewTx[T](this, false)
+	defer func() {
+		_ = tx.Close()
+	}()
+
+	err := fn(tx)
+	if err != nil {
+		return err
+	}
+
+	return tx.CommitSync()
+}
+
 func (this *Table[T]) Truncate() error {
 	this.mu.Lock()
 	defer this.mu.Unlock()
@@ -256,7 +289,7 @@ func (this *Table[T]) deleteKeys(tx *Tx[T], key ...string) error {
 			if len(this.fieldNames) > 0 {
 				valueBytes, closer, getErr := batch.Get(keyBytes)
 				if getErr != nil {
-					if IsKeyNotFound(getErr) {
+					if IsNotFound(getErr) {
 						return nil
 					}
 					return getErr
@@ -298,8 +331,12 @@ func (this *Table[T]) deleteKeys(tx *Tx[T], key ...string) error {
 	return nil
 }
 
-func (this *Table[T]) set(tx *Tx[T], key string, valueBytes []byte, value T, insertOnly bool) error {
+func (this *Table[T]) set(tx *Tx[T], key string, valueBytes []byte, value T, insertOnly bool, syncMode bool) error {
 	var keyBytes = this.FullKey(key)
+	var writeOptions = DefaultWriteOptions
+	if syncMode {
+		writeOptions = DefaultWriteSyncOptions
+	}
 
 	var batch = tx.batch
 
@@ -312,7 +349,7 @@ func (this *Table[T]) set(tx *Tx[T], key string, valueBytes []byte, value T, ins
 		if countFields > 0 {
 			oldValueBytes, closer, getErr := batch.Get(keyBytes)
 			if getErr != nil {
-				if !IsKeyNotFound(getErr) {
+				if !IsNotFound(getErr) {
 					return getErr
 				}
 			} else {
@@ -330,7 +367,7 @@ func (this *Table[T]) set(tx *Tx[T], key string, valueBytes []byte, value T, ins
 		}
 	}
 
-	setErr := batch.Set(keyBytes, valueBytes, DefaultWriteOptions)
+	setErr := batch.Set(keyBytes, valueBytes, writeOptions)
 	if setErr != nil {
 		return setErr
 	}
@@ -362,14 +399,14 @@ func (this *Table[T]) set(tx *Tx[T], key string, valueBytes []byte, value T, ins
 					// skip the field
 					continue
 				}
-				deleteFieldErr := batch.Delete(oldFieldKeyBytes, DefaultWriteOptions)
+				deleteFieldErr := batch.Delete(oldFieldKeyBytes, writeOptions)
 				if deleteFieldErr != nil {
 					return deleteFieldErr
 				}
 			}
 
 			// set new field key
-			setFieldErr := batch.Set(newFieldKeyBytes, nil, DefaultWriteOptions)
+			setFieldErr := batch.Set(newFieldKeyBytes, nil, writeOptions)
 			if setFieldErr != nil {
 				return setFieldErr
 			}

internal/utils/kvstore/table_counter.go

@@ -21,7 +21,7 @@ func (this *CounterTable[T]) Increase(key string, delta T) (newValue T, err erro
 	err = this.Table.WriteTx(func(tx *Tx[T]) error {
 		value, getErr := tx.Get(key)
 		if getErr != nil {
-			if !IsKeyNotFound(getErr) {
+			if !IsNotFound(getErr) {
 				return getErr
 			}
 		}

internal/utils/kvstore/table_test.go

@@ -45,7 +45,7 @@ func TestTable_Set(t *testing.T) {
 
 	value, err := table.Get("a")
 	if err != nil {
-		if kvstore.IsKeyNotFound(err) {
+		if kvstore.IsNotFound(err) {
 			t.Log("not found key")
 			return
 		}
@@ -81,7 +81,7 @@ func TestTable_Get(t *testing.T) {
 	for _, key := range []string{"a", "b", "c"} {
 		value, getErr := table.Get(key)
 		if getErr != nil {
-			if kvstore.IsKeyNotFound(getErr) {
+			if kvstore.IsNotFound(getErr) {
 				t.Log("not found key", key)
 				continue
 			}
@@ -146,7 +146,7 @@ func TestTable_Delete(t *testing.T) {
 
 	value, err := table.Get("a123")
 	if err != nil {
-		if !kvstore.IsKeyNotFound(err) {
+		if !kvstore.IsNotFound(err) {
 			t.Fatal(err)
 		}
 	} else {
@@ -173,7 +173,7 @@ func TestTable_Delete(t *testing.T) {
 
 	{
 		_, err = table.Get("a123")
-		a.IsTrue(kvstore.IsKeyNotFound(err))
+		a.IsTrue(kvstore.IsNotFound(err))
 	}
 }
 
@@ -357,7 +357,7 @@ func BenchmarkTable_Get(b *testing.B) {
 		for pb.Next() {
 			_, putErr := table.Get(types.String(rand.Int()))
 			if putErr != nil {
-				if kvstore.IsKeyNotFound(putErr) {
+				if kvstore.IsNotFound(putErr) {
 					continue
 				}
 				b.Fatal(putErr)

internal/utils/kvstore/tx.go

@@ -37,7 +37,24 @@ func (this *Tx[T]) Set(key string, value T) error {
 		return err
 	}
 
-	return this.table.set(this, key, valueBytes, value, false)
+	return this.table.set(this, key, valueBytes, value, false, false)
+}
+
+func (this *Tx[T]) SetSync(key string, value T) error {
+	if this.readOnly {
+		return errors.New("can not set value in readonly transaction")
+	}
+
+	if len(key) > KeyMaxLength {
+		return ErrKeyTooLong
+	}
+
+	valueBytes, err := this.table.encoder.Encode(value)
+	if err != nil {
+		return err
+	}
+
+	return this.table.set(this, key, valueBytes, value, false, true)
 }
 
 func (this *Tx[T]) Insert(key string, value T) error {
@@ -54,7 +71,7 @@ func (this *Tx[T]) Insert(key string, value T) error {
 		return err
 	}
 
-	return this.table.set(this, key, valueBytes, value, true)
+	return this.table.set(this, key, valueBytes, value, true, false)
 }
 
 func (this *Tx[T]) Get(key string) (value T, err error) {
@@ -78,6 +95,20 @@ func (this *Tx[T]) Close() error {
 }
 
 func (this *Tx[T]) Commit() (err error) {
+	return this.commit(DefaultWriteOptions)
+}
+
+func (this *Tx[T]) CommitSync() (err error) {
+	return this.commit(DefaultWriteSyncOptions)
+}
+
+func (this *Tx[T]) Query() *Query[T] {
+	var query = NewQuery[T]()
+	query.SetTx(this)
+	return query
+}
+
+func (this *Tx[T]) commit(opt *pebble.WriteOptions) (err error) {
 	defer func() {
 		var panicErr = recover()
 		if panicErr != nil {
@@ -88,11 +119,5 @@ func (this *Tx[T]) Commit() (err error) {
 		}
 	}()
 
-	return this.batch.Commit(DefaultWriteOptions)
-}
-
-func (this *Tx[T]) Query() *Query[T] {
-	var query = NewQuery[T]()
-	query.SetTx(this)
-	return query
+	return this.batch.Commit(opt)
 }

internal/utils/testutils/utils.go

@@ -2,7 +2,11 @@
 
 package testutils
 
-import "os"
+import (
+	"fmt"
+	"math/rand"
+	"os"
+)
 
 // IsSingleTesting 判断当前测试环境是否为单个函数测试
 func IsSingleTesting() bool {
@@ -12,4 +16,9 @@ func IsSingleTesting() bool {
 		}
 	}
 	return false
-}
+}
+
+// RandIP 生成一个随机IP用于测试
+func RandIP() string {
+	return fmt.Sprintf("%d.%d.%d.%d", rand.Int()%255, rand.Int()%255, rand.Int()%255, rand.Int()%255)
+}