github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/storage/labels/labels.go (about)

     1  package labels
     2  
     3  import (
     4  	"strings"
     5  
     6  	"github.com/dgraph-io/badger/v2"
     7  )
     8  
     9  type Labels struct {
    10  	db *badger.DB
    11  }
    12  
    13  func New(db *badger.DB) *Labels {
    14  	ll := &Labels{
    15  		db: db,
    16  	}
    17  	return ll
    18  }
    19  
    20  func (ll *Labels) PutLabels(labels map[string]string) error {
    21  	return ll.db.Update(func(txn *badger.Txn) error {
    22  		for k, v := range labels {
    23  			if err := txn.SetEntry(badger.NewEntry([]byte("l:"+k), nil)); err != nil {
    24  				return err
    25  			}
    26  			if err := txn.SetEntry(badger.NewEntry([]byte("v:"+k+":"+v), nil)); err != nil {
    27  				return err
    28  			}
    29  		}
    30  		return nil
    31  	})
    32  }
    33  
    34  func (ll *Labels) Put(key, val string) {
    35  	kk := "l:" + key
    36  	kv := "v:" + key + ":" + val
    37  	// ks := "h:" + key + ":" + val + ":" + stree
    38  	err := ll.db.Update(func(txn *badger.Txn) error {
    39  		return txn.SetEntry(badger.NewEntry([]byte(kk), []byte{}))
    40  	})
    41  	if err != nil {
    42  		// TODO: handle
    43  		panic(err)
    44  	}
    45  	err = ll.db.Update(func(txn *badger.Txn) error {
    46  		return txn.SetEntry(badger.NewEntry([]byte(kv), []byte{}))
    47  	})
    48  	if err != nil {
    49  		// TODO: handle
    50  		panic(err)
    51  	}
    52  	// err = ll.db.Update(func(txn *badger.Txn) error {
    53  	// 	return txn.SetEntry(badger.NewEntry([]byte(ks), []byte{}))
    54  	// })
    55  	// if err != nil {
    56  	// 	// TODO: handle
    57  	// 	panic(err)
    58  	// }
    59  }
    60  
    61  //revive:disable-next-line:get-return A callback is fine
    62  func (ll *Labels) GetKeys(cb func(k string) bool) {
    63  	err := ll.db.View(func(txn *badger.Txn) error {
    64  		opts := badger.DefaultIteratorOptions
    65  		opts.Prefix = []byte("l:")
    66  		opts.PrefetchValues = false
    67  		it := txn.NewIterator(opts)
    68  		defer it.Close()
    69  		for it.Rewind(); it.Valid(); it.Next() {
    70  			item := it.Item()
    71  			k := item.Key()
    72  			shouldContinue := cb(string(k[2:]))
    73  			if !shouldContinue {
    74  				return nil
    75  			}
    76  		}
    77  		return nil
    78  	})
    79  	if err != nil {
    80  		// TODO: handle
    81  		panic(err)
    82  	}
    83  }
    84  
    85  // Delete removes key value label pair from the storage.
    86  // If the pair can not be found, no error is returned.
    87  func (ll *Labels) Delete(key, value string) error {
    88  	return ll.db.Update(func(txn *badger.Txn) error {
    89  		return txn.Delete([]byte("v:" + key + ":" + value))
    90  	})
    91  }
    92  
    93  //revive:disable-next-line:get-return A callback is fine
    94  func (ll *Labels) GetValues(key string, cb func(v string) bool) {
    95  	err := ll.db.View(func(txn *badger.Txn) error {
    96  		opts := badger.DefaultIteratorOptions
    97  		opts.Prefix = []byte("v:" + key + ":")
    98  		opts.PrefetchValues = false
    99  		it := txn.NewIterator(opts)
   100  		defer it.Close()
   101  		for it.Rewind(); it.Valid(); it.Next() {
   102  			item := it.Item()
   103  			k := item.Key()
   104  			ks := string(k)
   105  			li := strings.LastIndex(ks, ":") + 1
   106  			shouldContinue := cb(ks[li:])
   107  			if !shouldContinue {
   108  				return nil
   109  			}
   110  		}
   111  		return nil
   112  	})
   113  	if err != nil {
   114  		// TODO: handle
   115  		panic(err)
   116  	}
   117  }