github.com/coocood/badger@v1.5.1-0.20200528065104-c02ac3616d04/cache/policy.go (about)

     1  /*
     2   * Copyright 2019 Dgraph Labs, Inc. and Contributors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package cache
    18  
    19  import (
    20  	"math"
    21  	"sync"
    22  
    23  	"github.com/coocood/badger/cache/z"
    24  )
    25  
    26  const (
    27  	// lfuSample is the number of items to sample when looking at eviction
    28  	// candidates. 5 seems to be the most optimal number [citation needed].
    29  	lfuSample = 5
    30  )
    31  
    32  type policy struct {
    33  	sync.Mutex
    34  	admit   *tinyLFU
    35  	evict   *sampledLFU
    36  	itemsCh chan []uint64
    37  	stop    chan struct{}
    38  	metrics *Metrics
    39  }
    40  
    41  func newPolicy(numCounters, maxCost int64) *policy {
    42  	p := &policy{
    43  		admit:   newTinyLFU(numCounters),
    44  		evict:   newSampledLFU(maxCost),
    45  		itemsCh: make(chan []uint64, 3),
    46  		stop:    make(chan struct{}),
    47  	}
    48  	go p.processItems()
    49  	return p
    50  }
    51  
    52  func (p *policy) CollectMetrics(metrics *Metrics) {
    53  	p.metrics = metrics
    54  	p.evict.metrics = metrics
    55  }
    56  
    57  type policyPair struct {
    58  	key  uint64
    59  	cost int64
    60  }
    61  
    62  func (p *policy) processItems() {
    63  	for {
    64  		select {
    65  		case items := <-p.itemsCh:
    66  			p.Lock()
    67  			p.admit.Push(items)
    68  			p.Unlock()
    69  		case <-p.stop:
    70  			return
    71  		}
    72  	}
    73  }
    74  
    75  func (p *policy) Push(keys []uint64) bool {
    76  	if len(keys) == 0 {
    77  		return true
    78  	}
    79  	select {
    80  	case p.itemsCh <- keys:
    81  		p.metrics.add(keepGets, keys[0], uint64(len(keys)))
    82  		return true
    83  	default:
    84  		p.metrics.add(dropGets, keys[0], uint64(len(keys)))
    85  		return false
    86  	}
    87  }
    88  
    89  func (p *policy) Add(key uint64, cost int64) ([]*item, bool) {
    90  	p.Lock()
    91  	defer p.Unlock()
    92  	// can't add an item bigger than entire cache
    93  	if cost > p.evict.maxCost {
    94  		return nil, false
    95  	}
    96  	// we don't need to go any further if the item is already in the cache
    97  	if has := p.evict.updateIfHas(key, cost); has {
    98  		return nil, true
    99  	}
   100  	// if we got this far, this key doesn't exist in the cache
   101  	//
   102  	// calculate the remaining room in the cache (usually bytes)
   103  	room := p.evict.roomLeft(cost)
   104  	if room >= 0 {
   105  		// there's enough room in the cache to store the new item without
   106  		// overflowing, so we can do that now and stop here
   107  		p.evict.add(key, cost)
   108  		p.metrics.add(costAdd, key, uint64(cost))
   109  		p.metrics.add(keyAdd, key, 1)
   110  		return nil, true
   111  	}
   112  	// incHits is the hit count for the incoming item
   113  	incHits := p.admit.Estimate(key)
   114  	// sample is the eviction candidate pool to be filled via random sampling
   115  	//
   116  	// TODO: perhaps we should use a min heap here. Right now our time
   117  	// complexity is N for finding the min. Min heap should bring it down to
   118  	// O(lg N).
   119  	sample := make([]*policyPair, 0, lfuSample)
   120  	// as items are evicted they will be appended to victims
   121  	victims := make([]*item, 0)
   122  	// delete victims until there's enough space or a minKey is found that has
   123  	// more hits than incoming item.
   124  	for ; room < 0; room = p.evict.roomLeft(cost) {
   125  		// fill up empty slots in sample
   126  		sample = p.evict.fillSample(sample)
   127  		// find minimally used item in sample
   128  		minKey, minHits, minId := uint64(0), int64(math.MaxInt64), 0
   129  		for i, pair := range sample {
   130  			// look up hit count for sample key
   131  			if hits := p.admit.Estimate(pair.key); hits < minHits {
   132  				minKey, minHits, minId = pair.key, hits, i
   133  			}
   134  		}
   135  		// if the incoming item isn't worth keeping in the policy, reject.
   136  		if incHits < minHits {
   137  			p.metrics.add(rejectSets, key, 1)
   138  			return victims, false
   139  		}
   140  		// delete the victim from metadata
   141  		p.evict.del(minKey)
   142  
   143  		// delete the victim from sample
   144  		sample[minId] = sample[len(sample)-1]
   145  		sample = sample[:len(sample)-1]
   146  		// store victim in evicted victims slice
   147  		victims = append(victims, &item{
   148  			key: minKey,
   149  		})
   150  	}
   151  	p.evict.add(key, cost)
   152  	p.metrics.add(costAdd, key, uint64(cost))
   153  	p.metrics.add(keyAdd, key, 1)
   154  	return victims, true
   155  }
   156  
   157  func (p *policy) Has(key uint64) bool {
   158  	p.Lock()
   159  	_, exists := p.evict.keyCosts[key]
   160  	p.Unlock()
   161  	return exists
   162  }
   163  
   164  func (p *policy) Del(key uint64) {
   165  	p.Lock()
   166  	p.evict.del(key)
   167  	p.Unlock()
   168  }
   169  
   170  func (p *policy) Cap() int64 {
   171  	p.Lock()
   172  	capacity := int64(p.evict.maxCost - p.evict.used)
   173  	p.Unlock()
   174  	return capacity
   175  }
   176  
   177  func (p *policy) Update(key uint64, cost int64) {
   178  	p.Lock()
   179  	p.evict.updateIfHas(key, cost)
   180  	p.Unlock()
   181  }
   182  
   183  func (p *policy) Cost(key uint64) int64 {
   184  	p.Lock()
   185  	if cost, found := p.evict.keyCosts[key]; found {
   186  		p.Unlock()
   187  		return cost
   188  	}
   189  	p.Unlock()
   190  	return -1
   191  }
   192  
   193  func (p *policy) Clear() {
   194  	p.Lock()
   195  	p.admit.clear()
   196  	p.evict.clear()
   197  	p.Unlock()
   198  }
   199  
   200  func (p *policy) Close() {
   201  	// block until p.processItems goroutine is returned
   202  	p.stop <- struct{}{}
   203  	close(p.stop)
   204  	close(p.itemsCh)
   205  }
   206  
   207  // sampledLFU is an eviction helper storing key-cost pairs.
   208  type sampledLFU struct {
   209  	keyCosts map[uint64]int64
   210  	maxCost  int64
   211  	used     int64
   212  	metrics  *Metrics
   213  }
   214  
   215  func newSampledLFU(maxCost int64) *sampledLFU {
   216  	return &sampledLFU{
   217  		keyCosts: make(map[uint64]int64),
   218  		maxCost:  maxCost,
   219  	}
   220  }
   221  
   222  func (p *sampledLFU) roomLeft(cost int64) int64 {
   223  	return p.maxCost - (p.used + cost)
   224  }
   225  
   226  func (p *sampledLFU) fillSample(in []*policyPair) []*policyPair {
   227  	if len(in) >= lfuSample {
   228  		return in
   229  	}
   230  	for key, cost := range p.keyCosts {
   231  		in = append(in, &policyPair{key, cost})
   232  		if len(in) >= lfuSample {
   233  			return in
   234  		}
   235  	}
   236  	return in
   237  }
   238  
   239  func (p *sampledLFU) del(key uint64) {
   240  	cost, ok := p.keyCosts[key]
   241  	if !ok {
   242  		return
   243  	}
   244  	p.used -= cost
   245  	delete(p.keyCosts, key)
   246  	p.metrics.add(costEvict, key, uint64(cost))
   247  	p.metrics.add(keyEvict, key, 1)
   248  }
   249  
   250  func (p *sampledLFU) add(key uint64, cost int64) {
   251  	p.keyCosts[key] = cost
   252  	p.used += cost
   253  }
   254  
   255  func (p *sampledLFU) updateIfHas(key uint64, cost int64) bool {
   256  	if prev, found := p.keyCosts[key]; found {
   257  		// update the cost of an existing key, but don't worry about evicting,
   258  		// evictions will be handled the next time a new item is added
   259  		p.metrics.add(keyUpdate, key, 1)
   260  		if prev > cost {
   261  			diff := prev - cost
   262  			p.metrics.add(costAdd, key, ^uint64(uint64(diff)-1))
   263  		} else if cost > prev {
   264  			diff := cost - prev
   265  			p.metrics.add(costAdd, key, uint64(diff))
   266  		}
   267  		p.used += cost - prev
   268  		p.keyCosts[key] = cost
   269  		return true
   270  	}
   271  	return false
   272  }
   273  
   274  func (p *sampledLFU) clear() {
   275  	p.used = 0
   276  	p.keyCosts = make(map[uint64]int64)
   277  }
   278  
   279  // tinyLFU is an admission helper that keeps track of access frequency using
   280  // tiny (4-bit) counters in the form of a count-min sketch.
   281  // tinyLFU is NOT thread safe.
   282  type tinyLFU struct {
   283  	freq    *cmSketch
   284  	door    *z.Bloom
   285  	incrs   int64
   286  	resetAt int64
   287  }
   288  
   289  func newTinyLFU(numCounters int64) *tinyLFU {
   290  	return &tinyLFU{
   291  		freq:    newCmSketch(numCounters),
   292  		door:    z.NewBloomFilter(float64(numCounters), 0.01),
   293  		resetAt: numCounters,
   294  	}
   295  }
   296  
   297  func (p *tinyLFU) Push(keys []uint64) {
   298  	for _, key := range keys {
   299  		p.Increment(key)
   300  	}
   301  }
   302  
   303  func (p *tinyLFU) Estimate(key uint64) int64 {
   304  	hits := p.freq.Estimate(key)
   305  	if p.door.Has(key) {
   306  		hits += 1
   307  	}
   308  	return hits
   309  }
   310  
   311  func (p *tinyLFU) Increment(key uint64) {
   312  	// flip doorkeeper bit if not already
   313  	if added := p.door.AddIfNotHas(key); !added {
   314  		// increment count-min counter if doorkeeper bit is already set.
   315  		p.freq.Increment(key)
   316  	}
   317  	p.incrs++
   318  	if p.incrs >= p.resetAt {
   319  		p.reset()
   320  	}
   321  }
   322  
   323  func (p *tinyLFU) reset() {
   324  	// Zero out incrs.
   325  	p.incrs = 0
   326  	// clears doorkeeper bits
   327  	p.door.Clear()
   328  	// halves count-min counters
   329  	p.freq.Reset()
   330  }
   331  
   332  func (p *tinyLFU) clear() {
   333  	p.incrs = 0
   334  	p.door.Clear()
   335  	p.freq.Clear()
   336  }