github.com/grafana/pyroscope@v1.18.0/pkg/metastore/compaction/scheduler/store/job_state_store.go (about)

     1  package store
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  
     7  	"go.etcd.io/bbolt"
     8  
     9  	"github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1/raft_log"
    10  	"github.com/grafana/pyroscope/pkg/iter"
    11  	"github.com/grafana/pyroscope/pkg/metastore/store"
    12  )
    13  
    14  var jobStateBucketName = []byte("compaction_job_state")
    15  
    16  var ErrInvalidJobState = errors.New("invalid job state entry")
    17  
    18  type JobStateStore struct{ bucketName []byte }
    19  
    20  func NewJobStateStore() *JobStateStore {
    21  	return &JobStateStore{bucketName: jobStateBucketName}
    22  }
    23  
    24  func (s JobStateStore) CreateBuckets(tx *bbolt.Tx) error {
    25  	_, err := tx.CreateBucketIfNotExists(s.bucketName)
    26  	return err
    27  }
    28  
    29  func (s JobStateStore) GetJobState(tx *bbolt.Tx, name string) (*raft_log.CompactionJobState, error) {
    30  	b := tx.Bucket(s.bucketName).Get([]byte(name))
    31  	if b == nil {
    32  		return nil, fmt.Errorf("loading job state %s: %w", name, store.ErrNotFound)
    33  	}
    34  	var v raft_log.CompactionJobState
    35  	if err := v.UnmarshalVT(b); err != nil {
    36  		return nil, fmt.Errorf("%w: %v", ErrInvalidJobState, err)
    37  	}
    38  	return &v, nil
    39  }
    40  
    41  func (s JobStateStore) StoreJobState(tx *bbolt.Tx, state *raft_log.CompactionJobState) error {
    42  	v, _ := state.MarshalVT()
    43  	return tx.Bucket(s.bucketName).Put([]byte(state.Name), v)
    44  }
    45  
    46  func (s JobStateStore) DeleteJobState(tx *bbolt.Tx, name string) error {
    47  	return tx.Bucket(s.bucketName).Delete([]byte(name))
    48  }
    49  
    50  func (s JobStateStore) ListEntries(tx *bbolt.Tx) iter.Iterator[*raft_log.CompactionJobState] {
    51  	return newJobEntriesIterator(tx.Bucket(s.bucketName))
    52  }
    53  
    54  type jobEntriesIterator struct {
    55  	iter *store.CursorIterator
    56  	cur  *raft_log.CompactionJobState
    57  	err  error
    58  }
    59  
    60  func newJobEntriesIterator(bucket *bbolt.Bucket) *jobEntriesIterator {
    61  	return &jobEntriesIterator{iter: store.NewCursorIter(bucket.Cursor())}
    62  }
    63  
    64  func (x *jobEntriesIterator) Next() bool {
    65  	if x.err != nil || !x.iter.Next() {
    66  		return false
    67  	}
    68  	e := x.iter.At()
    69  	var s raft_log.CompactionJobState
    70  	x.err = s.UnmarshalVT(e.Value)
    71  	if x.err != nil {
    72  		x.err = fmt.Errorf("%w: %v", ErrInvalidJobState, x.err)
    73  		return false
    74  	}
    75  	x.cur = &s
    76  	return true
    77  }
    78  
    79  func (x *jobEntriesIterator) At() *raft_log.CompactionJobState { return x.cur }
    80  
    81  func (x *jobEntriesIterator) Close() error { return x.iter.Close() }
    82  
    83  func (x *jobEntriesIterator) Err() error {
    84  	if err := x.iter.Err(); err != nil {
    85  		return err
    86  	}
    87  	return x.err
    88  }