forked from rancher/rancher
-
Notifications
You must be signed in to change notification settings - Fork 0
/
broadcaster.go
122 lines (104 loc) · 2.13 KB
/
broadcaster.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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
package broadcast
import (
"context"
"sync"
)
type ConnectFunc func() (chan map[string]interface{}, error)
type Broadcaster struct {
sync.Mutex
running bool
subs map[chan map[string]interface{}]struct{}
}
func (b *Broadcaster) Subscribe(ctx context.Context, connect ConnectFunc) (chan map[string]interface{}, error) {
b.Lock()
defer b.Unlock()
if !b.running {
if err := b.start(connect); err != nil {
return nil, err
}
}
sub := make(chan map[string]interface{}, 100)
if b.subs == nil {
b.subs = map[chan map[string]interface{}]struct{}{}
}
b.subs[sub] = struct{}{}
go func() {
<-ctx.Done()
b.Lock()
b.unsub(sub)
b.Unlock()
}()
return sub, nil
}
func (b *Broadcaster) unsub(sub chan map[string]interface{}) {
if _, ok := b.subs[sub]; ok {
close(sub)
delete(b.subs, sub)
}
}
func (b *Broadcaster) start(connect ConnectFunc) error {
c, err := connect()
if err != nil {
return err
}
go b.stream(c)
b.running = true
return nil
}
func (b *Broadcaster) stream(input chan map[string]interface{}) {
for item := range input {
b.Lock()
for sub := range b.subs {
newItem := cloneMap(item)
select {
case sub <- newItem:
default:
// Slow consumer, drop
go b.unsub(sub)
}
}
b.Unlock()
}
b.Lock()
for sub := range b.subs {
b.unsub(sub)
}
b.running = false
b.Unlock()
}
func cloneMap(data map[string]interface{}) map[string]interface{} {
if data == nil {
return nil
}
result := map[string]interface{}{}
for k, v := range data {
result[k] = cloneValue(v)
}
return result
}
func cloneValue(v interface{}) interface{} {
switch t := v.(type) {
case []interface{}:
return cloneSlice(t)
case []map[string]interface{}:
return cloneMapSlice(t)
case map[string]interface{}:
return cloneMap(t)
default:
return v
}
}
func cloneMapSlice(data []map[string]interface{}) []interface{} {
result := make([]interface{}, len(data))
for i := range data {
result[i] = cloneValue(data[i])
}
return result
}
func cloneSlice(data []interface{}) []interface{} {
result := make([]interface{}, len(data))
for i := range data {
result[i] = cloneValue(data[i])
}
return result
}