github.com/confluentinc/confluent-kafka-go@v1.9.2/schemaregistry/cache/lrucache.go (about)

     1  /**
     2   * Copyright 2022 Confluent Inc.
     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  	"container/list"
    21  	"fmt"
    22  	"sync"
    23  )
    24  
    25  const maxPreallocateCapacity = 10000
    26  
    27  // LRUCache is a Least Recently Used (LRU) Cache with given capacity
    28  type LRUCache struct {
    29  	cacheLock   sync.RWMutex
    30  	capacity    int
    31  	entries     map[interface{}]interface{}
    32  	lruElements map[interface{}]*list.Element
    33  	lruKeys     *list.List
    34  }
    35  
    36  // NewLRUCache creates a new Least Recently Used (LRU) Cache
    37  //
    38  // Parameters:
    39  //  * `capacity` - a positive integer indicating the max capacity of this cache
    40  //
    41  // Returns the new allocated LRU Cache and an error
    42  func NewLRUCache(capacity int) (c *LRUCache, err error) {
    43  	if capacity <= 0 {
    44  		return nil, fmt.Errorf("capacity must be a positive integer")
    45  	}
    46  	c = new(LRUCache)
    47  	c.capacity = capacity
    48  	if capacity <= maxPreallocateCapacity {
    49  		c.entries = make(map[interface{}]interface{}, capacity)
    50  		c.lruElements = make(map[interface{}]*list.Element, capacity)
    51  	} else {
    52  		c.entries = make(map[interface{}]interface{})
    53  		c.lruElements = make(map[interface{}]*list.Element)
    54  	}
    55  	c.lruKeys = list.New()
    56  	return
    57  }
    58  
    59  // Get returns the cache value associated with key
    60  //
    61  // Parameters:
    62  //  * `key` - the key to retrieve
    63  //
    64  // Returns the value associated with key and a bool that is `false`
    65  // if the key was not found
    66  func (c *LRUCache) Get(key interface{}) (value interface{}, ok bool) {
    67  	var element *list.Element
    68  	c.cacheLock.RLock()
    69  	value, ok = c.entries[key]
    70  	if ok {
    71  		element, ok = c.lruElements[key]
    72  	}
    73  	c.cacheLock.RUnlock()
    74  	if ok {
    75  		c.cacheLock.Lock()
    76  		c.lruKeys.MoveToFront(element)
    77  		c.cacheLock.Unlock()
    78  	} else {
    79  		value = nil
    80  	}
    81  	return value, ok
    82  }
    83  
    84  // Put puts a value in cache associated with key
    85  //
    86  // Parameters:
    87  //  * `key` - the key to put
    88  //  * `value` - the value to put
    89  func (c *LRUCache) Put(key interface{}, value interface{}) {
    90  	c.cacheLock.Lock()
    91  	_, ok := c.entries[key]
    92  	if !ok {
    93  		// delete in advance to avoid increasing map capacity
    94  		if c.lruKeys.Len() == c.capacity {
    95  			back := c.lruKeys.Back()
    96  			if back != nil {
    97  				value := c.lruKeys.Remove(back)
    98  				delete(c.lruElements, back)
    99  				delete(c.entries, value)
   100  			}
   101  		}
   102  		element := c.lruKeys.PushFront(key)
   103  		c.lruElements[key] = element
   104  	} else {
   105  		existingElement, okElement := c.lruElements[key]
   106  		if okElement {
   107  			c.lruKeys.MoveToFront(existingElement)
   108  		}
   109  	}
   110  	c.entries[key] = value
   111  	c.cacheLock.Unlock()
   112  }
   113  
   114  // Delete deletes the cache entry associated with key
   115  //
   116  // Parameters:
   117  //  * `key` - the key to delete
   118  func (c *LRUCache) Delete(key interface{}) {
   119  	c.cacheLock.RLock()
   120  	_, ok := c.entries[key]
   121  	c.cacheLock.RUnlock()
   122  	if ok {
   123  		c.cacheLock.Lock()
   124  		element, okElement := c.lruElements[key]
   125  		if okElement {
   126  			delete(c.lruElements, key)
   127  			c.lruKeys.Remove(element)
   128  		}
   129  		delete(c.entries, key)
   130  		c.cacheLock.Unlock()
   131  	}
   132  }
   133  
   134  // ToMap returns the current cache entries copied into a map
   135  func (c *LRUCache) ToMap() map[interface{}]interface{} {
   136  	ret := make(map[interface{}]interface{})
   137  	c.cacheLock.RLock()
   138  	for k, v := range c.entries {
   139  		ret[k] = v
   140  	}
   141  	c.cacheLock.RUnlock()
   142  	return ret
   143  }