/
contextMutex.go
83 lines (69 loc) · 1.78 KB
/
contextMutex.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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
package lock
import (
"context"
"github.com/Invicton-Labs/go-stackerr"
)
// CtxMutex is a mutex where Lock operations use a context
// that can be cancelled/deadlined to terminate the lock
// attempt.
type CtxMutex interface {
// Lock will wait until either the mutex can be locked, or
// the context is done (cancelled/deadlined). If the lock
// succeeds, it will return nil. If the context is done,
// it will return a stack-wrapped version of the context's error.
Lock(ctx context.Context) (err stackerr.Error)
// TryLock will attempt to lock the mutex, but will not
// wait if it cannot immediately do so.
TryLock() (locked bool)
// Unlock will unlock the mutex, and will panic if the mutex
// is not currently locked.
Unlock()
// TryUnlock will attempt to unlock the mutex, but will not
// panic if the mutex is not currently locked.
TryUnlock() (unlocked bool)
}
type ctxMutex struct {
ch chan struct{}
}
// NewCtxMutex creates a new CtxMutex
func NewCtxMutex() CtxMutex {
return &ctxMutex{
ch: make(chan struct{}, 1),
}
}
func (mu *ctxMutex) Lock(ctx context.Context) (err stackerr.Error) {
select {
case <-ctx.Done():
return stackerr.Wrap(ctx.Err())
case mu.ch <- struct{}{}:
return nil
}
}
func (mu *ctxMutex) TryLock() (locked bool) {
select {
case mu.ch <- struct{}{}:
return true
default:
return false
}
}
func (mu *ctxMutex) Unlock() {
select {
case <-mu.ch:
return
default:
panic("unlock of unlocked mutex")
}
}
func (mu *ctxMutex) TryUnlock() (unlocked bool) {
select {
case <-mu.ch:
return true
default:
return false
}
}
// Locked will return whether the mutex is currently locked.
func (mu *ctxMutex) Locked() (locked bool) {
return len(mu.ch) > 0
}