storage: improve PutBatch for MemoryStore

Make it look more like a real transaction, put/delete things with a single
lock. Make a copy of value in Put also, just for safety purposes, no one knows
how this value slice can be used after the Put.
This commit is contained in:
Roman Khimov 2019-10-07 20:13:59 +03:00
parent 3ada92944a
commit 13bf2618ef

View file

@ -55,22 +55,36 @@ func (s *MemoryStore) Get(key []byte) ([]byte, error) {
return nil, ErrKeyNotFound return nil, ErrKeyNotFound
} }
// put puts a key-value pair into the store, it's supposed to be called
// with mutex locked.
func (s *MemoryStore) put(key string, value []byte) {
s.mem[key] = value
delete(s.del, key)
}
// Put implements the Store interface. Never returns an error. // Put implements the Store interface. Never returns an error.
func (s *MemoryStore) Put(key, value []byte) error { func (s *MemoryStore) Put(key, value []byte) error {
s.mut.Lock()
newKey := string(key) newKey := string(key)
s.mem[newKey] = value vcopy := make([]byte, len(value))
delete(s.del, newKey) copy(vcopy, value)
s.mut.Lock()
s.put(newKey, vcopy)
s.mut.Unlock() s.mut.Unlock()
return nil return nil
} }
// drop deletes a key-valu pair from the store, it's supposed to be called
// with mutex locked.
func (s *MemoryStore) drop(key string) {
s.del[key] = true
delete(s.mem, key)
}
// Delete implements Store interface. Never returns an error. // Delete implements Store interface. Never returns an error.
func (s *MemoryStore) Delete(key []byte) error { func (s *MemoryStore) Delete(key []byte) error {
s.mut.Lock()
newKey := string(key) newKey := string(key)
s.del[newKey] = true s.mut.Lock()
delete(s.mem, newKey) s.drop(newKey)
s.mut.Unlock() s.mut.Unlock()
return nil return nil
} }
@ -78,11 +92,13 @@ func (s *MemoryStore) Delete(key []byte) error {
// PutBatch implements the Store interface. Never returns an error. // PutBatch implements the Store interface. Never returns an error.
func (s *MemoryStore) PutBatch(batch Batch) error { func (s *MemoryStore) PutBatch(batch Batch) error {
b := batch.(*MemoryBatch) b := batch.(*MemoryBatch)
s.mut.Lock()
defer s.mut.Unlock()
for k := range b.del { for k := range b.del {
_ = s.Delete([]byte(k)) s.drop(k)
} }
for k, v := range b.m { for k, v := range b.m {
_ = s.Put([]byte(k), v) s.put(k, v)
} }
return nil return nil
} }