bitswap: finish unsubscribing from the pubsub instance before shutting it down
Otherwise, we'll deadlock and leak a goroutine. This fix is kind of crappy but modifying the pubsub library would have been worse (and, really, it *is* reasonable to say "don't use the pubsub instance after shutting it down"). License: MIT Signed-off-by: Steven Allen <steven@stebalien.com>
Steven Allen committed
Feb 9, 2018 at 12:19 UTC
53958266d659acd2df33466c8d20ad46751322e3
1 file changed
+43
-3
exchange/bitswap/notifications/notifications.go
+43
-3
@@ -2,6 +2,7 @@ package notifications
2
3
import (
4
"context"
5
+ "sync"
6
7
blocks "gx/ipfs/Qmej7nf81hi2x2tvjRBF3mcp74sQyuDH4VMYDGd1YtXjb2/go-block-format"
8
@@ -18,18 +19,33 @@ type PubSub interface {
19
}
20
21
func New() PubSub {
21
- return &impl{*pubsub.New(bufferSize)}
22
+ return &impl{
23
+ wrapped: *pubsub.New(bufferSize),
24
+ cancel: make(chan struct{}),
25
+ }
26
}
27
28
type impl struct {
29
wrapped pubsub.PubSub
30
+
31
+ // These two fields make up a shutdown "lock".
32
+ // We need them as calling, e.g., `Unsubscribe` after calling `Shutdown`
33
+ // blocks forever and fixing this in pubsub would be rather invasive.
34
+ cancel chan struct{}
35
+ wg sync.WaitGroup
36
}
37
38
func (ps *impl) Publish(block blocks.Block) {
39
ps.wrapped.Pub(block, block.Cid().KeyString())
40
}
41
42
+// Not safe to call more than once.
43
func (ps *impl) Shutdown() {
44
+ // Interrupt in-progress subscriptions.
45
+ close(ps.cancel)
46
+ // Wait for them to finish.
47
+ ps.wg.Wait()
48
+ // shutdown the pubsub.
49
ps.wrapped.Shutdown()
50
}
51
@@ -44,12 +60,34 @@ func (ps *impl) Subscribe(ctx context.Context, keys ...*cid.Cid) <-chan blocks.B
60
close(blocksCh)
61
return blocksCh
62
}
63
+
64
+ // prevent shutdown
65
+ ps.wg.Add(1)
66
+
67
+ // check if shutdown *after* preventing shutdowns.
68
+ select {
69
+ case <-ps.cancel:
70
+ // abort, allow shutdown to continue.
71
+ ps.wg.Done()
72
+ close(blocksCh)
73
+ return blocksCh
74
+ default:
75
+ }
76
+
77
ps.wrapped.AddSubOnceEach(valuesCh, toStrings(keys)...)
78
go func() {
49
- defer close(blocksCh)
50
- defer ps.wrapped.Unsub(valuesCh) // with a len(keys) buffer, this is an optimization
79
+ defer func() {
80
+ ps.wrapped.Unsub(valuesCh)
81
+ close(blocksCh)
82
+
83
+ // Unblock shutdown.
84
+ ps.wg.Done()
85
+ }()
86
+
87
for {
88
select {
89
+ case <-ps.cancel:
90
+ return
91
case <-ctx.Done():
92
return
93
case val, ok := <-valuesCh:
@@ -61,6 +99,8 @@ func (ps *impl) Subscribe(ctx context.Context, keys ...*cid.Cid) <-chan blocks.B
99
return
100
}
101
select {
102
+ case <-ps.cancel:
103
+ return
104
case <-ctx.Done():
105
return
106
case blocksCh <- block: // continue