-
-
Notifications
You must be signed in to change notification settings - Fork 623
/
r.go
101 lines (93 loc) · 1.89 KB
/
r.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
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
package alloclim
import (
"context"
"errors"
"fmt"
"sync"
"github.com/anacrolix/chansync"
"github.com/anacrolix/log"
)
type Reservation struct {
l *Limiter
n int64
releaseOnce sync.Once
mu sync.Mutex
granted chansync.SetOnce
cancelled chansync.SetOnce
}
// Releases the alloc claim if the reservation has been granted. Does nothing if it was cancelled.
// Otherwise panics.
func (me *Reservation) Release() {
me.mu.Lock()
defer me.mu.Unlock()
switch {
default:
panic("not resolved")
case me.cancelled.IsSet():
return
case me.granted.IsSet():
}
me.releaseOnce.Do(func() {
me.l.addValue(me.n)
})
}
// Cancel the reservation, returns false if it was already granted. You must still release if that's
// the case. See Drop.
func (me *Reservation) Cancel() bool {
me.mu.Lock()
defer me.mu.Unlock()
if me.granted.IsSet() {
return false
}
if me.cancelled.Set() {
go me.l.doWakes()
}
return true
}
// If the reservation is granted, release it, otherwise cancel the reservation.
func (me *Reservation) Drop() {
me.mu.Lock()
defer me.mu.Unlock()
if me.granted.IsSet() {
me.releaseOnce.Do(func() {
me.l.addValue(me.n)
})
return
}
if me.cancelled.Set() {
go me.l.doWakes()
}
}
func (me *Reservation) wake() bool {
me.mu.Lock()
defer me.mu.Unlock()
if me.cancelled.IsSet() {
return false
}
return me.granted.Set()
}
func (me *Reservation) Wait(ctx context.Context) error {
if me.n > me.l.Max {
return log.WithLevel(
log.Warning,
fmt.Errorf("reservation for %v exceeds limiter max %v", me.n, me.l.Max),
)
}
select {
case <-ctx.Done():
case <-me.granted.Done():
case <-me.cancelled.Done():
}
defer me.mu.Unlock()
me.mu.Lock()
switch {
case me.granted.IsSet():
return nil
case me.cancelled.IsSet():
return errors.New("reservation cancelled")
case ctx.Err() != nil:
return ctx.Err()
default:
panic("unexpected")
}
}