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