give sessiongenerator a master context for easy cancelling
Jeromy committed
Dec 10, 2014 at 02:02 UTC
c79fad2e1256e30d4e1df5473fa07887ea817d7e
1 file changed
+19
-11
exchange/bitswap/testutils.go
+19
-11
@@ -17,24 +17,33 @@ import (
17
18
func NewSessionGenerator(
19
net tn.Network, rs mockrouting.Server) SessionGenerator {
20
+ ctx, cancel := context.WithCancel(context.TODO())
21
return SessionGenerator{
21
- net: net,
22
- rs: rs,
23
- ps: peer.NewPeerstore(),
24
- seq: 0,
22
+ ps: peer.NewPeerstore(),
23
+ net: net,
24
+ rs: rs,
25
+ seq: 0,
26
+ ctx: ctx,
27
+ cancel: cancel,
28
}
29
}
30
31
type SessionGenerator struct {
29
- seq int
30
- net tn.Network
31
- rs mockrouting.Server
32
- ps peer.Peerstore
32
+ seq int
33
+ net tn.Network
34
+ rs mockrouting.Server
35
+ ps peer.Peerstore
36
+ ctx context.Context
37
+ cancel context.CancelFunc
38
+}
39
+
40
+func (g *SessionGenerator) Stop() {
41
+ g.cancel()
42
}
43
44
func (g *SessionGenerator) Next() Instance {
45
g.seq++
37
- return session(g.net, g.rs, g.ps, []byte(string(g.seq)))
46
+ return session(g.ctx, g.net, g.rs, g.ps, []byte(string(g.seq)))
47
}
48
49
func (g *SessionGenerator) Instances(n int) []Instance {
@@ -67,7 +76,7 @@ func (i *Instance) SetBlockstoreLatency(t time.Duration) time.Duration {
76
// NB: It's easy make mistakes by providing the same peer ID to two different
77
// sessions. To safeguard, use the SessionGenerator to generate sessions. It's
78
// just a much better idea.
70
-func session(net tn.Network, rs mockrouting.Server, ps peer.Peerstore, id peer.ID) Instance {
79
+func session(ctx context.Context, net tn.Network, rs mockrouting.Server, ps peer.Peerstore, id peer.ID) Instance {
80
p := ps.WithID(id)
81
82
adapter := net.Adapter(p)
@@ -82,7 +91,6 @@ func session(net tn.Network, rs mockrouting.Server, ps peer.Peerstore, id peer.I
91
}
92
93
const alwaysSendToPeer = true
85
- ctx := context.TODO()
94
95
bs := New(ctx, p, adapter, htc, bstore, alwaysSendToPeer)
96