github.com/castai/kvisor@v1.7.1-0.20240516114728-b3572a2607b5/pkg/bucketcache/bucketcache.go (about)

     1  package bucketcache
     2  
     3  import (
     4  	"github.com/elastic/go-freelru"
     5  )
     6  
     7  type BucketCache[K comparable, V any] struct {
     8  	cache         freelru.Cache[K, []V]
     9  	maxBucketSize int
    10  }
    11  
    12  func New[K comparable, V any](cacheSize uint32, maxBucketSize uint32, hash freelru.HashKeyCallback[K]) (*BucketCache[K, V], error) {
    13  	cache, err := freelru.NewSynced[K, []V](cacheSize, hash)
    14  	if err != nil {
    15  		return nil, err
    16  	}
    17  
    18  	return &BucketCache[K, V]{
    19  		cache:         cache,
    20  		maxBucketSize: int(maxBucketSize),
    21  	}, nil
    22  }
    23  
    24  func (b *BucketCache[K, V]) AddToBucket(k K, val V) bool {
    25  	return b.addToCache(k, val, false)
    26  }
    27  
    28  func (b *BucketCache[K, V]) ForceAddToBucket(k K, val V) {
    29  	b.addToCache(k, val, true)
    30  }
    31  
    32  func (b *BucketCache[K, V]) addToCache(k K, val V, force bool) bool {
    33  	bucket, found := b.cache.Get(k)
    34  	if !found {
    35  		b.cache.Add(k, []V{val})
    36  		return true
    37  	}
    38  
    39  	if len(bucket) >= b.maxBucketSize {
    40  		if force {
    41  			bucket[0] = val
    42  			b.cache.Add(k, bucket)
    43  			return true
    44  		}
    45  
    46  		return false
    47  	}
    48  
    49  	b.cache.Add(k, append(bucket, val))
    50  
    51  	return true
    52  }
    53  
    54  func (b *BucketCache[K, V]) GetBucket(k K) []V {
    55  	res, _ := b.cache.Get(k)
    56  	return res
    57  }