github.com/fiatjaf/generic-ristretto@v0.0.1/stress_test.go (about)

     1  package ristretto
     2  
     3  import (
     4  	"container/heap"
     5  	"fmt"
     6  	"math/rand"
     7  	"runtime"
     8  	"sync"
     9  	"testing"
    10  	"time"
    11  
    12  	"github.com/fiatjaf/generic-ristretto/sim"
    13  	"github.com/stretchr/testify/require"
    14  )
    15  
    16  func TestStressSetGet(t *testing.T) {
    17  	c, err := NewCache(&Config[int, int]{
    18  		NumCounters:        1000,
    19  		MaxCost:            100,
    20  		IgnoreInternalCost: true,
    21  		BufferItems:        64,
    22  		Metrics:            true,
    23  	})
    24  	require.NoError(t, err)
    25  
    26  	for i := 0; i < 100; i++ {
    27  		c.Set(i, i, 1)
    28  	}
    29  	time.Sleep(wait)
    30  	wg := &sync.WaitGroup{}
    31  	for i := 0; i < runtime.GOMAXPROCS(0); i++ {
    32  		wg.Add(1)
    33  		go func() {
    34  			r := rand.New(rand.NewSource(time.Now().UnixNano()))
    35  			for a := 0; a < 1000; a++ {
    36  				k := r.Int() % 10
    37  				if val, ok := c.Get(k); &val == nil || !ok {
    38  					err = fmt.Errorf("expected %d but got nil", k)
    39  					break
    40  				} else if val != 0 && val != k {
    41  					err = fmt.Errorf("expected %d but got %d", k, val)
    42  					break
    43  				}
    44  			}
    45  			wg.Done()
    46  		}()
    47  	}
    48  	wg.Wait()
    49  	require.NoError(t, err)
    50  	require.Equal(t, 1.0, c.Metrics.Ratio())
    51  }
    52  
    53  func TestStressHitRatio(t *testing.T) {
    54  	key := sim.NewZipfian(1.0001, 1, 1000)
    55  	c, err := NewCache(&Config[uint64, uint64]{
    56  		NumCounters: 1000,
    57  		MaxCost:     100,
    58  		BufferItems: 64,
    59  		Metrics:     true,
    60  	})
    61  	require.NoError(t, err)
    62  
    63  	o := NewClairvoyant(100)
    64  	for i := 0; i < 10000; i++ {
    65  		k, err := key()
    66  		require.NoError(t, err)
    67  
    68  		if _, ok := o.Get(k); !ok {
    69  			o.Set(k, k, 1)
    70  		}
    71  		if _, ok := c.Get(k); !ok {
    72  			c.Set(k, k, 1)
    73  		}
    74  	}
    75  	t.Logf("actual: %.2f, optimal: %.2f", c.Metrics.Ratio(), o.Metrics().Ratio())
    76  }
    77  
    78  // Clairvoyant is a mock cache providing us with optimal hit ratios to compare
    79  // with Ristretto's. It looks ahead and evicts the absolute least valuable item,
    80  // which we try to approximate in a real cache.
    81  type Clairvoyant struct {
    82  	capacity uint64
    83  	hits     map[uint64]uint64
    84  	access   []uint64
    85  }
    86  
    87  func NewClairvoyant(capacity uint64) *Clairvoyant {
    88  	return &Clairvoyant{
    89  		capacity: capacity,
    90  		hits:     make(map[uint64]uint64),
    91  		access:   make([]uint64, 0),
    92  	}
    93  }
    94  
    95  // Get just records the cache access so that we can later take this event into
    96  // consideration when calculating the absolute least valuable item to evict.
    97  func (c *Clairvoyant) Get(key interface{}) (interface{}, bool) {
    98  	c.hits[key.(uint64)]++
    99  	c.access = append(c.access, key.(uint64))
   100  	return nil, false
   101  }
   102  
   103  // Set isn't important because it is only called after a Get (in the case of our
   104  // hit ratio benchmarks, at least).
   105  func (c *Clairvoyant) Set(key, value interface{}, cost int64) bool {
   106  	return false
   107  }
   108  
   109  func (c *Clairvoyant) Metrics() *Metrics {
   110  	stat := newMetrics()
   111  	look := make(map[uint64]struct{}, c.capacity)
   112  	data := &clairvoyantHeap{}
   113  	heap.Init(data)
   114  	for _, key := range c.access {
   115  		if _, has := look[key]; has {
   116  			stat.add(hit, 0, 1)
   117  			continue
   118  		}
   119  		if uint64(data.Len()) >= c.capacity {
   120  			victim := heap.Pop(data)
   121  			delete(look, victim.(*clairvoyantItem).key)
   122  		}
   123  		stat.add(miss, 0, 1)
   124  		look[key] = struct{}{}
   125  		heap.Push(data, &clairvoyantItem{key, c.hits[key]})
   126  	}
   127  	return stat
   128  }
   129  
   130  type clairvoyantItem struct {
   131  	key  uint64
   132  	hits uint64
   133  }
   134  
   135  type clairvoyantHeap []*clairvoyantItem
   136  
   137  func (h clairvoyantHeap) Len() int           { return len(h) }
   138  func (h clairvoyantHeap) Less(i, j int) bool { return h[i].hits < h[j].hits }
   139  func (h clairvoyantHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
   140  
   141  func (h *clairvoyantHeap) Push(x interface{}) {
   142  	*h = append(*h, x.(*clairvoyantItem))
   143  }
   144  
   145  func (h *clairvoyantHeap) Pop() interface{} {
   146  	old := *h
   147  	n := len(old)
   148  	x := old[n-1]
   149  	*h = old[0 : n-1]
   150  	return x
   151  }