From 7a12a78359a4da9e6f4195290686a4080a569119 Mon Sep 17 00:00:00 2001
From: GoEdgeLab <goedgecloud@gmail.com>
Date: Thu, 28 Mar 2024 17:18:00 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E4=BB=A3=E7=A0=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/utils/kvstore/table.go   | 52 +++++++++++++++++++++----------
 internal/utils/kvstore/tx.go      | 19 ++++++++++-
 internal/utils/mem/system_test.go | 17 ++++++++++
 3 files changed, 70 insertions(+), 18 deletions(-)
 create mode 100644 internal/utils/mem/system_test.go

diff --git a/internal/utils/kvstore/table.go b/internal/utils/kvstore/table.go
index 70d37b7..ccee18f 100644
--- a/internal/utils/kvstore/table.go
+++ b/internal/utils/kvstore/table.go
@@ -75,7 +75,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)
+		return this.set(tx, key, valueBytes, value, false)
+	})
+}
+
+func (this *Table[T]) Insert(key string, value T) error {
+	if len(key) > KeyMaxLength {
+		return ErrKeyTooLong
+	}
+
+	valueBytes, err := this.encoder.Encode(value)
+	if err != nil {
+		return err
+	}
+
+	return this.WriteTx(func(tx *Tx[T]) error {
+		return this.set(tx, key, valueBytes, value, true)
 	})
 }
 
@@ -283,7 +298,7 @@ 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) error {
+func (this *Table[T]) set(tx *Tx[T], key string, valueBytes []byte, value T, insertOnly bool) error {
 	var keyBytes = this.FullKey(key)
 
 	var batch = tx.batch
@@ -292,23 +307,26 @@ func (this *Table[T]) set(tx *Tx[T], key string, valueBytes []byte, value T) err
 	var oldValue T
 	var oldFound bool
 	var countFields = len(this.fieldNames)
-	if countFields > 0 {
-		oldValueBytes, closer, getErr := batch.Get(keyBytes)
-		if getErr != nil {
-			if !IsKeyNotFound(getErr) {
-				return getErr
-			}
-		} else {
-			defer func() {
-				_ = closer.Close()
-			}()
 
-			var decodeErr error
-			oldValue, decodeErr = this.encoder.Decode(oldValueBytes)
-			if decodeErr != nil {
-				return decodeErr
+	if !insertOnly {
+		if countFields > 0 {
+			oldValueBytes, closer, getErr := batch.Get(keyBytes)
+			if getErr != nil {
+				if !IsKeyNotFound(getErr) {
+					return getErr
+				}
+			} else {
+				defer func() {
+					_ = closer.Close()
+				}()
+
+				var decodeErr error
+				oldValue, decodeErr = this.encoder.Decode(oldValueBytes)
+				if decodeErr != nil {
+					return decodeErr
+				}
+				oldFound = true
 			}
-			oldFound = true
 		}
 	}
 
diff --git a/internal/utils/kvstore/tx.go b/internal/utils/kvstore/tx.go
index 0cc2463..7df329a 100644
--- a/internal/utils/kvstore/tx.go
+++ b/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)
+	return this.table.set(this, key, valueBytes, value, false)
+}
+
+func (this *Tx[T]) Insert(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, true)
 }
 
 func (this *Tx[T]) Get(key string) (value T, err error) {
diff --git a/internal/utils/mem/system_test.go b/internal/utils/mem/system_test.go
new file mode 100644
index 0000000..b05e0d5
--- /dev/null
+++ b/internal/utils/mem/system_test.go
@@ -0,0 +1,17 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package memutils_test
+
+import (
+	"github.com/TeaOSLab/EdgeNode/internal/utils/mem"
+	"testing"
+)
+
+func TestSystemMemoryGB(t *testing.T) {
+	t.Log(memutils.SystemMemoryGB())
+	t.Log(memutils.SystemMemoryGB())
+	t.Log(memutils.SystemMemoryGB())
+	t.Log(memutils.SystemMemoryBytes())
+	t.Log(memutils.SystemMemoryBytes())
+	t.Log(memutils.SystemMemoryBytes()>>30, "GB")
+}