go.uber.org/cadence@v1.2.9/internal/common/cache/lru.go (about)

     1  // Copyright (c) 2017 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package cache
    22  
    23  import (
    24  	"container/list"
    25  	"errors"
    26  	"sync"
    27  	"time"
    28  )
    29  
    30  var (
    31  	// ErrCacheFull is returned if Put fails due to cache being filled with pinned elements
    32  	ErrCacheFull = errors.New("Cache capacity is fully occupied with pinned elements")
    33  )
    34  
    35  // lru is a concurrent fixed size cache that evicts elements in lru order
    36  type lru struct {
    37  	mut      sync.Mutex
    38  	byAccess *list.List
    39  	byKey    map[string]*list.Element
    40  	maxSize  int
    41  	ttl      time.Duration
    42  	pin      bool
    43  	rmFunc   RemovedFunc
    44  	// We use this instead of time.Now() in order to make testing easier
    45  	now func() time.Time
    46  }
    47  
    48  // New creates a new cache with the given options
    49  func New(maxSize int, opts *Options) Cache {
    50  	if opts == nil {
    51  		opts = &Options{}
    52  	}
    53  
    54  	return &lru{
    55  		byAccess: list.New(),
    56  		byKey:    make(map[string]*list.Element, opts.InitialCapacity),
    57  		ttl:      opts.TTL,
    58  		maxSize:  maxSize,
    59  		pin:      opts.Pin,
    60  		rmFunc:   opts.RemovedFunc,
    61  		now:      time.Now,
    62  	}
    63  }
    64  
    65  // NewLRU creates a new LRU cache of the given size, setting initial capacity
    66  // to the max size
    67  func NewLRU(maxSize int) Cache {
    68  	return New(maxSize, nil)
    69  }
    70  
    71  // NewLRUWithInitialCapacity creates a new LRU cache with an initial capacity
    72  // and a max size
    73  func NewLRUWithInitialCapacity(initialCapacity, maxSize int) Cache {
    74  	return New(maxSize, &Options{
    75  		InitialCapacity: initialCapacity,
    76  	})
    77  }
    78  
    79  // Exist checks if a given key exists in the cache
    80  func (c *lru) Exist(key string) bool {
    81  	c.mut.Lock()
    82  	defer c.mut.Unlock()
    83  	_, ok := c.byKey[key]
    84  	return ok
    85  }
    86  
    87  // Get retrieves the value stored under the given key
    88  func (c *lru) Get(key string) interface{} {
    89  	c.mut.Lock()
    90  	defer c.mut.Unlock()
    91  
    92  	elt := c.byKey[key]
    93  	if elt == nil {
    94  		return nil
    95  	}
    96  
    97  	cacheEntry := elt.Value.(*cacheEntry)
    98  
    99  	if c.pin {
   100  		cacheEntry.refCount++
   101  	}
   102  
   103  	if cacheEntry.refCount == 0 && !cacheEntry.expiration.IsZero() && c.now().After(cacheEntry.expiration) {
   104  		// Entry has expired
   105  		if c.rmFunc != nil {
   106  			go c.rmFunc(cacheEntry.value)
   107  		}
   108  		c.byAccess.Remove(elt)
   109  		delete(c.byKey, cacheEntry.key)
   110  		return nil
   111  	}
   112  
   113  	c.byAccess.MoveToFront(elt)
   114  	return cacheEntry.value
   115  }
   116  
   117  // Put puts a new value associated with a given key, returning the existing value (if present)
   118  func (c *lru) Put(key string, value interface{}) interface{} {
   119  	if c.pin {
   120  		panic("Cannot use Put API in Pin mode. Use Delete and PutIfNotExist if necessary")
   121  	}
   122  	val, _ := c.putInternal(key, value, true)
   123  	return val
   124  }
   125  
   126  // PutIfNotExist puts a value associated with a given key if it does not exist
   127  func (c *lru) PutIfNotExist(key string, value interface{}) (interface{}, error) {
   128  	existing, err := c.putInternal(key, value, false)
   129  	if err != nil {
   130  		return nil, err
   131  	}
   132  
   133  	if existing == nil {
   134  		// This is a new value
   135  		return value, err
   136  	}
   137  
   138  	return existing, err
   139  }
   140  
   141  // Delete deletes a key, value pair associated with a key
   142  func (c *lru) Delete(key string) {
   143  	c.mut.Lock()
   144  	defer c.mut.Unlock()
   145  
   146  	elt := c.byKey[key]
   147  	if elt != nil {
   148  		entry := c.byAccess.Remove(elt).(*cacheEntry)
   149  		if c.rmFunc != nil {
   150  			go c.rmFunc(entry.value)
   151  		}
   152  		delete(c.byKey, key)
   153  	}
   154  }
   155  
   156  // Release decrements the ref count of a pinned element.
   157  func (c *lru) Release(key string) {
   158  	c.mut.Lock()
   159  	defer c.mut.Unlock()
   160  
   161  	elt := c.byKey[key]
   162  	cacheEntry := elt.Value.(*cacheEntry)
   163  	cacheEntry.refCount--
   164  }
   165  
   166  // Size returns the number of entries currently in the lru, useful if cache is not full
   167  func (c *lru) Size() int {
   168  	c.mut.Lock()
   169  	defer c.mut.Unlock()
   170  
   171  	return len(c.byKey)
   172  }
   173  
   174  // Put puts a new value associated with a given key, returning the existing value (if present)
   175  // allowUpdate flag is used to control overwrite behavior if the value exists
   176  func (c *lru) putInternal(key string, value interface{}, allowUpdate bool) (interface{}, error) {
   177  	c.mut.Lock()
   178  	defer c.mut.Unlock()
   179  
   180  	elt := c.byKey[key]
   181  	if elt != nil {
   182  		entry := elt.Value.(*cacheEntry)
   183  		existing := entry.value
   184  		if allowUpdate {
   185  			entry.value = value
   186  		}
   187  		if c.ttl != 0 {
   188  			entry.expiration = c.now().Add(c.ttl)
   189  		}
   190  		c.byAccess.MoveToFront(elt)
   191  		if c.pin {
   192  			entry.refCount++
   193  		}
   194  		return existing, nil
   195  	}
   196  
   197  	entry := &cacheEntry{
   198  		key:   key,
   199  		value: value,
   200  	}
   201  
   202  	if c.pin {
   203  		entry.refCount++
   204  	}
   205  
   206  	if c.ttl != 0 {
   207  		entry.expiration = c.now().Add(c.ttl)
   208  	}
   209  
   210  	c.byKey[key] = c.byAccess.PushFront(entry)
   211  	if len(c.byKey) == c.maxSize {
   212  		oldest := c.byAccess.Back().Value.(*cacheEntry)
   213  
   214  		if oldest.refCount > 0 {
   215  			// Cache is full with pinned elements
   216  			// revert the insert and return
   217  			c.byAccess.Remove(c.byAccess.Front())
   218  			delete(c.byKey, key)
   219  			return nil, ErrCacheFull
   220  		}
   221  
   222  		c.byAccess.Remove(c.byAccess.Back())
   223  		if c.rmFunc != nil {
   224  			go c.rmFunc(oldest.value)
   225  		}
   226  		delete(c.byKey, oldest.key)
   227  	}
   228  
   229  	return nil, nil
   230  }
   231  
   232  type cacheEntry struct {
   233  	key        string
   234  	expiration time.Time
   235  	value      interface{}
   236  	refCount   int
   237  }