使用KV存储实现指标统计

internal/utils/agents/manager.go

@@ -199,7 +199,7 @@ func (this *Manager) loadDB() error {
 	var sqlitePath = Tea.Root + "/data/agents.db"
 	_, sqliteErr := os.Stat(sqlitePath)
 	var db DB
-	if sqliteErr == nil {
+	if sqliteErr == nil || !teaconst.EnableKVCacheStore {
 		db = NewSQLiteDB(sqlitePath)
 	} else {
 		db = NewKVDB()

internal/utils/byte/utils.go

@@ -18,8 +18,8 @@ func Append(b []byte, b2 ...byte) []byte {
 	return append(Copy(b), b2...)
 }
 
-// Contact bytes
-func Contact(b []byte, b2 ...[]byte) []byte {
+// Concat bytes
+func Concat(b []byte, b2 ...[]byte) []byte {
 	b = Copy(b)
 	for _, b3 := range b2 {
 		b = append(b, b3...)

internal/utils/byte/utils_test.go

@@ -39,7 +39,7 @@ func TestConcat(t *testing.T) {
 	var prefix []byte
 	prefix = append(prefix, 1, 2, 3)
 
-	var b = byteutils.Contact(prefix, []byte{4, 5, 6}, []byte{7})
+	var b = byteutils.Concat(prefix, []byte{4, 5, 6}, []byte{7})
 	t.Log(b)
 
 	a.IsTrue(bytes.Equal(b, []byte{1, 2, 3, 4, 5, 6, 7}))

internal/utils/kvstore/db.go

@@ -52,6 +52,15 @@ func (this *DB) Store() *Store {
 	return this.store
 }
 
+// Truncate the database
+func (this *DB) Truncate() error {
+	this.mu.Lock()
+	defer this.mu.Unlock()
+
+	var start = []byte(this.Namespace())
+	return this.store.rawDB.DeleteRange(start, append(start, 0xFF), DefaultWriteOptions)
+}
+
 func (this *DB) Close() error {
 	this.mu.Lock()
 	defer this.mu.Unlock()

internal/utils/kvstore/store.go

@@ -162,14 +162,22 @@ func (this *Store) Delete(keyBytes []byte) error {
 }
 
 func (this *Store) NewDB(dbName string) (*DB, error) {
+	this.mu.Lock()
+	defer this.mu.Unlock()
+
+	// check existence
+	for _, db := range this.dbs {
+		if db.name == dbName {
+			return db, nil
+		}
+	}
+
+	// create new
 	db, err := NewDB(this, dbName)
 	if err != nil {
 		return nil, err
 	}
 
-	this.mu.Lock()
-	defer this.mu.Unlock()
-
 	this.dbs = append(this.dbs, db)
 	return db, nil
 }

internal/utils/kvstore/table.go

@@ -6,6 +6,7 @@ import (
 	"bytes"
 	"encoding/binary"
 	"errors"
+	"fmt"
 	"github.com/cockroachdb/pebble"
 	"github.com/iwind/TeaGo/types"
 	"sync"
@@ -213,6 +214,10 @@ func (this *Table[T]) Truncate() error {
 	return this.db.store.rawDB.DeleteRange(this.Namespace(), append(this.Namespace(), 0xFF), DefaultWriteOptions)
 }
 
+func (this *Table[T]) DeleteRange(start string, end string) error {
+	return this.db.store.rawDB.DeleteRange(this.FullKeyBytes([]byte(start)), this.FullKeyBytes([]byte(end)), DefaultWriteOptions)
+}
+
 func (this *Table[T]) Query() *Query[T] {
 	var query = NewQuery[T]()
 	query.SetTable(this)
@@ -275,6 +280,7 @@ func (this *Table[T]) DecodeFieldKey(fieldName string, fieldKey []byte) (fieldVa
 }
 
 func (this *Table[T]) Close() error {
+	// nothing to do
 	return nil
 }
 
@@ -300,7 +306,7 @@ func (this *Table[T]) deleteKeys(tx *Tx[T], key ...string) error {
 
 				value, decodeErr := this.encoder.Decode(valueBytes)
 				if decodeErr != nil {
-					return decodeErr
+					return fmt.Errorf("decode value failed: %w", decodeErr)
 				}
 
 				for _, fieldName := range this.fieldNames {
@@ -360,7 +366,7 @@ func (this *Table[T]) set(tx *Tx[T], key string, valueBytes []byte, value T, ins
 				var decodeErr error
 				oldValue, decodeErr = this.encoder.Decode(oldValueBytes)
 				if decodeErr != nil {
-					return decodeErr
+					return fmt.Errorf("decode value failed: %w", decodeErr)
 				}
 				oldFound = true
 			}
@@ -431,7 +437,7 @@ func (this *Table[T]) getWithKeyBytes(tx *Tx[T], keyBytes []byte) (value T, err
 
 	resultValue, decodeErr := this.encoder.Decode(valueBytes)
 	if decodeErr != nil {
-		return value, decodeErr
+		return value, fmt.Errorf("decode value failed: %w", decodeErr)
 	}
 	value = resultValue
 	return

internal/utils/kvstore/value_encoder_nil.go

@@ -0,0 +1,22 @@
+// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package kvstore
+
+type NilValueEncoder[T []byte] struct {
+}
+
+func NewNilValueEncoder[T []byte]() *NilValueEncoder[T] {
+	return &NilValueEncoder[T]{}
+}
+
+func (this *NilValueEncoder[T]) Encode(value T) ([]byte, error) {
+	return nil, nil
+}
+
+func (this *NilValueEncoder[T]) EncodeField(value T, fieldName string) ([]byte, error) {
+	return nil, nil
+}
+
+func (this *NilValueEncoder[T]) Decode(valueData []byte) (value T, err error) {
+	return
+}