KV存储增加panic处理

internal/utils/kvstore/query_test.go

@@ -5,6 +5,7 @@ package kvstore_test
 import (
 	"fmt"
 	"github.com/TeaOSLab/EdgeNode/internal/utils/kvstore"
+	"github.com/iwind/TeaGo/assert"
 	"runtime"
 	"testing"
 	"time"
@@ -25,6 +26,7 @@ func TestQuery_FindAll(t *testing.T) {
 		//Desc().
 		FindAll(func(tx *kvstore.Tx[*testCachedItem], item kvstore.Item[*testCachedItem]) (goNext bool, err error) {
 			t.Log("key:", item.Key, "value:", item.Value)
+
 			return true, nil
 		})
 	if err != nil {
@@ -46,6 +48,12 @@ func TestQuery_FindAll_Break(t *testing.T) {
 		FindAll(func(tx *kvstore.Tx[*testCachedItem], item kvstore.Item[*testCachedItem]) (goNext bool, err error) {
 			t.Log("key:", item.Key, "value:", item.Value)
 			count++
+
+			if count > 2 {
+				// break test
+				_ = table.DB().Store().Close()
+			}
+
 			return count < 3, nil
 		})
 	if err != nil {
@@ -53,6 +61,34 @@ func TestQuery_FindAll_Break(t *testing.T) {
 	}
 }
 
+func TestQuery_FindAll_Break_Closed(t *testing.T) {
+	var table = testOpenStoreTable[*testCachedItem](t, "cache_items", &testCacheItemEncoder[*testCachedItem]{})
+
+	var a = assert.NewAssertion(t)
+
+	var before = time.Now()
+	defer func() {
+		t.Log("cost:", time.Since(before).Seconds()*1000, "ms")
+	}()
+
+	var count int
+	err := table.
+		Query().
+		FindAll(func(tx *kvstore.Tx[*testCachedItem], item kvstore.Item[*testCachedItem]) (goNext bool, err error) {
+			t.Log("key:", item.Key, "value:", item.Value)
+			count++
+
+			if count > 2 {
+				// break test
+				_ = table.DB().Store().Close()
+			}
+
+			return count < 3, nil
+		})
+	t.Log("expected error:", err)
+	a.IsTrue(err != nil)
+}
+
 func TestQuery_FindAll_Desc(t *testing.T) {
 	var table = testOpenStoreTable[*testCachedItem](t, "cache_items", &testCacheItemEncoder[*testCachedItem]{})
 

internal/utils/kvstore/store.go

@@ -26,7 +26,7 @@ type Store struct {
 
 	dbs []*DB
 
-	locker sync.Mutex
+	mu sync.Mutex
 }
 
 // NewStore create store with name
@@ -126,8 +126,8 @@ func (this *Store) NewDB(dbName string) (*DB, error) {
 		return nil, err
 	}
 
-	this.locker.Lock()
-	defer this.locker.Unlock()
+	this.mu.Lock()
+	defer this.mu.Unlock()
 
 	this.dbs = append(this.dbs, db)
 	return db, nil
@@ -142,7 +142,7 @@ func (this *Store) Close() error {
 		return nil
 	}
 
-	this.locker.Lock()
+	this.mu.Lock()
 	var lastErr error
 	for _, db := range this.dbs {
 		err := db.Close()
@@ -151,12 +151,16 @@ func (this *Store) Close() error {
 		}
 	}
 
-	this.locker.Unlock()
+	this.mu.Unlock()
 
 	if this.rawDB != nil {
 		this.isClosed = true
-		return this.rawDB.Close()
+		err := this.rawDB.Close()
+		if err != nil {
+			return err
+		}
 	}
+
 	return lastErr
 }
 

internal/utils/kvstore/table.go

@@ -60,6 +60,10 @@ func (this *Table[T]) SetDB(db *DB) {
 	this.db = db
 }
 
+func (this *Table[T]) DB() *DB {
+	return this.db
+}
+
 func (this *Table[T]) Set(key string, value T) error {
 	if len(key) > KeyMaxLength {
 		return ErrKeyTooLong

internal/utils/kvstore/tx.go

@@ -4,6 +4,7 @@ package kvstore
 
 import (
 	"errors"
+	"fmt"
 	"github.com/cockroachdb/pebble"
 )
 
@@ -59,7 +60,17 @@ func (this *Tx[T]) Close() error {
 	return this.batch.Close()
 }
 
-func (this *Tx[T]) Commit() error {
+func (this *Tx[T]) Commit() (err error) {
+	defer func() {
+		var panicErr = recover()
+		if panicErr != nil {
+			resultErr, ok := panicErr.(error)
+			if ok {
+				err = fmt.Errorf("commit batch failed: %w", resultErr)
+			}
+		}
+	}()
+
 	return this.batch.Commit(DefaultWriteOptions)
 }