2018-03-17 11:53:21 +00:00
|
|
|
package storage
|
|
|
|
|
2018-03-21 16:11:04 +00:00
|
|
|
import (
|
|
|
|
"encoding/hex"
|
2019-09-18 15:20:41 +00:00
|
|
|
"strings"
|
2019-08-27 13:32:52 +00:00
|
|
|
"sync"
|
2018-03-21 16:11:04 +00:00
|
|
|
)
|
|
|
|
|
2018-03-17 11:53:21 +00:00
|
|
|
// MemoryStore is an in-memory implementation of a Store, mainly
|
|
|
|
// used for testing. Do not use MemoryStore in production.
|
|
|
|
type MemoryStore struct {
|
2019-08-27 13:32:52 +00:00
|
|
|
*sync.RWMutex
|
2018-03-17 11:53:21 +00:00
|
|
|
mem map[string][]byte
|
|
|
|
}
|
|
|
|
|
|
|
|
// MemoryBatch a in-memory batch compatible with MemoryStore.
|
|
|
|
type MemoryBatch struct {
|
|
|
|
m map[*[]byte][]byte
|
|
|
|
}
|
|
|
|
|
|
|
|
// Put implements the Batch interface.
|
|
|
|
func (b *MemoryBatch) Put(k, v []byte) {
|
|
|
|
key := &k
|
|
|
|
b.m[key] = v
|
|
|
|
}
|
|
|
|
|
|
|
|
// Len implements the Batch interface.
|
|
|
|
func (b *MemoryBatch) Len() int {
|
|
|
|
return len(b.m)
|
|
|
|
}
|
|
|
|
|
|
|
|
// NewMemoryStore creates a new MemoryStore object.
|
|
|
|
func NewMemoryStore() *MemoryStore {
|
|
|
|
return &MemoryStore{
|
2019-08-27 13:32:52 +00:00
|
|
|
RWMutex: new(sync.RWMutex),
|
|
|
|
mem: make(map[string][]byte),
|
2018-03-17 11:53:21 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
// Get implements the Store interface.
|
|
|
|
func (s *MemoryStore) Get(key []byte) ([]byte, error) {
|
2019-08-27 13:32:52 +00:00
|
|
|
s.RLock()
|
|
|
|
defer s.RUnlock()
|
2018-03-21 16:11:04 +00:00
|
|
|
if val, ok := s.mem[makeKey(key)]; ok {
|
2018-03-17 11:53:21 +00:00
|
|
|
return val, nil
|
|
|
|
}
|
|
|
|
return nil, ErrKeyNotFound
|
|
|
|
}
|
|
|
|
|
2019-02-09 15:53:58 +00:00
|
|
|
// Put implements the Store interface.
|
2018-03-17 11:53:21 +00:00
|
|
|
func (s *MemoryStore) Put(key, value []byte) error {
|
2019-08-27 13:32:52 +00:00
|
|
|
s.Lock()
|
2018-03-21 16:11:04 +00:00
|
|
|
s.mem[makeKey(key)] = value
|
2019-08-27 13:32:52 +00:00
|
|
|
s.Unlock()
|
2018-03-17 11:53:21 +00:00
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2019-02-09 15:53:58 +00:00
|
|
|
// PutBatch implements the Store interface.
|
2018-03-17 11:53:21 +00:00
|
|
|
func (s *MemoryStore) PutBatch(batch Batch) error {
|
|
|
|
b := batch.(*MemoryBatch)
|
|
|
|
for k, v := range b.m {
|
2019-01-25 11:20:35 +00:00
|
|
|
if err := s.Put(*k, v); err != nil {
|
|
|
|
return err
|
|
|
|
}
|
2018-03-17 11:53:21 +00:00
|
|
|
}
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2019-02-09 15:53:58 +00:00
|
|
|
// Seek implements the Store interface.
|
2018-03-17 11:53:21 +00:00
|
|
|
func (s *MemoryStore) Seek(key []byte, f func(k, v []byte)) {
|
2019-09-18 15:20:41 +00:00
|
|
|
for k, v := range s.mem {
|
|
|
|
if strings.Contains(k, hex.EncodeToString(key)) {
|
|
|
|
decodeString, _ := hex.DecodeString(k)
|
|
|
|
f(decodeString, v)
|
|
|
|
}
|
|
|
|
}
|
2018-03-17 11:53:21 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
// Batch implements the Batch interface and returns a compatible Batch.
|
|
|
|
func (s *MemoryStore) Batch() Batch {
|
|
|
|
return &MemoryBatch{
|
|
|
|
m: make(map[*[]byte][]byte),
|
|
|
|
}
|
|
|
|
}
|
2018-03-21 16:11:04 +00:00
|
|
|
|
2019-09-16 15:52:47 +00:00
|
|
|
// Close implements Store interface and clears up memory.
|
|
|
|
func (s *MemoryStore) Close() error {
|
|
|
|
s.Lock()
|
|
|
|
s.mem = nil
|
|
|
|
s.Unlock()
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2018-03-21 16:11:04 +00:00
|
|
|
func makeKey(k []byte) string {
|
|
|
|
return hex.EncodeToString(k)
|
|
|
|
}
|