diff --git a/pkg/core/storage/boltdb_store.go b/pkg/core/storage/boltdb_store.go index ffa1d4ecb..18283bc63 100644 --- a/pkg/core/storage/boltdb_store.go +++ b/pkg/core/storage/boltdb_store.go @@ -24,25 +24,6 @@ type BoltDBStore struct { db *bbolt.DB } -// BoltDBBatch simple batch implementation to satisfy the Store interface. -type BoltDBBatch struct { - mem map[*[]byte][]byte -} - -// Len implements the Batch interface. -func (b *BoltDBBatch) Len() int { - return len(b.mem) -} - -// Put implements the Batch interface. -func (b *BoltDBBatch) Put(k, v []byte) { - vcopy := make([]byte, len(v)) - copy(vcopy, v) - kcopy := make([]byte, len(k)) - copy(kcopy, k) - b.mem[&kcopy] = vcopy -} - // NewBoltDBStore returns a new ready to use BoltDB storage with created bucket. func NewBoltDBStore(cfg BoltDBOptions) (*BoltDBStore, error) { var opts *bbolt.Options // should be exposed via BoltDBOptions if anything needed @@ -94,7 +75,7 @@ func (s *BoltDBStore) Get(key []byte) (val []byte, err error) { func (s *BoltDBStore) PutBatch(batch Batch) error { return s.db.Batch(func(tx *bbolt.Tx) error { b := tx.Bucket(Bucket) - for k, v := range batch.(*BoltDBBatch).mem { + for k, v := range batch.(*MemoryBatch).m { err := b.Put(*k, v) if err != nil { return err @@ -122,9 +103,7 @@ func (s *BoltDBStore) Seek(key []byte, f func(k, v []byte)) { // Batch implements the Batch interface and returns a boltdb // compatible Batch. func (s *BoltDBStore) Batch() Batch { - return &BoltDBBatch{ - mem: make(map[*[]byte][]byte), - } + return newMemoryBatch() } // Close releases all db resources. diff --git a/pkg/core/storage/boltdb_store_test.go b/pkg/core/storage/boltdb_store_test.go index ac23d136a..e5cf397b1 100644 --- a/pkg/core/storage/boltdb_store_test.go +++ b/pkg/core/storage/boltdb_store_test.go @@ -12,14 +12,14 @@ import ( func TestBoltDBBatch(t *testing.T) { boltDB := BoltDBStore{} - want := &BoltDBBatch{mem: map[*[]byte][]byte{}} + want := &MemoryBatch{m: map[*[]byte][]byte{}} if got := boltDB.Batch(); !reflect.DeepEqual(got, want) { t.Errorf("BoltDB Batch() = %v, want %v", got, want) } } func TestBoltDBBatch_Len(t *testing.T) { - batch := &BoltDBBatch{mem: map[*[]byte][]byte{}} + batch := &MemoryBatch{m: map[*[]byte][]byte{}} want := len(map[*[]byte][]byte{}) assert.Equal(t, want, batch.Len()) } diff --git a/pkg/core/storage/memory_store.go b/pkg/core/storage/memory_store.go index 197202c1f..45484d546 100644 --- a/pkg/core/storage/memory_store.go +++ b/pkg/core/storage/memory_store.go @@ -20,8 +20,11 @@ type MemoryBatch struct { // Put implements the Batch interface. func (b *MemoryBatch) Put(k, v []byte) { - key := &k - b.m[key] = v + vcopy := make([]byte, len(v)) + copy(vcopy, v) + kcopy := make([]byte, len(k)) + copy(kcopy, k) + b.m[&kcopy] = vcopy } // Len implements the Batch interface. @@ -76,6 +79,11 @@ func (s *MemoryStore) Seek(key []byte, f func(k, v []byte)) { // Batch implements the Batch interface and returns a compatible Batch. func (s *MemoryStore) Batch() Batch { + return newMemoryBatch() +} + +// newMemoryBatch returns new memory batch. +func newMemoryBatch() *MemoryBatch { return &MemoryBatch{ m: make(map[*[]byte][]byte), }