package store import ( "errors" "io" "math/rand" "sort" "sync" "github.com/cockroachdb/pebble" ) type InMemKVDB struct { open bool sortedKeys []string store map[string][]byte storeMx sync.Mutex } type Operation int const ( SetOperation Operation = iota DeleteOperation ) type InMemKVDBOperation struct { op Operation key []byte value []byte } type InMemKVDBTransaction struct { id int changes []InMemKVDBOperation db *InMemKVDB } type InMemKVDBIterator struct { db *InMemKVDB start []byte end []byte pos int open bool } func (i *InMemKVDBIterator) Key() []byte { if !i.open { return nil } i.db.storeMx.Lock() if _, ok := i.db.store[i.db.sortedKeys[i.pos]]; !ok { return nil } i.db.storeMx.Unlock() return []byte(i.db.sortedKeys[i.pos]) } func (i *InMemKVDBIterator) First() bool { if !i.open { return false } i.db.storeMx.Lock() found := false idx := sort.SearchStrings(i.db.sortedKeys, string(i.start)) final := sort.SearchStrings(i.db.sortedKeys, string(i.end)) if idx < final { i.pos = idx found = true } i.db.storeMx.Unlock() return found } func (i *InMemKVDBIterator) Next() bool { if !i.open { return false } i.db.storeMx.Lock() found := false if _, ok := i.db.store[i.db.sortedKeys[i.pos]]; ok { final := sort.SearchStrings(i.db.sortedKeys, string(i.end)) if i.pos < final { i.pos = i.pos + 1 found = true } } i.db.storeMx.Unlock() return found } func (i *InMemKVDBIterator) Prev() bool { if !i.open { return false } i.db.storeMx.Lock() found := false if _, ok := i.db.store[i.db.sortedKeys[i.pos]]; ok { start := sort.SearchStrings(i.db.sortedKeys, string(i.start)) if i.pos-1 > start { i.pos = i.pos - 1 found = true } } i.db.storeMx.Unlock() return found } func (i *InMemKVDBIterator) Valid() bool { if !i.open { return false } i.db.storeMx.Lock() start := sort.SearchStrings(i.db.sortedKeys, string(i.start)) final := sort.SearchStrings(i.db.sortedKeys, string(i.end)) i.db.storeMx.Unlock() return i.pos < final && i.pos >= start } func (i *InMemKVDBIterator) Value() []byte { if !i.open { return nil } i.db.storeMx.Lock() value := i.db.store[i.db.sortedKeys[i.pos]] i.db.storeMx.Unlock() return value } func (i *InMemKVDBIterator) Close() error { if !i.open { return errors.New("already closed iterator") } i.open = false return nil } func (i *InMemKVDBIterator) SeekLT(lt []byte) bool { if !i.open { return false } i.db.storeMx.Lock() found := false if _, ok := i.db.store[i.db.sortedKeys[i.pos]]; ok { idx := sort.SearchStrings(i.db.sortedKeys, string(lt)) start := sort.SearchStrings(i.db.sortedKeys, string(i.start)) if idx >= start { i.pos = idx + 1 found = true } } i.db.storeMx.Unlock() return found } func (t *InMemKVDBTransaction) Set(key []byte, value []byte) error { if !t.db.open { return errors.New("inmem db closed") } t.changes = append(t.changes, InMemKVDBOperation{ op: SetOperation, key: key, value: value, }) return nil } func (t *InMemKVDBTransaction) Commit() error { if !t.db.open { return errors.New("inmem db closed") } var err error loop: for _, op := range t.changes { switch op.op { case SetOperation: err = t.db.Set(op.key, op.value) if err != nil { break loop } case DeleteOperation: err = t.db.Delete(op.key) if err != nil { break loop } } } return err } func (t *InMemKVDBTransaction) Delete(key []byte) error { if !t.db.open { return errors.New("inmem db closed") } t.changes = append(t.changes, InMemKVDBOperation{ op: DeleteOperation, key: key, }) return nil } func (t *InMemKVDBTransaction) Abort() error { return nil } func NewInMemKVDB() *InMemKVDB { return &InMemKVDB{ open: true, store: map[string][]byte{}, sortedKeys: []string{}, } } func (d *InMemKVDB) Get(key []byte) ([]byte, io.Closer, error) { if !d.open { return nil, nil, errors.New("inmem db closed") } d.storeMx.Lock() b, ok := d.store[string(key)] d.storeMx.Unlock() if !ok { return nil, nil, pebble.ErrNotFound } return b, io.NopCloser(nil), nil } func (d *InMemKVDB) Set(key, value []byte) error { if !d.open { return errors.New("inmem db closed") } d.storeMx.Lock() _, ok := d.store[string(key)] if !ok { i := sort.SearchStrings(d.sortedKeys, string(key)) if len(d.sortedKeys) > i { d.sortedKeys = append(d.sortedKeys[:i+1], d.sortedKeys[i:]...) d.sortedKeys[i] = string(key) } else { d.sortedKeys = append(d.sortedKeys, string(key)) } } d.store[string(key)] = value d.storeMx.Unlock() return nil } func (d *InMemKVDB) Delete(key []byte) error { if !d.open { return errors.New("inmem db closed") } d.storeMx.Lock() _, ok := d.store[string(key)] if ok { i := sort.SearchStrings(d.sortedKeys, string(key)) if len(d.sortedKeys)-1 > i { d.sortedKeys = append(d.sortedKeys[:i], d.sortedKeys[i+1:]...) } else { d.sortedKeys = d.sortedKeys[:i] } } delete(d.store, string(key)) d.storeMx.Unlock() return nil } func (d *InMemKVDB) NewBatch() Transaction { if !d.open { return nil } id := rand.Int() return &InMemKVDBTransaction{ id: id, db: d, changes: []InMemKVDBOperation{}, } } func (d *InMemKVDB) NewIter(lowerBound []byte, upperBound []byte) (Iterator, error) { if !d.open { return nil, errors.New("inmem db closed") } return &InMemKVDBIterator{ open: true, db: d, start: lowerBound, end: upperBound, pos: -1, }, nil } func (d *InMemKVDB) Compact(start, end []byte, parallelize bool) error { if !d.open { return errors.New("inmem db closed") } return nil } func (d *InMemKVDB) Close() error { if !d.open { return errors.New("inmem db closed") } d.open = false return nil } func (d *InMemKVDB) DeleteRange(start, end []byte) error { if !d.open { return errors.New("inmem db closed") } iter, err := d.NewIter(start, end) if err != nil { return err } for iter.First(); iter.Valid(); iter.Next() { err = d.Delete(iter.Key()) if err != nil { return err } } return nil } var _ KVDB = (*InMemKVDB)(nil)