1// Copyright (C) 2020 Storj Labs, Inc. 2// See LICENSE for copying information. 3 4package drpccache 5 6import ( 7 "context" 8 "sync" 9) 10 11type cacheKey struct{} 12 13// Cache is a per stream cache. 14type Cache struct { 15 mu sync.Mutex 16 values map[interface{}]interface{} 17} 18 19// New returns a new cache. 20func New() *Cache { return &Cache{} } 21 22// FromContext returns a cache from a context. 23// 24// Example usage: 25// 26// cache := drpccache.FromContext(stream.Context()) 27// if cache != nil { 28// value := cache.LoadOrCreate("initialized", func() (interface{}) { 29// return 42 30// }) 31// } 32func FromContext(ctx context.Context) *Cache { 33 cache, _ := ctx.Value(cacheKey{}).(*Cache) 34 return cache 35} 36 37// WithContext returns a context with the value cache associated with the context. 38func WithContext(parent context.Context, cache *Cache) context.Context { 39 return context.WithValue(parent, cacheKey{}, cache) 40} 41 42// init ensures that the values map exist. 43func (cache *Cache) init() { 44 if cache.values == nil { 45 cache.values = map[interface{}]interface{}{} 46 } 47} 48 49// Clear clears the cache. 50func (cache *Cache) Clear() { 51 cache.mu.Lock() 52 defer cache.mu.Unlock() 53 54 cache.values = nil 55} 56 57// Store sets the value at a key. 58func (cache *Cache) Store(key, value interface{}) { 59 cache.mu.Lock() 60 defer cache.mu.Unlock() 61 62 cache.init() 63 cache.values[key] = value 64} 65 66// Load returns the value with the given key. 67func (cache *Cache) Load(key interface{}) interface{} { 68 cache.mu.Lock() 69 defer cache.mu.Unlock() 70 71 if cache.values == nil { 72 return nil 73 } 74 75 return cache.values[key] 76} 77 78// LoadOrCreate returns the value with the given key. 79func (cache *Cache) LoadOrCreate(key interface{}, fn func() interface{}) interface{} { 80 cache.mu.Lock() 81 defer cache.mu.Unlock() 82 83 cache.init() 84 85 value, ok := cache.values[key] 86 if !ok { 87 value = fn() 88 cache.values[key] = value 89 } 90 91 return value 92} 93