-
Notifications
You must be signed in to change notification settings - Fork 1.9k
/
funcs.go
89 lines (78 loc) · 1.92 KB
/
funcs.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
package structs
import (
"sync"
"github.com/hashicorp/nomad/nomad/structs"
)
// AllocBroadcaster implements an allocation broadcast channel.
// The zero value is a usable unbuffered channel.
type AllocBroadcaster struct {
m sync.Mutex
listeners map[int]chan<- *structs.Allocation // lazy init
nextId int
capacity int
closed bool
}
// NewAllocBroadcaster returns a new AllocBroadcaster with the given capacity (0 means unbuffered).
func NewAllocBroadcaster(n int) *AllocBroadcaster {
return &AllocBroadcaster{capacity: n}
}
// AllocListener implements a listening endpoint for an allocation broadcast channel.
type AllocListener struct {
// Ch receives the broadcast messages.
Ch <-chan *structs.Allocation
b *AllocBroadcaster
id int
}
// Send broadcasts a message to the channel. Send returns whether the message
// was sent to all channels.
func (b *AllocBroadcaster) Send(v *structs.Allocation) bool {
b.m.Lock()
defer b.m.Unlock()
if b.closed {
return false
}
sent := true
for _, l := range b.listeners {
select {
case l <- v:
default:
sent = false
}
}
return sent
}
// Close closes the channel, disabling the sending of further messages.
func (b *AllocBroadcaster) Close() {
b.m.Lock()
defer b.m.Unlock()
if b.closed {
return
}
b.closed = true
for _, l := range b.listeners {
close(l)
}
}
// Listen returns a Listener for the broadcast channel.
func (b *AllocBroadcaster) Listen() *AllocListener {
b.m.Lock()
defer b.m.Unlock()
if b.listeners == nil {
b.listeners = make(map[int]chan<- *structs.Allocation)
}
for b.listeners[b.nextId] != nil {
b.nextId++
}
ch := make(chan *structs.Allocation, b.capacity)
if b.closed {
close(ch)
}
b.listeners[b.nextId] = ch
return &AllocListener{ch, b, b.nextId}
}
// Close closes the Listener, disabling the receival of further messages.
func (l *AllocListener) Close() {
l.b.m.Lock()
defer l.b.m.Unlock()
delete(l.b.listeners, l.id)
}