forked from saiya/dsps
-
Notifications
You must be signed in to change notification settings - Fork 2
/
dispatcher.go
95 lines (77 loc) · 2.37 KB
/
dispatcher.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
package stub
import (
"context"
"sync"
"github.com/m3dev/dsps/server/storage/redis/internal/pubsub"
)
// RedisPubSubDispatcherStub implements RedisPubSubDispatcher
type RedisPubSubDispatcherStub struct {
lock sync.Mutex
isClosed bool
hooks map[pubsub.RedisChannelID][]func(pubsub.RedisPubSubPromise)
promises map[pubsub.RedisChannelID][]pubsub.RedisPubSubPromise
}
// NewRedisPubSubDispatcherStub returns new RedisPubSubDispatcherStub
func NewRedisPubSubDispatcherStub() *RedisPubSubDispatcherStub {
return &RedisPubSubDispatcherStub{
hooks: make(map[pubsub.RedisChannelID][]func(pubsub.RedisPubSubPromise)),
promises: make(map[pubsub.RedisChannelID][]pubsub.RedisPubSubPromise),
}
}
// Shutdown implements RedisPubSubDispatcher
func (d *RedisPubSubDispatcherStub) Shutdown(ctx context.Context) {
func() {
d.lock.Lock()
defer d.lock.Unlock()
d.isClosed = true
}()
d.Reject(pubsub.ErrClosed)
}
// HookAwaitOnce is to register hook to future Await() call.
func (d *RedisPubSubDispatcherStub) HookAwaitOnce(channel pubsub.RedisChannelID, f func(pubsub.RedisPubSubPromise)) {
d.lock.Lock()
defer d.lock.Unlock()
d.hooks[channel] = append(d.hooks[channel], f)
}
// Await implements RedisPubSubDispatcher
func (d *RedisPubSubDispatcherStub) Await(ctx context.Context, channel pubsub.RedisChannelID) (pubsub.RedisPubSubAwaiter, pubsub.AwaitCancelFunc) {
p := pubsub.NewPromise()
var hooks []func(pubsub.RedisPubSubPromise)
defer func() { // Make sure to call hooks on the end of function
// Call hooks outside of lock, otherwise may cause deadlock (recursive lock)
for _, hook := range hooks {
hook(p)
}
}()
d.lock.Lock()
defer d.lock.Unlock()
hooks = d.hooks[channel]
delete(d.hooks, channel)
if d.isClosed {
p.Reject(pubsub.ErrClosed)
return p, func(error) {}
}
d.promises[channel] = append(d.promises[channel], p)
return p, func(err error) {
p.Reject(err)
// Because this is stub implementation, not take care memory leak.
}
}
// Reject all promises
func (d *RedisPubSubDispatcherStub) Reject(err error) {
d.lock.Lock()
defer d.lock.Unlock()
for _, chain := range d.promises {
for _, p := range chain {
p.Reject(err)
}
}
}
// Resolve all promises of the channel
func (d *RedisPubSubDispatcherStub) Resolve(channel pubsub.RedisChannelID) {
d.lock.Lock()
defer d.lock.Unlock()
for _, p := range d.promises[channel] {
p.Resolve()
}
}