1 package keystore 2 3 import ( 4 "io" 5 "sort" 6 "sync" 7 8 "vimagination.zapto.org/byteio" 9 "vimagination.zapto.org/memio" 10 ) 11 12 // MemStore implements Store and does so entirely in memory 13 type MemStore struct { 14 mu sync.RWMutex 15 data map[string]memio.Buffer 16 } 17 18 // NewMemStore creates a new memory-backed key-value store 19 func NewMemStore() *MemStore { 20 ms := new(MemStore) 21 ms.init() 22 return ms 23 } 24 25 func (ms *MemStore) init() { 26 ms.data = make(map[string]memio.Buffer) 27 } 28 29 // Get retrieves the key data from memory 30 func (ms *MemStore) Get(key string, r io.ReaderFrom) error { 31 d := ms.get(key) 32 if d == nil { 33 return ErrUnknownKey 34 } 35 _, err := r.ReadFrom(&d) 36 return err 37 } 38 39 // GetAll retrieves data for all of the keys given. Useful to reduce locking. 40 // Unknown Key errors are not returned, only errors from the ReaderFrom's 41 func (ms *MemStore) GetAll(data map[string]io.ReaderFrom) error { 42 var err error 43 ms.mu.RLock() 44 for k, d := range data { 45 buf, ok := ms.data[k] 46 if !ok { 47 continue 48 } 49 if _, err = d.ReadFrom(&buf); err != nil { 50 return err 51 } 52 } 53 ms.mu.RUnlock() 54 return err 55 } 56 57 func (ms *MemStore) get(key string) memio.Buffer { 58 ms.mu.RLock() 59 d := ms.data[key] 60 ms.mu.RUnlock() 61 return d 62 } 63 64 // Set stores the key data in memory 65 func (ms *MemStore) Set(key string, w io.WriterTo) error { 66 d := make(memio.Buffer, 0) 67 if _, err := w.WriteTo(&d); err != nil && err != io.EOF { 68 return err 69 } 70 ms.set(key, d) 71 return nil 72 } 73 74 // SetAll set data for all of the keys given. Useful to reduce locking. 75 // Will return the first error found, so may not set all data. 76 func (ms *MemStore) SetAll(data map[string]io.WriterTo) error { 77 var err error 78 ms.mu.Lock() 79 for k, d := range data { 80 var buf memio.Buffer 81 if _, err = d.WriteTo(&buf); err != nil { 82 break 83 } 84 ms.data[k] = buf 85 } 86 ms.mu.Unlock() 87 return err 88 } 89 90 func (ms *MemStore) set(key string, d memio.Buffer) { 91 ms.mu.Lock() 92 ms.data[key] = d 93 ms.mu.Unlock() 94 } 95 96 // Remove deletes the key data from memory 97 func (ms *MemStore) Remove(key string) error { 98 ms.mu.Lock() 99 _, ok := ms.data[key] 100 if !ok { 101 ms.mu.Unlock() 102 return ErrUnknownKey 103 } 104 delete(ms.data, key) 105 ms.mu.Unlock() 106 return nil 107 } 108 109 // RemoveAll will attempt to remove all keys given. It does not return an error 110 // if a key doesn't exist 111 func (ms *MemStore) RemoveAll(keys ...string) { 112 ms.mu.Lock() 113 for _, key := range keys { 114 delete(ms.data, key) 115 } 116 ms.mu.Unlock() 117 } 118 119 // WriteTo implements the io.WriterTo interface allowing a MemStore to be 120 // be stored in another Store 121 func (ms *MemStore) WriteTo(w io.Writer) (int64, error) { 122 lw := byteio.StickyLittleEndianWriter{Writer: w} 123 ms.mu.RLock() 124 for key, value := range ms.data { 125 lw.WriteStringX(key) 126 lw.WriteUintX(uint64(len(value))) 127 lw.Write(value) 128 } 129 ms.mu.RUnlock() 130 return lw.Count, lw.Err 131 } 132 133 // ReadFrom implements the io.ReaderFrom interface allowing a MemStore to be 134 // be retrieved in another Store 135 func (ms *MemStore) ReadFrom(r io.Reader) (int64, error) { 136 lr := byteio.StickyLittleEndianReader{Reader: r} 137 ms.mu.Lock() 138 for { 139 key := lr.ReadStringX() 140 if lr.Err == io.EOF { 141 lr.Err = nil 142 break 143 } 144 buf := make(memio.Buffer, lr.ReadUintX()) 145 lr.Read(buf) 146 if lr.Err != nil { 147 if lr.Err == io.EOF { 148 lr.Err = io.ErrUnexpectedEOF 149 } 150 break 151 } 152 ms.data[key] = buf 153 } 154 ms.mu.Unlock() 155 return lr.Count, lr.Err 156 } 157 158 // Keys returns a sorted slice of all of the keys 159 func (ms *MemStore) Keys() []string { 160 ms.mu.RLock() 161 s := make([]string, 0, len(ms.data)) 162 for key := range ms.data { 163 s = append(s, key) 164 } 165 ms.mu.RUnlock() 166 sort.Strings(s) 167 return s 168 } 169 170 // Exists returns true when the key exists within the store 171 func (ms *MemStore) Exists(key string) bool { 172 ms.mu.RLock() 173 _, ok := ms.data[key] 174 ms.mu.RUnlock() 175 return ok 176 } 177 178 // Rename moves data from an existing key to a new, unused key 179 func (ms *MemStore) Rename(oldkey, newkey string) error { 180 ms.mu.Lock() 181 var err error 182 if d, ok := ms.data[oldkey]; !ok { 183 err = ErrUnknownKey 184 } else if _, ok = ms.data[newkey]; ok { 185 err = ErrKeyExists 186 } else { 187 ms.data[newkey] = d 188 delete(ms.data, oldkey) 189 } 190 ms.mu.Unlock() 191 return err 192 }