github.com/grafana/pyroscope@v1.18.0/pkg/segmentwriter/client/distributor/placement/adaptiveplacement/store.go (about)

     1  package adaptiveplacement
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"errors"
     7  	"time"
     8  
     9  	"github.com/thanos-io/objstore"
    10  
    11  	"github.com/grafana/pyroscope/pkg/segmentwriter/client/distributor/placement/adaptiveplacement/adaptive_placementpb"
    12  )
    13  
    14  const (
    15  	pathRoot      = "adaptive_placement/"
    16  	rulesFilePath = pathRoot + "placement_rules.binpb"
    17  	statsFilePath = pathRoot + "placement_stats.binpb"
    18  )
    19  
    20  var (
    21  	ErrRulesNotFound = errors.New("placement rules not found")
    22  	ErrStatsNotFound = errors.New("placement stats not found")
    23  )
    24  
    25  type StoreReader interface {
    26  	LoadRules(context.Context) (*adaptive_placementpb.PlacementRules, error)
    27  	LoadStats(context.Context) (*adaptive_placementpb.DistributionStats, error)
    28  }
    29  
    30  type StoreWriter interface {
    31  	StoreRules(context.Context, *adaptive_placementpb.PlacementRules) error
    32  	StoreStats(context.Context, *adaptive_placementpb.DistributionStats) error
    33  }
    34  
    35  type Store interface {
    36  	StoreReader
    37  	StoreWriter
    38  }
    39  
    40  type BucketStore struct{ bucket objstore.Bucket }
    41  
    42  func NewStore(bucket objstore.Bucket) *BucketStore { return &BucketStore{bucket: bucket} }
    43  
    44  func (s *BucketStore) LoadRules(ctx context.Context) (*adaptive_placementpb.PlacementRules, error) {
    45  	var rules adaptive_placementpb.PlacementRules
    46  	if err := s.get(ctx, rulesFilePath, &rules); err != nil {
    47  		if s.bucket.IsObjNotFoundErr(err) {
    48  			return nil, ErrRulesNotFound
    49  		}
    50  		return nil, err
    51  	}
    52  	return &rules, nil
    53  }
    54  
    55  func (s *BucketStore) LoadStats(ctx context.Context) (*adaptive_placementpb.DistributionStats, error) {
    56  	var stats adaptive_placementpb.DistributionStats
    57  	if err := s.get(ctx, statsFilePath, &stats); err != nil {
    58  		if s.bucket.IsObjNotFoundErr(err) {
    59  			return nil, ErrStatsNotFound
    60  		}
    61  		return nil, err
    62  	}
    63  	return &stats, nil
    64  }
    65  
    66  func (s *BucketStore) StoreRules(ctx context.Context, rules *adaptive_placementpb.PlacementRules) error {
    67  	return s.put(ctx, rulesFilePath, rules)
    68  }
    69  
    70  func (s *BucketStore) StoreStats(ctx context.Context, stats *adaptive_placementpb.DistributionStats) error {
    71  	return s.put(ctx, statsFilePath, stats)
    72  }
    73  
    74  type vtProtoMessage interface {
    75  	UnmarshalVT([]byte) error
    76  	MarshalVT() ([]byte, error)
    77  }
    78  
    79  func (s *BucketStore) get(ctx context.Context, name string, m vtProtoMessage) error {
    80  	r, err := s.bucket.Get(ctx, name)
    81  	if err != nil {
    82  		return err
    83  	}
    84  	defer func() {
    85  		_ = r.Close()
    86  	}()
    87  	var buf bytes.Buffer
    88  	if _, err = buf.ReadFrom(r); err != nil {
    89  		return err
    90  	}
    91  	return m.UnmarshalVT(buf.Bytes())
    92  }
    93  
    94  func (s *BucketStore) put(ctx context.Context, name string, m vtProtoMessage) error {
    95  	b, err := m.MarshalVT()
    96  	if err != nil {
    97  		return err
    98  	}
    99  	return s.bucket.Upload(ctx, name, bytes.NewReader(b))
   100  }
   101  
   102  // EmptyStore is a Store implementation that always returns
   103  // empty rules and stats, and doesn't store anything.
   104  type EmptyStore struct{}
   105  
   106  func NewEmptyStore() *EmptyStore { return new(EmptyStore) }
   107  
   108  func (e *EmptyStore) LoadRules(context.Context) (*adaptive_placementpb.PlacementRules, error) {
   109  	return &adaptive_placementpb.PlacementRules{CreatedAt: time.Now().UnixNano()}, nil
   110  }
   111  
   112  func (e *EmptyStore) LoadStats(context.Context) (*adaptive_placementpb.DistributionStats, error) {
   113  	return &adaptive_placementpb.DistributionStats{CreatedAt: time.Now().UnixNano()}, nil
   114  }
   115  
   116  func (e *EmptyStore) StoreRules(context.Context, *adaptive_placementpb.PlacementRules) error {
   117  	return nil
   118  }
   119  
   120  func (e *EmptyStore) StoreStats(context.Context, *adaptive_placementpb.DistributionStats) error {
   121  	return nil
   122  }