github.com/etecs-ru/ristretto@v0.9.1/ttl.go (about)

     1  /*
     2   * Copyright 2020 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 ristretto
    18  
    19  import (
    20  	"sync"
    21  	"time"
    22  )
    23  
    24  // TODO: find the optimal value or make it configurable.
    25  var bucketDurationSecs = int64(5) //nolint:gomnd,unused,varcheck,gochecknoglobals,lll,deadcode,revive // adopt fork, do not touch it
    26  
    27  func storageBucket(t time.Time) int64 {
    28  	return (t.Unix() / bucketDurationSecs) + 1
    29  }
    30  
    31  func cleanupBucket(t time.Time) int64 {
    32  	// The bucket to cleanup is always behind the storage bucket by one so that
    33  	// no elements in that bucket (which might not have expired yet) are deleted.
    34  	return storageBucket(t) - 1
    35  }
    36  
    37  // bucket type is a map of key to conflict.
    38  type bucket map[uint64]uint64
    39  
    40  // expirationMap is a map of bucket number to the corresponding bucket.
    41  type expirationMap struct {
    42  	buckets map[int64]bucket
    43  	sync.RWMutex
    44  }
    45  
    46  func newExpirationMap() *expirationMap {
    47  	return &expirationMap{
    48  		buckets: make(map[int64]bucket),
    49  	}
    50  }
    51  
    52  func (m *expirationMap) add(key, conflict uint64, expiration time.Time) {
    53  	if m == nil {
    54  		return
    55  	}
    56  
    57  	// Items that don't expire don't need to be in the expiration map.
    58  	if expiration.IsZero() {
    59  		return
    60  	}
    61  
    62  	bucketNum := storageBucket(expiration)
    63  	m.Lock()
    64  	defer m.Unlock()
    65  
    66  	b, ok := m.buckets[bucketNum]
    67  	if !ok {
    68  		b = make(bucket)
    69  		m.buckets[bucketNum] = b
    70  	}
    71  	b[key] = conflict
    72  }
    73  
    74  func (m *expirationMap) update(key, conflict uint64, oldExpTime, newExpTime time.Time) {
    75  	if m == nil {
    76  		return
    77  	}
    78  	if oldExpTime.IsZero() && newExpTime.IsZero() {
    79  		return
    80  	}
    81  
    82  	m.Lock()
    83  	defer m.Unlock()
    84  
    85  	oldBucketNum := storageBucket(oldExpTime)
    86  	newBucketNum := storageBucket(newExpTime)
    87  	if oldBucketNum == newBucketNum {
    88  		// No change.
    89  		return
    90  	}
    91  
    92  	oldBucket, ok := m.buckets[oldBucketNum]
    93  	if ok {
    94  		delete(oldBucket, key)
    95  	}
    96  
    97  	newBucket, ok := m.buckets[newBucketNum]
    98  	if !ok {
    99  		newBucket = make(bucket)
   100  		m.buckets[newBucketNum] = newBucket
   101  	}
   102  	newBucket[key] = conflict
   103  }
   104  
   105  func (m *expirationMap) del(key uint64, expiration time.Time) {
   106  	if m == nil {
   107  		return
   108  	}
   109  
   110  	bucketNum := storageBucket(expiration)
   111  	m.Lock()
   112  	defer m.Unlock()
   113  	_, ok := m.buckets[bucketNum]
   114  	if !ok {
   115  		return
   116  	}
   117  	delete(m.buckets[bucketNum], key)
   118  }
   119  
   120  // cleanup removes all the items in the bucket that was just completed. It deletes
   121  // those items from the store, and calls the onEvict function on those items.
   122  // This function is meant to be called periodically.
   123  func (m *expirationMap) cleanup(store *shardedMap, policy *lfuPolicy, onEvict itemCallback) {
   124  	if m == nil {
   125  		return
   126  	}
   127  
   128  	m.Lock()
   129  	now := time.Now()
   130  	bucketNum := cleanupBucket(now)
   131  	keys := m.buckets[bucketNum]
   132  	delete(m.buckets, bucketNum)
   133  	m.Unlock()
   134  
   135  	for key, conflict := range keys {
   136  		// Sanity check. Verify that the store agrees that this key is expired.
   137  		if store.Expiration(key).After(now) {
   138  			continue
   139  		}
   140  
   141  		cost := policy.Cost(key)
   142  		policy.Del(key)
   143  		_, value := store.Del(key, conflict)
   144  
   145  		if onEvict != nil {
   146  			onEvict(&Item{
   147  				Key:      key,
   148  				Conflict: conflict,
   149  				Value:    value,
   150  				Cost:     cost,
   151  			})
   152  		}
   153  	}
   154  }