diff --git a/db.go b/db.go index 517baf1..5ec79cf 100644 --- a/db.go +++ b/db.go @@ -302,6 +302,48 @@ func (db *DB) Put(key []byte, value []byte) error { return nil } +// BatchPut set multiple keys at one time, it updates the value for the existing key. +func (db *DB) BatchPut(kvs ...[]byte) error { + db.mu.Lock() + defer db.mu.Unlock() + + for i := 0; i < len(kvs); i += 2 { + key := kvs[i] + value := kvs[i+1] + + if len(key) > MaxKeyLength { + return errKeyTooLarge + } + if len(value) > MaxValueLength { + return errValueTooLarge + } + h := db.hash(key) + db.metrics.Puts.Add(1) + + segID, offset, err := db.datalog.put(key, value) + if err != nil { + return err + } + + sl := slot{ + hash: h, + segmentID: segID, + keySize: uint16(len(key)), + valueSize: uint32(len(value)), + offset: offset, + } + + if err := db.put(sl, key); err != nil { + return err + } + } + + if db.syncWrites { + return db.sync() + } + return nil +} + func (db *DB) del(h uint32, key []byte, writeWAL bool) error { err := db.index.delete(h, func(sl slot) (b bool, e error) { if uint16(len(key)) != sl.keySize { diff --git a/db_test.go b/db_test.go index a0cec32..c9df262 100644 --- a/db_test.go +++ b/db_test.go @@ -10,6 +10,7 @@ import ( "log" "os" "path/filepath" + "strconv" "strings" "testing" "time" @@ -116,6 +117,37 @@ func TestEmpty(t *testing.T) { assert.Nil(t, db.Close()) } +func TestBatchPut(t *testing.T) { + opts := &Options{ + BackgroundSyncInterval: -1, + FileSystem: testFS, + maxSegmentSize: 1024, + } + db, err := createTestDB(opts) + assert.Nil(t, err) + defer db.Close() + + count := 1000 + bucket := 100 + kvs := [][]byte{} + for i := 0; i < count; i++ { + key := []byte(strconv.Itoa(i)) + kvs = append(kvs, key, key) + if len(kvs) == bucket { + err := db.BatchPut(kvs...) + assert.Equal(t, nil, err) + kvs = kvs[:0] + } + } + + for i := 0; i < count; i++ { + expected := []byte(strconv.Itoa(i)) + actual, err := db.Get(expected) + assert.Equal(t, nil, err) + assert.Equal(t, expected, actual) + } +} + func TestFull(t *testing.T) { opts := &Options{ BackgroundSyncInterval: -1, @@ -414,6 +446,32 @@ func BenchmarkPut(b *testing.B) { assert.Nil(b, db.Close()) } +func BenchmarkBatchPut(b *testing.B) { + db, err := createTestDB(nil) + assert.Nil(b, err) + b.ResetTimer() + k := []byte{1} + kvs := [][]byte{} + for i := 0; i < b.N; i++ { + kvs = append(kvs, k, k) + if len(kvs) < 100 { + continue + } + + if err := db.BatchPut(kvs...); err != nil { + b.Fail() + } + kvs = kvs[:0] + } + + if len(kvs) > 0 { + if err := db.BatchPut(kvs...); err != nil { + b.Fail() + } + } + assert.Nil(b, db.Close()) +} + func BenchmarkGet(b *testing.B) { db, err := createTestDB(nil) assert.Nil(b, err)