summaryrefslogtreecommitdiff
path: root/syncutil/cachemap.go
blob: 8eab4bc690e4aeeb80f64ec1554c7fc539b8cba7 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
// Copyright (C) 2023  Luke Shumaker <lukeshu@lukeshu.com>
//
// SPDX-License-Identifier: GPL-2.0-or-later

package syncutil

import (
	"sync"

	"git.lukeshu.com/go/containers/typedsync"
)

type cacheVal[V any] struct {
	wg sync.WaitGroup
	v  V
}

// The techniques used by CacheMap are similar to the techniques used
// by encoding/json's type cache.

// A CacheMap is similar to a typedsync.Map, but also has a
// .LoadOrCompute sibling to .LoadOrStore.  This is useful for caching
// the results of computation where it is undesirable for concurrent
// calls to duplicate work.
//
// Compared to a plain typedsync.Map, a CacheMap has both more space
// overhead and more time overhead.
type CacheMap[K mapkey, V any] struct {
	inner typedsync.Map[K, *cacheVal[V]]
	pool  typedsync.Pool[*cacheVal[V]]
}

func (m *CacheMap[K, V]) Delete(key K) {
	m.inner.Delete(key)
}

// Load returns the value stored in the map for a key.  If the value
// for that key is actively being computed by LoadOrCompute, this
// blocks until the value has been computed.
func (m *CacheMap[K, V]) Load(key K) (value V, ok bool) {
	_value, ok := m.inner.Load(key)
	if !ok {
		var zero V
		return zero, false
	}
	_value.wg.Wait()
	return _value.v, true
}

// LoadAndDelete deletes the value for a key, returning the previous
// value if any.  The loaded result reports whether the key was
// present.  If the value for that key was actively being computed by
// LoadOrCompute, this immediately removes the partial value from the
// CacheMap, but does not return until the computation has finished.
func (m *CacheMap[K, V]) LoadAndDelete(key K) (value V, loaded bool) {
	_value, ok := m.inner.LoadAndDelete(key)
	if !ok {
		var zero V
		return zero, false
	}
	_value.wg.Wait()
	return _value.v, true
}

// LoadOrStore returns the existing value for the key if
// present. Otherwise, it stores and returns the given value.  The
// loaded result is true if the value was loaded, false if stored.  If
// the value for that key is actively being computed by LoadOrCompute,
// this blocks until the value has been computed.
func (m *CacheMap[K, V]) LoadOrStore(key K, value V) (actual V, loaded bool) {
	_value, _ := m.pool.Get()
	if _value == nil {
		_value = new(cacheVal[V])
	}
	_value.v = value
	_actual, loaded := m.inner.LoadOrStore(key, _value)
	if loaded {
		*_value = cacheVal[V]{}
		m.pool.Put(_value)
		_actual.wg.Wait()
	}
	return _actual.v, loaded
}

// LoadOrCompute returns the existing value for the key if present.
// Otherwise, it computes and stores a value using the given function.
// The loaded result is true if the value was loaded, false if
// computed.  If a prior call to LoadOrCompute is still computing the
// value for that key, a latter call blocks until the computation is
// complete, and then returns the initial call's value.
func (m *CacheMap[K, V]) LoadOrCompute(key K, fn func(K) V) (actual V, loaded bool) {
	_value, _ := m.pool.Get()
	if _value == nil {
		_value = new(cacheVal[V])
	}
	_value.wg.Add(1)
	_actual, loaded := m.inner.LoadOrStore(key, _value)
	if loaded {
		*_value = cacheVal[V]{}
		m.pool.Put(_value)
		_actual.wg.Wait()
	} else {
		_actual.v = fn(key)
		_actual.wg.Done()
	}
	return _actual.v, loaded
}

func (m *CacheMap[K, V]) Store(key K, value V) {
	_value, _ := m.pool.Get()
	if _value == nil {
		_value = new(cacheVal[V])
	}
	_value.v = value
	m.inner.Store(key, _value)
}