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

     1  //go:build !windows
     2  // +build !windows
     3  
     4  package cache
     5  
     6  import (
     7  	"fmt"
     8  	"io"
     9  	"os"
    10  	"path/filepath"
    11  	"strconv"
    12  
    13  	"github.com/dgraph-io/badger/v2"
    14  	"github.com/dgraph-io/badger/v2/options"
    15  	. "github.com/onsi/ginkgo/v2"
    16  	. "github.com/onsi/gomega"
    17  	"github.com/prometheus/client_golang/prometheus"
    18  	"github.com/prometheus/client_golang/prometheus/promauto"
    19  	"github.com/pyroscope-io/pyroscope/pkg/config"
    20  	"github.com/pyroscope-io/pyroscope/pkg/testing"
    21  )
    22  
    23  type fakeCodec struct{}
    24  
    25  const fakeCodecEmptyStub = "empty"
    26  
    27  func (fakeCodec) New(_ string) interface{} { return fakeCodecEmptyStub }
    28  
    29  func (fakeCodec) Serialize(w io.Writer, _ string, v interface{}) error {
    30  	_, err := w.Write([]byte(v.(string)))
    31  	return err
    32  }
    33  
    34  func (fakeCodec) Deserialize(r io.Reader, _ string) (interface{}, error) {
    35  	b, err := io.ReadAll(r)
    36  	return string(b), err
    37  }
    38  
    39  var _ = Describe("cache", func() {
    40  	var c *Cache
    41  	testing.WithConfig(func(cfg **config.Config) {
    42  		JustBeforeEach(func() {
    43  			badgerPath := filepath.Join((*cfg).Server.StoragePath)
    44  			err := os.MkdirAll(badgerPath, 0o755)
    45  			Expect(err).ToNot(HaveOccurred())
    46  
    47  			badgerOptions := badger.DefaultOptions(badgerPath)
    48  			badgerOptions = badgerOptions.WithTruncate(false)
    49  			badgerOptions = badgerOptions.WithSyncWrites(false)
    50  			badgerOptions = badgerOptions.WithCompression(options.ZSTD)
    51  
    52  			db, err := badger.Open(badgerOptions)
    53  			Expect(err).ToNot(HaveOccurred())
    54  
    55  			reg := prometheus.NewRegistry()
    56  			c = New(Config{
    57  				DB:     db,
    58  				Codec:  fakeCodec{},
    59  				Prefix: "p:",
    60  				Metrics: &Metrics{
    61  					MissesCounter: promauto.With(reg).NewCounter(prometheus.CounterOpts{
    62  						Name: "cache_test_miss",
    63  					}),
    64  					ReadsCounter: promauto.With(reg).NewCounter(prometheus.CounterOpts{
    65  						Name: "storage_test_read",
    66  					}),
    67  					DBWrites: promauto.With(reg).NewHistogram(prometheus.HistogramOpts{
    68  						Name: "storage_test_write",
    69  					}),
    70  					DBReads: promauto.With(reg).NewHistogram(prometheus.HistogramOpts{
    71  						Name: "storage_test_reads",
    72  					}),
    73  				},
    74  			})
    75  		})
    76  	})
    77  
    78  	It("works properly", func() {
    79  		done := make(chan interface{})
    80  		go func() {
    81  			for i := 0; i < 200; i++ {
    82  				c.Put(fmt.Sprintf("foo-%d", i), fmt.Sprintf("bar-%d", i))
    83  			}
    84  
    85  			v, err := c.GetOrCreate("foo-199")
    86  			Expect(err).ToNot(HaveOccurred())
    87  			Expect(v).To(Equal("bar-199"))
    88  
    89  			v, err = c.GetOrCreate("foo-1")
    90  			Expect(err).ToNot(HaveOccurred())
    91  			Expect(v).To(Equal("bar-1"))
    92  
    93  			v, err = c.GetOrCreate("foo-1234")
    94  			Expect(err).ToNot(HaveOccurred())
    95  			Expect(v).To(Equal(fakeCodecEmptyStub))
    96  			c.Flush()
    97  
    98  			close(done)
    99  		}()
   100  		Eventually(done, 3).Should(BeClosed())
   101  	})
   102  
   103  	Context("discard prefix", func() {
   104  		It("removes data from cache and disk", func() {
   105  			const (
   106  				prefixToDelete = "0:"
   107  				prefixToKeep   = "1:"
   108  				n              = 5 * defaultBatchSize
   109  			)
   110  
   111  			for i := 0; i < n; i++ {
   112  				v := strconv.Itoa(i)
   113  				c.Put(prefixToDelete+v, v)
   114  				c.Put(prefixToKeep+v, v)
   115  			}
   116  
   117  			k := prefixToDelete + strconv.Itoa(0)
   118  			_, ok := c.Lookup(k)
   119  			Expect(ok).To(BeTrue())
   120  			c.Flush()
   121  
   122  			v, err := c.GetOrCreate(k)
   123  			Expect(err).ToNot(HaveOccurred())
   124  			Expect(v).To(Equal("0"))
   125  			_, ok = c.Lookup(k)
   126  			Expect(ok).To(BeTrue())
   127  
   128  			Expect(c.DiscardPrefix(prefixToDelete)).ToNot(HaveOccurred())
   129  
   130  			v, err = c.GetOrCreate(k)
   131  			Expect(err).ToNot(HaveOccurred())
   132  			Expect(v).To(Equal(fakeCodecEmptyStub))
   133  
   134  			k = prefixToKeep + strconv.Itoa(0)
   135  			v, err = c.GetOrCreate(k)
   136  			Expect(err).ToNot(HaveOccurred())
   137  			Expect(v).To(Equal("0"))
   138  		})
   139  	})
   140  })