refac(exch:bitswap) always notify strategy when message sent
Brian Tiger Chow committed
Sep 18, 2014 at 19:01 UTC
d82a2517d1ccaa8fa5680836ccc2d7bc6ff59912
1 file changed
+23
-22
exchange/bitswap/bitswap.go
+23
-22
@@ -79,6 +79,9 @@ type bitswap struct {
79
// GetBlock attempts to retrieve a particular block from peers, within timeout.
80
func (bs *bitswap) Block(k u.Key, timeout time.Duration) (
81
*blocks.Block, error) {
82
+ ctx, _ := context.WithTimeout(context.Background(), timeout)
83
+
84
+ // TODO replace timeout with ctx in routing interface
85
begin := time.Now()
86
tleft := timeout - time.Now().Sub(begin)
87
provs_ch := bs.routing.FindProvidersAsync(k, 20, timeout)
@@ -90,7 +93,7 @@ func (bs *bitswap) Block(k u.Key, timeout time.Duration) (
93
go func() {
94
for p := range provs_ch {
95
go func(pr *peer.Peer) {
93
- blk, err := bs.getBlock(k, pr, tleft)
96
+ blk, err := bs.getBlock(ctx, k, pr)
97
if err != nil {
98
return
99
}
@@ -111,19 +114,14 @@ func (bs *bitswap) Block(k u.Key, timeout time.Duration) (
114
}
115
}
116
114
-func (bs *bitswap) getBlock(k u.Key, p *peer.Peer, timeout time.Duration) (*blocks.Block, error) {
117
+func (bs *bitswap) getBlock(ctx context.Context, k u.Key, p *peer.Peer) (*blocks.Block, error) {
118
116
- ctx, _ := context.WithTimeout(context.Background(), timeout)
119
blockChannel := bs.notifications.Subscribe(ctx, k)
120
121
message := bsmsg.New()
122
message.AppendWanted(k)
123
122
- // FIXME(brian): register the accountant on the service wrapper to ensure
123
- // that accounting is _always_ performed when SendMessage and
124
- // ReceiveMessage are called
125
- bs.sender.SendMessage(ctx, p, message)
126
- bs.strategy.MessageSent(p, message)
124
+ bs.send(ctx, p, message)
125
126
block, ok := <-blockChannel
127
if !ok {
@@ -132,11 +130,13 @@ func (bs *bitswap) getBlock(k u.Key, p *peer.Peer, timeout time.Duration) (*bloc
130
return &block, nil
131
}
132
135
-func (bs *bitswap) sendToPeersThatWant(block blocks.Block) {
133
+func (bs *bitswap) sendToPeersThatWant(ctx context.Context, block blocks.Block) {
134
for _, p := range bs.strategy.Peers() {
135
if bs.strategy.BlockIsWantedByPeer(block.Key(), p) {
136
if bs.strategy.ShouldSendBlockToPeer(block.Key(), p) {
139
- go bs.send(p, block)
137
+ message := bsmsg.New()
138
+ message.AppendBlock(block)
139
+ go bs.send(ctx, p, message)
140
}
141
}
142
}
@@ -145,16 +145,17 @@ func (bs *bitswap) sendToPeersThatWant(block blocks.Block) {
145
// HasBlock announces the existance of a block to bitswap, potentially sending
146
// it to peers (Partners) whose WantLists include it.
147
func (bs *bitswap) HasBlock(blk blocks.Block) error {
148
- go bs.sendToPeersThatWant(blk)
148
+ ctx := context.TODO()
149
+ go bs.sendToPeersThatWant(ctx, blk)
150
return bs.routing.Provide(blk.Key())
151
}
152
153
// TODO(brian): handle errors
154
func (bs *bitswap) ReceiveMessage(
154
- ctx context.Context, sender *peer.Peer, incoming bsmsg.BitSwapMessage) (
155
+ ctx context.Context, p *peer.Peer, incoming bsmsg.BitSwapMessage) (
156
*peer.Peer, bsmsg.BitSwapMessage, error) {
157
157
- bs.strategy.MessageReceived(sender, incoming)
158
+ bs.strategy.MessageReceived(p, incoming)
159
160
if incoming.Blocks() != nil {
161
for _, block := range incoming.Blocks() {
@@ -165,26 +166,26 @@ func (bs *bitswap) ReceiveMessage(
166
167
if incoming.Wantlist() != nil {
168
for _, key := range incoming.Wantlist() {
168
- if bs.strategy.ShouldSendBlockToPeer(key, sender) {
169
+ if bs.strategy.ShouldSendBlockToPeer(key, p) {
170
block, errBlockNotFound := bs.blockstore.Get(key)
171
if errBlockNotFound != nil {
172
// TODO(brian): log/return the error
173
continue
174
}
174
- go bs.send(sender, *block)
175
+ message := bsmsg.New()
176
+ message.AppendBlock(*block)
177
+ go bs.send(ctx, p, message)
178
}
179
}
180
}
181
return nil, nil, errors.New("TODO implement")
182
}
183
181
-// TODO(brian): get a return value
182
-func (bs *bitswap) send(p *peer.Peer, b blocks.Block) {
183
- message := bsmsg.New()
184
- message.AppendBlock(b)
185
- // FIXME(brian): pass ctx
186
- bs.sender.SendMessage(context.Background(), p, message)
187
- bs.strategy.MessageSent(p, message)
184
+// send strives to ensure that accounting is always performed when a message is
185
+// sent
186
+func (bs *bitswap) send(ctx context.Context, p *peer.Peer, m bsmsg.BitSwapMessage) {
187
+ bs.sender.SendMessage(context.Background(), p, m)
188
+ bs.strategy.MessageSent(p, m)
189
}
190
191
func numBytes(b blocks.Block) int {