From 076d5fe84c95416b51fa6dfa808582ab62fe3e7b Mon Sep 17 00:00:00 2001
From: Mikhail Rakhmanov <rakhmanov.m@gmail.com>
Date: Sat, 3 Aug 2024 03:28:22 +0200
Subject: [PATCH] Implement GetAppend (#68)

---
 db.go      | 28 ++++++++++++++++++++++++++++
 db_test.go | 37 +++++++++++++++++++++++++++++++++++--
 2 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/db.go b/db.go
index e2c908e..d12e846 100644
--- a/db.go
+++ b/db.go
@@ -215,6 +215,34 @@ func (db *DB) Get(key []byte) ([]byte, error) {
 	return retValue, nil
 }
 
+// GetAppend returns the value for the given key (appended into buffer) stored in the DB or nil if the key doesn't exist
+func (db *DB) GetAppend(key, buf []byte) ([]byte, error) {
+	h := db.hash(key)
+	db.metrics.Gets.Add(1)
+	db.mu.RLock()
+	defer db.mu.RUnlock()
+	var retValue []byte
+	err := db.index.get(h, func(sl slot) (bool, error) {
+		if uint16(len(key)) != sl.keySize {
+			return false, nil
+		}
+		slKey, value, err := db.datalog.readKeyValue(sl)
+		if err != nil {
+			return true, err
+		}
+		if bytes.Equal(key, slKey) {
+			retValue = append(buf, value...)
+			return true, nil
+		}
+		db.metrics.HashCollisions.Add(1)
+		return false, nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	return retValue, nil
+}
+
 // Has returns true if the DB contains the given key.
 func (db *DB) Has(key []byte) (bool, error) {
 	h := db.hash(key)
diff --git a/db_test.go b/db_test.go
index a0cec32..8a71bce 100644
--- a/db_test.go
+++ b/db_test.go
@@ -117,6 +117,18 @@ func TestEmpty(t *testing.T) {
 }
 
 func TestFull(t *testing.T) {
+	fullTest(t, func(db *DB, key []byte) ([]byte, error) {
+		return db.Get(key)
+	})
+	var buf []byte
+	fullTest(t, func(db *DB, key []byte) ([]byte, error) {
+		var err error
+		buf, err = db.GetAppend(key, buf[:0])
+		return buf, err
+	})
+}
+
+func fullTest(t *testing.T, getFunc func(db *DB, key []byte) ([]byte, error)) {
 	opts := &Options{
 		BackgroundSyncInterval: -1,
 		FileSystem:             testFS,
@@ -165,7 +177,7 @@ func TestFull(t *testing.T) {
 			if has, err := db.Has([]byte{0, i}); has || err != nil {
 				t.Fatal(has, err)
 			}
-			v, err := db.Get([]byte{i})
+			v, err := getFunc(db, []byte{i})
 			if err != nil {
 				t.Fatal(err)
 			}
@@ -418,10 +430,11 @@ func BenchmarkGet(b *testing.B) {
 	db, err := createTestDB(nil)
 	assert.Nil(b, err)
 	k := []byte{1}
-	if err := db.Put(k, k); err != nil {
+	if err := db.Put(k, make([]byte, 1024)); err != nil {
 		b.Fail()
 	}
 	b.ResetTimer()
+	b.ReportAllocs()
 	for i := 0; i < b.N; i++ {
 		if _, err := db.Get(k); err != nil {
 			b.Fatal()
@@ -430,6 +443,26 @@ func BenchmarkGet(b *testing.B) {
 	assert.Nil(b, db.Close())
 }
 
+func BenchmarkGetAppend(b *testing.B) {
+	db, err := createTestDB(nil)
+	assert.Nil(b, err)
+	k := []byte{1}
+	if err := db.Put(k, make([]byte, 1024)); err != nil {
+		b.Fail()
+	}
+	b.ResetTimer()
+	b.ReportAllocs()
+	buf := make([]byte, 0, 1024)
+	for i := 0; i < b.N; i++ {
+		value, err := db.GetAppend(k, buf[:0])
+		if err != nil {
+			b.Fatal()
+		}
+		buf = value
+	}
+	assert.Nil(b, db.Close())
+}
+
 func BenchmarkBucket_UnmarshalBinary(b *testing.B) {
 	testBucket := bucket{
 		slots: [slotsPerBucket]slot{},