v1.5.1 增强程序稳定性

EdgeNode/internal/utils/kvstore/query.go

@@ -5,7 +5,6 @@ package kvstore
 import (
 	"bytes"
 	"errors"
-	"fmt"
 	byteutils "github.com/TeaOSLab/EdgeNode/internal/utils/byte"
 )
 
@@ -66,6 +65,9 @@ func (this *Query[T]) SetTable(table *Table[T]) *Query[T] {
 
 func (this *Query[T]) SetTx(tx *Tx[T]) *Query[T] {
 	this.tx = tx
+	if tx != nil {
+		this.table = tx.table
+	}
 	return this
 }
 
@@ -128,7 +130,12 @@ func (this *Query[T]) FieldPrefix(fieldName string, fieldPrefix string) *Query[T
 }
 
 func (this *Query[T]) FieldOffset(fieldOffset []byte) *Query[T] {
-	this.fieldOffsetKey = fieldOffset
+	if len(fieldOffset) == 0 {
+		this.fieldOffsetKey = nil
+		return this
+	}
+
+	this.fieldOffsetKey = append([]byte(nil), fieldOffset...)
 	return this
 }
 
@@ -172,28 +179,12 @@ func (this *Query[T]) FindAll(fn IteratorFunc[T]) (err error) {
 	defer func() {
 		var panicErr = recover()
 		if panicErr != nil {
-			resultErr, ok := panicErr.(error)
-			if ok {
-				err = fmt.Errorf("execute query failed: %w", resultErr)
-			}
+			err = wrapRecoveredPanic("execute query failed", panicErr)
 		}
 	}()
 
 	if this.tx != nil {
-		defer func() {
-			_ = this.tx.Close()
-		}()
-
-		var itErr error
-		if len(this.fieldName) == 0 {
-			itErr = this.iterateKeys(fn)
-		} else {
-			itErr = this.iterateFields(fn)
-		}
-		if itErr != nil {
-			return itErr
-		}
-		return this.tx.Commit()
+		return this.findAllWithTx(this.tx, fn)
 	}
 
 	if this.table != nil {
@@ -205,19 +196,29 @@ func (this *Query[T]) FindAll(fn IteratorFunc[T]) (err error) {
 		}
 
 		return txFn(func(tx *Tx[T]) error {
-			this.tx = tx
-
-			if len(this.fieldName) == 0 {
-				return this.iterateKeys(fn)
-			}
-			return this.iterateFields(fn)
+			return this.findAllWithTx(tx, fn)
 		})
 	}
 
 	return errors.New("current query require 'table' or 'tx'")
 }
 
-func (this *Query[T]) iterateKeys(fn IteratorFunc[T]) error {
+func (this *Query[T]) findAllWithTx(tx *Tx[T], fn IteratorFunc[T]) error {
+	if tx == nil {
+		return errors.New("current query require valid tx")
+	}
+
+	if this.table == nil {
+		this.table = tx.table
+	}
+
+	if len(this.fieldName) == 0 {
+		return this.iterateKeys(tx, fn)
+	}
+	return this.iterateFields(tx, fn)
+}
+
+func (this *Query[T]) iterateKeys(tx *Tx[T], fn IteratorFunc[T]) error {
 	if this.limit == 0 {
 		return nil
 	}
@@ -262,7 +263,7 @@ func (this *Query[T]) iterateKeys(fn IteratorFunc[T]) error {
 
 	var hasOffsetKey = len(this.offsetKey) > 0
 
-	it, itErr := this.tx.NewIterator(opt)
+	it, itErr := tx.NewIterator(opt)
 	if itErr != nil {
 		return itErr
 	}
@@ -297,7 +298,7 @@ func (this *Query[T]) iterateKeys(fn IteratorFunc[T]) error {
 			}
 		}
 
-		goNext, callbackErr := fn(this.tx, Item[T]{
+		goNext, callbackErr := fn(tx, Item[T]{
 			Key:   string(keyBytes[prefixLen:]),
 			Value: value,
 		})
@@ -346,10 +347,10 @@ func (this *Query[T]) iterateKeys(fn IteratorFunc[T]) error {
 		}
 	}
 
-	return nil
+	return it.Error()
 }
 
-func (this *Query[T]) iterateFields(fn IteratorFunc[T]) error {
+func (this *Query[T]) iterateFields(tx *Tx[T], fn IteratorFunc[T]) error {
 	if this.limit == 0 {
 		return nil
 	}
@@ -390,7 +391,7 @@ func (this *Query[T]) iterateFields(fn IteratorFunc[T]) error {
 		opt.UpperBound = byteutils.Append(prefix, 0xFF)
 	}
 
-	it, itErr := this.tx.NewIterator(opt)
+	it, itErr := tx.NewIterator(opt)
 	if itErr != nil {
 		return itErr
 	}
@@ -427,10 +428,10 @@ func (this *Query[T]) iterateFields(fn IteratorFunc[T]) error {
 
 		var resultItem = Item[T]{
 			Key:      string(keyBytes),
-			FieldKey: fieldKeyBytes,
+			FieldKey: append([]byte(nil), fieldKeyBytes...),
 		}
 		if !this.keysOnly {
-			value, getErr := this.table.getWithKeyBytes(this.tx, this.table.FullKeyBytes(keyBytes))
+			value, getErr := this.table.getWithKeyBytes(tx, this.table.FullKeyBytes(keyBytes))
 			if getErr != nil {
 				if IsNotFound(getErr) {
 					return true, nil
@@ -441,7 +442,7 @@ func (this *Query[T]) iterateFields(fn IteratorFunc[T]) error {
 			resultItem.Value = value
 		}
 
-		goNextItem, err = fn(this.tx, resultItem)
+		goNextItem, err = fn(tx, resultItem)
 		if err != nil {
 			if IsSkipError(err) {
 				return true, nil
@@ -487,7 +488,7 @@ func (this *Query[T]) iterateFields(fn IteratorFunc[T]) error {
 		}
 	}
 
-	return nil
+	return it.Error()
 }
 
 func (this *Query[T]) matchOperators(fieldValueBytes []byte) bool {