github.com/fiatjaf/generic-ristretto@v0.0.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  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[V any] struct {
    44  	sync.RWMutex
    45  	buckets map[int64]bucket
    46  }
    47  
    48  func newExpirationMap[V any]() *expirationMap[V] {
    49  	return &expirationMap[V]{
    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  
    81  	m.Lock()
    82  	defer m.Unlock()
    83  
    84  	oldBucketNum := storageBucket(oldExpTime)
    85  	oldBucket, ok := m.buckets[oldBucketNum]
    86  	if ok {
    87  		delete(oldBucket, key)
    88  	}
    89  
    90  	newBucketNum := storageBucket(newExpTime)
    91  	newBucket, ok := m.buckets[newBucketNum]
    92  	if !ok {
    93  		newBucket = make(bucket)
    94  		m.buckets[newBucketNum] = newBucket
    95  	}
    96  	newBucket[key] = conflict
    97  }
    98  
    99  func (m *expirationMap[_]) del(key uint64, expiration time.Time) {
   100  	if m == nil {
   101  		return
   102  	}
   103  
   104  	bucketNum := storageBucket(expiration)
   105  	m.Lock()
   106  	defer m.Unlock()
   107  	_, ok := m.buckets[bucketNum]
   108  	if !ok {
   109  		return
   110  	}
   111  	delete(m.buckets[bucketNum], key)
   112  }
   113  
   114  // cleanup removes all the items in the bucket that was just completed. It deletes
   115  // those items from the store, and calls the onEvict function on those items.
   116  // This function is meant to be called periodically.
   117  func (m *expirationMap[V]) cleanup(store store[V], policy policy[V], onEvict func(item *Item[V])) {
   118  	if m == nil {
   119  		return
   120  	}
   121  
   122  	m.Lock()
   123  	now := time.Now()
   124  	bucketNum := cleanupBucket(now)
   125  	keys := m.buckets[bucketNum]
   126  	delete(m.buckets, bucketNum)
   127  	m.Unlock()
   128  
   129  	for key, conflict := range keys {
   130  		// Sanity check. Verify that the store agrees that this key is expired.
   131  		if store.Expiration(key).After(now) {
   132  			continue
   133  		}
   134  
   135  		cost := policy.Cost(key)
   136  		policy.Del(key)
   137  		_, value := store.Del(key, conflict)
   138  
   139  		if onEvict != nil {
   140  			onEvict(&Item[V]{Key: key,
   141  				Conflict: conflict,
   142  				Value:    value,
   143  				Cost:     cost,
   144  			})
   145  		}
   146  	}
   147  }