github.com/git-lfs/git-lfs@v2.5.2+incompatible/tools/kv/keyvaluestore.go (about) 1 package kv 2 3 import ( 4 "encoding/gob" 5 "fmt" 6 "io" 7 "os" 8 "sync" 9 ) 10 11 // Store provides an in-memory key/value store which is persisted to 12 // a file. The file handle itself is not kept locked for the duration; it is 13 // only locked during load and save, to make it concurrency friendly. When 14 // saving, the store uses optimistic locking to determine whether the db on disk 15 // has been modified by another process; in which case it loads the latest 16 // version and re-applies modifications made during this session. This means 17 // the Lost Update db concurrency issue is possible; so don't use this if you 18 // need more DB integrity than Read Committed isolation levels. 19 type Store struct { 20 // Locks the entire store 21 mu sync.RWMutex 22 filename string 23 log []change 24 25 // This is the persistent data 26 // version for optimistic locking, this field is incremented with every Save() 27 version int64 28 db map[string]interface{} 29 } 30 31 // Type of operation; set or remove 32 type operation int 33 34 const ( 35 // Set a value for a key 36 setOperation = operation(iota) 37 // Removed a value for a key 38 removeOperation = operation(iota) 39 ) 40 41 type change struct { 42 op operation 43 key string 44 value interface{} 45 } 46 47 // NewStore creates a new key/value store and initialises it with contents from 48 // the named file, if it exists 49 func NewStore(filepath string) (*Store, error) { 50 kv := &Store{filename: filepath, db: make(map[string]interface{})} 51 return kv, kv.loadAndMergeIfNeeded() 52 } 53 54 // Set updates the key/value store in memory 55 // Changes are not persisted until you call Save() 56 func (k *Store) Set(key string, value interface{}) { 57 k.mu.Lock() 58 defer k.mu.Unlock() 59 60 k.db[key] = value 61 k.logChange(setOperation, key, value) 62 } 63 64 // Remove removes the key and its value from the store in memory 65 // Changes are not persisted until you call Save() 66 func (k *Store) Remove(key string) { 67 k.mu.Lock() 68 defer k.mu.Unlock() 69 70 delete(k.db, key) 71 k.logChange(removeOperation, key, nil) 72 } 73 74 // RemoveAll removes all entries from the store 75 // These changes are not persisted until you call Save() 76 func (k *Store) RemoveAll() { 77 k.mu.Lock() 78 defer k.mu.Unlock() 79 80 // Log all changes 81 for key, _ := range k.db { 82 k.logChange(removeOperation, key, nil) 83 } 84 k.db = make(map[string]interface{}) 85 } 86 87 // Visit walks through the entire store via a function; return false from 88 // your visitor function to halt the walk 89 func (k *Store) Visit(cb func(string, interface{}) bool) { 90 // Read-only lock 91 k.mu.RLock() 92 defer k.mu.RUnlock() 93 94 for k, v := range k.db { 95 if !cb(k, v) { 96 break 97 } 98 } 99 } 100 101 // Append a change to the log; mutex must already be locked 102 func (k *Store) logChange(op operation, key string, value interface{}) { 103 k.log = append(k.log, change{op, key, value}) 104 } 105 106 // Get retrieves a value from the store, or nil if it is not present 107 func (k *Store) Get(key string) interface{} { 108 // Read-only lock 109 k.mu.RLock() 110 defer k.mu.RUnlock() 111 112 // zero value of interface{} is nil so this does what we want 113 return k.db[key] 114 } 115 116 // Save persists the changes made to disk 117 // If any changes have been written by other code they will be merged 118 func (k *Store) Save() error { 119 k.mu.Lock() 120 defer k.mu.Unlock() 121 122 // Short-circuit if we have no changes 123 if len(k.log) == 0 { 124 return nil 125 } 126 127 // firstly peek at version; open read/write to keep lock between check & write 128 f, err := os.OpenFile(k.filename, os.O_RDWR|os.O_CREATE, 0664) 129 if err != nil { 130 return err 131 } 132 133 defer f.Close() 134 135 // Only try to merge if > 0 bytes, ignore empty files (decoder will fail) 136 if stat, _ := f.Stat(); stat.Size() > 0 { 137 k.loadAndMergeReaderIfNeeded(f) 138 // Now we overwrite the file 139 f.Seek(0, os.SEEK_SET) 140 f.Truncate(0) 141 } 142 143 k.version++ 144 145 enc := gob.NewEncoder(f) 146 if err := enc.Encode(k.version); err != nil { 147 return fmt.Errorf("Error while writing version data to %v: %v", k.filename, err) 148 } 149 if err := enc.Encode(k.db); err != nil { 150 return fmt.Errorf("Error while writing new key/value data to %v: %v", k.filename, err) 151 } 152 // Clear log now that it's saved 153 k.log = nil 154 155 return nil 156 } 157 158 // Reads as little as possible from the passed in file to determine if the 159 // contents are different from the version already held. If so, reads the 160 // contents and merges with any outstanding changes. If not, stops early without 161 // reading the rest of the file 162 func (k *Store) loadAndMergeIfNeeded() error { 163 stat, err := os.Stat(k.filename) 164 if err != nil { 165 if os.IsNotExist(err) { 166 return nil // missing is OK 167 } 168 return err 169 } 170 // Do nothing if empty file 171 if stat.Size() == 0 { 172 return nil 173 } 174 175 f, err := os.OpenFile(k.filename, os.O_RDONLY, 0664) 176 if err == nil { 177 defer f.Close() 178 return k.loadAndMergeReaderIfNeeded(f) 179 } else { 180 return err 181 } 182 } 183 184 // As loadAndMergeIfNeeded but lets caller decide how to manage file handles 185 func (k *Store) loadAndMergeReaderIfNeeded(f io.Reader) error { 186 var versionOnDisk int64 187 // Decode *only* the version field to check whether anyone else has 188 // modified the db; gob serializes structs in order so it will always be 1st 189 dec := gob.NewDecoder(f) 190 err := dec.Decode(&versionOnDisk) 191 if err != nil { 192 return fmt.Errorf("Problem checking version of key/value data from %v: %v", k.filename, err) 193 } 194 // Totally uninitialised Version == 0, saved versions are always >=1 195 if versionOnDisk != k.version { 196 // Reload data & merge 197 var dbOnDisk map[string]interface{} 198 err = dec.Decode(&dbOnDisk) 199 if err != nil { 200 return fmt.Errorf("Problem reading updated key/value data from %v: %v", k.filename, err) 201 } 202 k.reapplyChanges(dbOnDisk) 203 k.version = versionOnDisk 204 } 205 return nil 206 } 207 208 // reapplyChanges replays the changes made since the last load onto baseDb 209 // and stores the result as our own DB 210 func (k *Store) reapplyChanges(baseDb map[string]interface{}) { 211 for _, change := range k.log { 212 switch change.op { 213 case setOperation: 214 baseDb[change.key] = change.value 215 case removeOperation: 216 delete(baseDb, change.key) 217 } 218 } 219 // Note, log is not cleared here, that only happens on Save since it's a 220 // list of unsaved changes 221 k.db = baseDb 222 223 } 224 225 // RegisterTypeForStorage registers a custom type (e.g. a struct) for 226 // use in the key value store. This is necessary if you intend to pass custom 227 // structs to Store.Set() rather than primitive types. 228 func RegisterTypeForStorage(val interface{}) { 229 gob.Register(val) 230 }