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
|
// Copyright (c) Tailscale Inc & contributors
// SPDX-License-Identifier: BSD-3-Clause
package dist
import (
"sync"
"tailscale.com/util/deephash"
)
// MemoizedFn is a function that memoize.Do can call.
type MemoizedFn[T any] func() (T, error)
// Memoize runs MemoizedFns and remembers their results.
type Memoize[O any] struct {
mu sync.Mutex
cond *sync.Cond
outs map[deephash.Sum]O
errs map[deephash.Sum]error
inflight map[deephash.Sum]bool
}
// Do runs fn and returns its result.
// fn is only run once per unique key. Subsequent Do calls with the same key
// return the memoized result of the first call, even if fn is a different
// function.
func (m *Memoize[O]) Do(key any, fn MemoizedFn[O]) (ret O, err error) {
m.mu.Lock()
defer m.mu.Unlock()
if m.cond == nil {
m.cond = sync.NewCond(&m.mu)
m.outs = map[deephash.Sum]O{}
m.errs = map[deephash.Sum]error{}
m.inflight = map[deephash.Sum]bool{}
}
k := deephash.Hash(&key)
for m.inflight[k] {
m.cond.Wait()
}
if err := m.errs[k]; err != nil {
var ret O
return ret, err
}
if ret, ok := m.outs[k]; ok {
return ret, nil
}
m.inflight[k] = true
m.mu.Unlock()
defer func() {
m.mu.Lock()
delete(m.inflight, k)
if err != nil {
m.errs[k] = err
} else {
m.outs[k] = ret
}
m.cond.Broadcast()
}()
ret, err = fn()
if err != nil {
var ret O
return ret, err
}
return ret, nil
}
// once is like memoize, but for functions that don't return non-error values.
type once struct {
m Memoize[any]
}
// Do runs fn.
// fn is only run once per unique key. Subsequent Do calls with the same key
// return the memoized result of the first call, even if fn is a different
// function.
func (o *once) Do(key any, fn func() error) error {
_, err := o.m.Do(key, func() (any, error) {
return nil, fn()
})
return err
}
|
