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