forked from grailbio/base
-
Notifications
You must be signed in to change notification settings - Fork 0
/
once.go
69 lines (60 loc) · 1.85 KB
/
once.go
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
// Copyright 2018 GRAIL, Inc. All rights reserved.
// Use of this source code is governed by the Apache 2.0
// license that can be found in the LICENSE file.
// Package once contains utilities for managing actions that
// must be performed exactly once.
package once
import (
"sync"
"sync/atomic"
)
// Task manages a computation that must be run at most once.
// It's similar to sync.Once, except it also handles and returns errors.
type Task struct {
mu sync.Mutex
done uint32
err error
}
// Do run the function do at most once. Successive invocations of Do
// guarantee exactly one invocation of the function do. Do returns
// the error of do's invocation.
func (o *Task) Do(do func() error) error {
if atomic.LoadUint32(&o.done) == 1 {
return o.err
}
o.mu.Lock()
defer o.mu.Unlock()
if atomic.LoadUint32(&o.done) == 0 {
o.err = do()
atomic.StoreUint32(&o.done, 1)
}
return o.err
}
// Done returns whether the task is done.
func (o *Task) Done() bool {
o.mu.Lock()
defer o.mu.Unlock()
return 1 == atomic.LoadUint32(&o.done)
}
// Reset resets the task effectively making it possible for `Do` to invoke the underlying do func again.
// Reset will only reset the task if it was already completed.
func (o *Task) Reset() {
o.mu.Lock()
defer o.mu.Unlock()
atomic.CompareAndSwapUint32(&o.done, 1, 0)
}
// Map coordinates actions that must happen exactly once, keyed
// by user-defined keys.
type Map sync.Map
// Perform the provided action named by a key. Do invokes the action
// exactly once for each key, and returns any errors produced by the
// provided action.
func (m *Map) Do(key interface{}, do func() error) error {
taskv, _ := (*sync.Map)(m).LoadOrStore(key, new(Task))
task := taskv.(*Task)
return task.Do(do)
}
// Forget forgets past computations associated with the provided key.
func (m *Map) Forget(key interface{}) {
(*sync.Map)(m).Delete(key)
}