github.com/grafana/pyroscope@v1.18.0/pkg/settings/bucket.go (about)

     1  package settings
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"encoding/json"
     7  	"slices"
     8  	"strings"
     9  	"sync"
    10  
    11  	"github.com/pkg/errors"
    12  	"github.com/thanos-io/objstore"
    13  
    14  	settingsv1 "github.com/grafana/pyroscope/api/gen/proto/go/settings/v1"
    15  )
    16  
    17  var (
    18  	oldSettingErr    = errors.New("newer update already written")
    19  	settingsFilename = "tenant_settings.json"
    20  )
    21  
    22  // newMemoryStore will create a settings store with an in-memory bucket.
    23  func newMemoryStore() store {
    24  	return newBucketStore(objstore.NewInMemBucket())
    25  }
    26  
    27  // NewBucketStore will create a settings store with an objstore bucket.
    28  func newBucketStore(bucket objstore.Bucket) store {
    29  	return &bucketStore{
    30  		store:  make(map[string]map[string]*settingsv1.Setting),
    31  		bucket: bucket,
    32  	}
    33  }
    34  
    35  type bucketStore struct {
    36  	rw sync.Mutex
    37  
    38  	// store is kv pairs, indexed first by tenant id.
    39  	store map[string]map[string]*settingsv1.Setting
    40  
    41  	// bucket is an object store bucket.
    42  	bucket objstore.Bucket
    43  }
    44  
    45  func (s *bucketStore) Get(ctx context.Context, tenantID string) ([]*settingsv1.Setting, error) {
    46  	s.rw.Lock()
    47  	defer s.rw.Unlock()
    48  
    49  	err := s.unsafeLoad(ctx)
    50  	if err != nil {
    51  		return nil, err
    52  	}
    53  
    54  	tenantSettings := s.store[tenantID]
    55  
    56  	settings := make([]*settingsv1.Setting, 0, len(s.store[tenantID]))
    57  	for _, setting := range tenantSettings {
    58  		settings = append(settings, setting)
    59  	}
    60  
    61  	slices.SortFunc(settings, func(a, b *settingsv1.Setting) int {
    62  		return strings.Compare(a.Name, b.Name)
    63  	})
    64  	return settings, nil
    65  }
    66  
    67  func (s *bucketStore) Set(ctx context.Context, tenantID string, setting *settingsv1.Setting) (*settingsv1.Setting, error) {
    68  	s.rw.Lock()
    69  	defer s.rw.Unlock()
    70  
    71  	err := s.unsafeLoad(ctx)
    72  	if err != nil {
    73  		return nil, err
    74  	}
    75  
    76  	_, ok := s.store[tenantID]
    77  	if !ok {
    78  		s.store[tenantID] = make(map[string]*settingsv1.Setting, 1)
    79  	}
    80  
    81  	oldSetting, ok := s.store[tenantID][setting.Name]
    82  	if ok && oldSetting.ModifiedAt > setting.ModifiedAt {
    83  		return nil, errors.Wrapf(oldSettingErr, "failed to update %s", setting.Name)
    84  	}
    85  	s.store[tenantID][setting.Name] = setting
    86  
    87  	err = s.unsafeFlush(ctx)
    88  	if err != nil {
    89  		return nil, err
    90  	}
    91  
    92  	return setting, nil
    93  }
    94  
    95  func (s *bucketStore) Delete(ctx context.Context, tenantID string, name string, modifiedAtMs int64) error {
    96  	s.rw.Lock()
    97  	defer s.rw.Unlock()
    98  
    99  	err := s.unsafeLoad(ctx)
   100  	if err != nil {
   101  		return err
   102  	}
   103  
   104  	tenantSettings, ok := s.store[tenantID]
   105  	if !ok {
   106  		return nil
   107  	}
   108  
   109  	setting, ok := tenantSettings[name]
   110  	if !ok {
   111  		return nil
   112  	}
   113  
   114  	if setting.ModifiedAt > modifiedAtMs {
   115  		return errors.Wrapf(oldSettingErr, "failed to delete %s", name)
   116  	}
   117  
   118  	delete(tenantSettings, name)
   119  
   120  	err = s.unsafeFlush(ctx)
   121  	if err != nil {
   122  		return err
   123  	}
   124  
   125  	return nil
   126  }
   127  
   128  func (s *bucketStore) Close() error {
   129  	return s.bucket.Close()
   130  }
   131  
   132  // unsafeFlush will flush the store to object storage. This is not thread-safe,
   133  // the store's write mutex should be acquired first.
   134  func (s *bucketStore) unsafeFlush(ctx context.Context) error {
   135  	data, err := json.Marshal(s.store)
   136  	if err != nil {
   137  		return err
   138  	}
   139  
   140  	err = s.bucket.Upload(ctx, settingsFilename, bytes.NewReader(data))
   141  	if err != nil {
   142  		return err
   143  	}
   144  	return nil
   145  }
   146  
   147  // unsafeLoad will read the store in object storage into memory, if it exists.
   148  // This is not thread-safe, the store's write mutex should be acquired first.
   149  func (s *bucketStore) unsafeLoad(ctx context.Context) error {
   150  	reader, err := s.bucket.Get(ctx, settingsFilename)
   151  	if err != nil {
   152  		if s.bucket.IsObjNotFoundErr(err) {
   153  			// It is OK if we don't find the file.
   154  			return nil
   155  		}
   156  		return err
   157  	}
   158  
   159  	err = json.NewDecoder(reader).Decode(&s.store)
   160  	if err != nil {
   161  		reader.Close()
   162  		return err
   163  	}
   164  
   165  	err = reader.Close()
   166  	if err != nil {
   167  		return err
   168  	}
   169  	return nil
   170  }