git-lfs/tools/keyvaluestore.go

package tools

import (
	"encoding/gob"
	"fmt"
	"os"
	"sync"
)

// KeyValueStore provides an in-memory key/value store which is persisted to
// a file. The file handle itself is not kept locked for the duration; it is
// only locked during load and save, to make it concurrency friendly. When
// saving, the store uses optimistic locking to determine whether the db on disk
// has been modified by another process; in which case it loads the latest
// version and re-applies modifications made during this session. This means
// the Lost Update db concurrency issue is possible; so don't use this if you
// need more DB integrity than Read Committed isolation levels.
type KeyValueStore struct {
	mu       sync.RWMutex
	filename string
	data     *keyValueStoreData
	log      []keyValueChange
}

type keyValueOperation int

const (
	// Set a value for a key
	keyValueSetOperation = keyValueOperation(iota)
	// Removed a value for a key
	keyValueRemoveOperation = keyValueOperation(iota)
)

type keyValueChange struct {
	operation keyValueOperation
	key       string
	value     interface{}
}

// Actual data that we persist
type keyValueStoreData struct {
	// version for optimistic locking, this field is incremented with every Save()
	// MUST BE FIRST in the struct, gob serializes in order and we need to peek
	Version int64
	Db      map[string]interface{}
}

// NewKeyValueStore creates a new store and initialises it with contents from
// the named file, if it exists
func NewKeyValueStore(filepath string) (*KeyValueStore, error) {

	var kvdata *keyValueStoreData
	stat, _ := os.Stat(filepath)
	f, err := os.OpenFile(filepath, os.O_RDONLY, 0664)
	// Only try to load if file has some bytes in it
	if err == nil && stat.Size() > 0 {
		defer f.Close()
		dec := gob.NewDecoder(f)
		err := dec.Decode(&kvdata)
		if err != nil {
			return nil, fmt.Errorf("Problem loading key/value data from %v: %v", filepath, err)
		}
	} else if err != nil && !os.IsNotExist(err) {
		return nil, err
	}
	if kvdata == nil {
		kvdata = &keyValueStoreData{1, make(map[string]interface{}, 100)}
	}

	return &KeyValueStore{filename: filepath, data: kvdata}, nil

}

// Set updates the key/value store in memory
// Changes are not persisted until you call Save()
func (k *KeyValueStore) Set(key string, value interface{}) {
	k.mu.Lock()
	defer k.mu.Unlock()

	k.data.Db[key] = value
	k.log = append(k.log, keyValueChange{keyValueSetOperation, key, value})
}

// Remove removes the key and its value from the store in memory
// Changes are not persisted until you call Save()
func (k *KeyValueStore) Remove(key string) {
	k.mu.Lock()
	defer k.mu.Unlock()

	delete(k.data.Db, key)
	k.log = append(k.log, keyValueChange{keyValueRemoveOperation, key, nil})
}

// Get retrieves a value from the store, or nil if it is not present
func (k *KeyValueStore) Get(key string) interface{} {
	// Read-only lock
	k.mu.RLock()
	defer k.mu.RUnlock()

	// zero value of interface{} is nil so this does what we want
	return k.data.Db[key]
}

// Save persists the changes made to disk
func (k *KeyValueStore) Save() error {
	k.mu.Lock()
	defer k.mu.Unlock()

	// Short-circuit if we have no changes
	if len(k.log) == 0 {
		return nil
	}

	// firstly peek at version; open read/write to keep lock between check & write
	f, err := os.OpenFile(k.filename, os.O_RDWR|os.O_CREATE, 0664)
	if err != nil {
		return err
	}
	stat, _ := os.Stat(k.filename)

	defer f.Close()

	// Only try to peek if > 0 bytes, ignore empty files (decoder will fail)
	if stat.Size() > 0 {
		var versionOnDisk int64
		// Decode *only* the version field to check whether anyone else has
		// modified the db; gob serializes structs in order so it will always be 1st
		dec := gob.NewDecoder(f)
		err = dec.Decode(&versionOnDisk)
		if err != nil {
			return fmt.Errorf("Problem checking version of key/value data from %v: %v", k.filename, err)
		}
		if versionOnDisk != k.data.Version {
			// Reload data & merge
			var dbOnDisk map[string]interface{}
			err = dec.Decode(&dbOnDisk)
			if err != nil {
				return fmt.Errorf("Problem reading updated key/value data from %v: %v", k.filename, err)
			}
			k.reapplyChanges(dbOnDisk)
		}
		// Now we overwrite
		f.Seek(0, os.SEEK_SET)
	}

	k.data.Version++

	enc := gob.NewEncoder(f)
	err = enc.Encode(k.data)
	if err != nil {
		return fmt.Errorf("Error while writing new key/value data to %v: %v", k.filename, err)
	}
	// Clear log now that it's saved
	k.log = nil

	return nil
}

// reapplyChanges replays the changes made since the last load onto baseDb
// and stores the result as our own DB
func (k *KeyValueStore) reapplyChanges(baseDb map[string]interface{}) {
	for _, change := range k.log {
		switch change.operation {
		case keyValueSetOperation:
			baseDb[change.key] = change.value
		case keyValueRemoveOperation:
			delete(baseDb, change.key)
		}
	}
	k.data.Db = baseDb

}