Ristretto 高性能 go 语言缓存

// Push appends an item in the ring buffer and drains (copies items and
// sends to Consumer) if full.
func (s *ringStripe) Push(item uint64) {
  s.data = append(s.data, item)
  // if we should drain
  if len(s.data) >= s.capa {
    // Send elements to consumer. Create a new one.
    if s.cons.Push(s.data) {
      s.data = make([]uint64, 0, s.capa)
    } else {
      s.data = s.data[:0]
    }
  }
}

// ringBuffer stores multiple buffers (stripes) and distributes Pushed items
// between them to lower contention.
//
// This implements the "batching" process described in the BP-Wrapper paper
// (section III part A).
type ringBuffer struct {
  pool *sync.Pool
}
// newRingBuffer returns a striped ring buffer. The Consumer in ringConfig will
// be called when individual stripes are full and need to drain their elements.
func newRingBuffer(cons ringConsumer, capa int64) *ringBuffer {
  // LOSSY buffers use a very simple sync.Pool for concurrently reusing
  // stripes. We do lose some stripes due to GC (unheld items in sync.Pool
  // are cleared), but the performance gains generally outweigh the small
  // percentage of elements lost. The performance primarily comes from
  // low-level runtime functions used in the standard library that aren't
  // available to us (such as runtime_procPin()).
  return &ringBuffer{
    pool: &sync.Pool{
      New: func() interface{} { return newRingStripe(cons, capa) },
    },
  }
}
// Push adds an element to one of the internal stripes and possibly drains if
// the stripe becomes full.
func (b *ringBuffer) Push(item uint64) {
  // reuse or create a new stripe
  stripe := b.pool.Get().(*ringStripe)
  stripe.Push(item)
  b.pool.Put(stripe)
}

// TODO: Figure out a way to re-use memhash for the second uint64 hash, we
//       already know that appending bytes isn't reliable for generating a
//       second hash (see Ristretto PR #88).
//
//       We also know that while the Go runtime has a runtime memhash128
//       function, it's not possible to use it to generate [2]uint64 or
//       anything resembling a 128bit hash, even though that's exactly what
//       we need in this situation.
func KeyToHash(key interface{}) (uint64, uint64) {
  if key == nil {
    return 0, 0
  }
  switch k := key.(type) {
  case uint64:
    return k, 0
  case string:
    raw := []byte(k)
    return MemHash(raw), xxhash.Sum64(raw)
  case []byte:
    return MemHash(k), xxhash.Sum64(k)
  case byte:
    return uint64(k), 0
  case int:
    return uint64(k), 0
  case int32:
    return uint64(k), 0
  case uint32:
    return uint64(k), 0
  case int64:
    return uint64(k), 0
  default:
    panic("Key type not supported")
  }
}

type lockedMap struct {
  sync.RWMutex
  data map[uint64]storeItem
}

const numShards uint64 = 256

type shardedMap struct {
  shards []*lockedMap
}

func newShardedMap() *shardedMap {
  sm := &shardedMap{
    shards: make([]*lockedMap, int(numShards)),
  }
  for i := range sm.shards {
    sm.shards[i] = newLockedMap()
  }
  return sm
}

func (sm *shardedMap) Get(key, conflict uint64) (interface{}, bool) {
  return sm.shards[key%numShards].Get(key, conflict)
}

func (sm *shardedMap) Set(key, conflict uint64, value interface{}) {
  sm.shards[key%numShards].Set(key, conflict, value)
}

func (c *Cache) Set(key, value interface{}, cost int64) bool {
  if c == nil || key == nil {
    return false
  }
  keyHash, conflictHash := c.keyToHash(key)
  i := &item{
    flag:     itemNew,
    key:      keyHash,
    conflict: conflictHash,
    value:    value,
    cost:     cost,
  }
  // attempt to immediately update hashmap value and set flag to update so the
  // cost is eventually updated
  if c.store.Update(keyHash, conflictHash, i.value) {
    i.flag = itemUpdate
  }
  // attempt to send item to policy
  select {
  case c.setBuf <- i:
    return true
  default:
    c.Metrics.add(dropSets, keyHash, 1)
    return false
  }
}

// Snippet from the Admission and Eviction Algorithm
incHits := p.admit.Estimate(key)
for ; room < 0; room = p.evict.roomLeft(cost) {
    sample = p.evict.fillSample(sample)
    minKey, minHits, minId := uint64(0), int64(math.MaxInt64), 0
    for i, pair := range sample {
        if hits := p.admit.Estimate(pair.key); hits < minHits {
            minKey, minHits, minId = pair.key, hits, i
        }
    }
    if incHits < minHits {
        p.stats.Add(rejectSets, key, 1)
        return victims, false
    }
    p.evict.del(minKey)
    sample[minId] = sample[len(sample)-1]
    sample = sample[:len(sample)-1]
    victims = append(victims, minKey)
}

type sampledLFU struct {
  keyCosts map[uint64]int64
  maxCost  int64
  used     int64
  metrics  *Metrics
}

func (p *sampledLFU) fillSample(in []*policyPair) []*policyPair {
  if len(in) >= lfuSample {
    return in
  }
  for key, cost := range p.keyCosts {
    in = append(in, &policyPair{key, cost})
    if len(in) >= lfuSample {
      return in
    }
  }
  return in
}