package main import ( "github.com/jmhodges/levigo" "github.com/matttproud/prometheus/coding" "io" ) type LevigoCloser interface { Close() error } type LevigoPersistence struct { cache *levigo.Cache filterPolicy *levigo.FilterPolicy options *levigo.Options storage *levigo.DB readOptions *levigo.ReadOptions writeOptions *levigo.WriteOptions } func NewLevigoPersistence(storageRoot string, cacheCapacity, bitsPerBloomFilterEncoded int) (*LevigoPersistence, error) { options := levigo.NewOptions() options.SetCreateIfMissing(true) options.SetParanoidChecks(true) cache := levigo.NewLRUCache(cacheCapacity) options.SetCache(cache) filterPolicy := levigo.NewBloomFilter(bitsPerBloomFilterEncoded) options.SetFilterPolicy(filterPolicy) storage, openErr := levigo.Open(storageRoot, options) readOptions := levigo.NewReadOptions() writeOptions := levigo.NewWriteOptions() writeOptions.SetSync(true) emission := &LevigoPersistence{ cache: cache, filterPolicy: filterPolicy, options: options, readOptions: readOptions, storage: storage, writeOptions: writeOptions, } return emission, openErr } func (l *LevigoPersistence) Close() error { if l.storage != nil { l.storage.Close() } defer func() { if l.filterPolicy != nil { l.filterPolicy.Close() } }() defer func() { if l.cache != nil { l.cache.Close() } }() defer func() { if l.options != nil { l.options.Close() } }() defer func() { if l.readOptions != nil { l.readOptions.Close() } }() defer func() { if l.writeOptions != nil { l.writeOptions.Close() } }() return nil } func (l *LevigoPersistence) Get(value coding.Encoder) ([]byte, error) { var key []byte var keyError error if key, keyError = value.Encode(); keyError == nil { return l.storage.Get(l.readOptions, key) } return nil, keyError } func (l *LevigoPersistence) Has(value coding.Encoder) (bool, error) { if value, getError := l.Get(value); getError != nil { return false, getError } else if value == nil { return false, nil } return true, nil } func (l *LevigoPersistence) Drop(value coding.Encoder) error { var key []byte var keyError error if key, keyError = value.Encode(); keyError == nil { if deleteError := l.storage.Delete(l.writeOptions, key); deleteError != nil { return deleteError } return nil } return keyError } func (l *LevigoPersistence) Put(key, value coding.Encoder) error { var keyEncoded []byte var keyError error if keyEncoded, keyError = key.Encode(); keyError == nil { if valueEncoded, valueError := value.Encode(); valueError == nil { if putError := l.storage.Put(l.writeOptions, keyEncoded, valueEncoded); putError != nil { return putError } } else { return valueError } return nil } return keyError } func (l *LevigoPersistence) GetAll() ([]Pair, error) { snapshot := l.storage.NewSnapshot() defer l.storage.ReleaseSnapshot(snapshot) readOptions := levigo.NewReadOptions() defer readOptions.Close() readOptions.SetSnapshot(snapshot) iterator := l.storage.NewIterator(readOptions) defer iterator.Close() iterator.SeekToFirst() result := make([]Pair, 0) for iterator := iterator; iterator.Valid(); iterator.Next() { result = append(result, Pair{Left: iterator.Key(), Right: iterator.Value()}) } iteratorError := iterator.GetError() if iteratorError != nil { return nil, iteratorError } return result, nil } type iteratorCloser struct { iterator *levigo.Iterator readOptions *levigo.ReadOptions snapshot *levigo.Snapshot storage *levigo.DB } func (i *iteratorCloser) Close() error { defer func() { if i.storage != nil { if i.snapshot != nil { i.storage.ReleaseSnapshot(i.snapshot) } } }() defer func() { if i.iterator != nil { i.iterator.Close() } }() defer func() { if i.readOptions != nil { i.readOptions.Close() } }() return nil } func (l *LevigoPersistence) GetIterator() (*levigo.Iterator, io.Closer, error) { snapshot := l.storage.NewSnapshot() readOptions := levigo.NewReadOptions() readOptions.SetSnapshot(snapshot) iterator := l.storage.NewIterator(readOptions) closer := &iteratorCloser{ iterator: iterator, readOptions: readOptions, snapshot: snapshot, storage: l.storage, } return iterator, closer, nil }