github.com/observiq/carbon@v0.9.11-0.20200820160507-1b872e368a5e/operator/helper/persister.go (about)

     1  package helper
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/observiq/carbon/operator"
     7  	"go.etcd.io/bbolt"
     8  )
     9  
    10  // Persister is a helper used to persist data
    11  type Persister interface {
    12  	Get(string) []byte
    13  	Set(string, []byte)
    14  	Sync() error
    15  	Load() error
    16  }
    17  
    18  // ScopedBBoltPersister is a persister that uses a database for the backend
    19  type ScopedBBoltPersister struct {
    20  	scope    []byte
    21  	db       operator.Database
    22  	cache    map[string][]byte
    23  	cacheMux sync.Mutex
    24  }
    25  
    26  // NewScopedDBPersister returns a new ScopedBBoltPersister
    27  func NewScopedDBPersister(db operator.Database, scope string) *ScopedBBoltPersister {
    28  	return &ScopedBBoltPersister{
    29  		scope: []byte(scope),
    30  		db:    db,
    31  		cache: make(map[string][]byte),
    32  	}
    33  }
    34  
    35  // Get retrieves a key from the cache
    36  func (p *ScopedBBoltPersister) Get(key string) []byte {
    37  	p.cacheMux.Lock()
    38  	defer p.cacheMux.Unlock()
    39  	return p.cache[key]
    40  }
    41  
    42  // Set saves a key in the cache
    43  func (p *ScopedBBoltPersister) Set(key string, val []byte) {
    44  	p.cacheMux.Lock()
    45  	p.cache[key] = val
    46  	p.cacheMux.Unlock()
    47  }
    48  
    49  // OffsetsBucket is the scope provided to offset persistence
    50  var OffsetsBucket = []byte(`offsets`)
    51  
    52  // Sync saves the cache to the backend, ensuring values are
    53  // safely written to disk before returning
    54  func (p *ScopedBBoltPersister) Sync() error {
    55  	return p.db.Update(func(tx *bbolt.Tx) error {
    56  		offsetBucket, err := tx.CreateBucketIfNotExists(OffsetsBucket)
    57  		if err != nil {
    58  			return err
    59  		}
    60  
    61  		bucket, err := offsetBucket.CreateBucketIfNotExists(p.scope)
    62  		if err != nil {
    63  			return err
    64  		}
    65  
    66  		p.cacheMux.Lock()
    67  		for k, v := range p.cache {
    68  			err := bucket.Put([]byte(k), v)
    69  			if err != nil {
    70  				return err
    71  			}
    72  		}
    73  		p.cacheMux.Unlock()
    74  
    75  		return nil
    76  	})
    77  }
    78  
    79  // Load populates the cache with the values from the database,
    80  // overwriting anything currently in the cache.
    81  func (p *ScopedBBoltPersister) Load() error {
    82  	p.cacheMux.Lock()
    83  	defer p.cacheMux.Unlock()
    84  	p.cache = make(map[string][]byte)
    85  
    86  	return p.db.Update(func(tx *bbolt.Tx) error {
    87  		offsetBucket, err := tx.CreateBucketIfNotExists(OffsetsBucket)
    88  		if err != nil {
    89  			return err
    90  		}
    91  
    92  		bucket, err := offsetBucket.CreateBucketIfNotExists(p.scope)
    93  		if err != nil {
    94  			return err
    95  		}
    96  
    97  		return bucket.ForEach(func(k, v []byte) error {
    98  			p.cache[string(k)] = v
    99  			return nil
   100  		})
   101  	})
   102  }