refactor(bitswap:message) add wrapper for proto
Brian Tiger Chow committed
Sep 11, 2014 at 03:09 UTC
421fe5bf76078e8fc666dce926e3c10baa44dbd6
3 files changed
+122
-12
bitswap/bitswap.go
+9
-12
@@ -115,13 +115,12 @@ func (bs *BitSwap) GetBlock(k u.Key, timeout time.Duration) (
115
func (bs *BitSwap) getBlock(k u.Key, p *peer.Peer, timeout time.Duration) ([]byte, error) {
116
u.DOut("[%s] getBlock '%s' from [%s]\n", bs.peer.ID.Pretty(), k.Pretty(), p.ID.Pretty())
117
118
- pmes := new(PBMessage)
119
- pmes.Wantlist = []string{string(k)}
118
+ message := newMessage()
119
+ message.AppendWanted(k)
120
121
after := time.After(timeout)
122
resp := bs.listener.Listen(string(k), 1, timeout)
123
- smes := swarm.NewMessage(p, pmes)
124
- bs.meschan.Outgoing <- smes
123
+ bs.meschan.Outgoing <- message.ToSwarm(p)
124
125
select {
126
case resp_mes := <-resp:
@@ -149,11 +148,9 @@ func (bs *BitSwap) HaveBlock(blk *blocks.Block) error {
148
}
149
150
func (bs *BitSwap) SendBlock(p *peer.Peer, b *blocks.Block) {
152
- pmes := new(PBMessage)
153
- pmes.Blocks = [][]byte{b.Data}
154
-
155
- swarm_mes := swarm.NewMessage(p, pmes)
156
- bs.meschan.Outgoing <- swarm_mes
151
+ message := newMessage()
152
+ message.AppendBlock(b)
153
+ bs.meschan.Outgoing <- message.ToSwarm(p)
154
}
155
156
func (bs *BitSwap) handleMessages() {
@@ -257,14 +254,14 @@ func (bs *BitSwap) GetLedger(p *peer.Peer) *Ledger {
254
}
255
256
func (bs *BitSwap) SendWantList(wl KeySet) error {
260
- pmes := new(PBMessage)
257
+ message := newMessage()
258
for k, _ := range wl {
262
- pmes.Wantlist = append(pmes.Wantlist, string(k))
259
+ message.AppendWanted(k)
260
}
261
262
// Lets just ping everybody all at once
263
for _, ledger := range bs.partners {
267
- bs.meschan.Outgoing <- swarm.NewMessage(ledger.Partner, pmes)
264
+ bs.meschan.Outgoing <- message.ToSwarm(ledger.Partner)
265
}
266
267
return nil
bitswap/message.go
new
+38
@@ -0,0 +1,38 @@
1
+package bitswap
2
+
3
+import (
4
+ blocks "github.com/jbenet/go-ipfs/blocks"
5
+ peer "github.com/jbenet/go-ipfs/peer"
6
+ swarm "github.com/jbenet/go-ipfs/swarm"
7
+ u "github.com/jbenet/go-ipfs/util"
8
+)
9
+
10
+// message wraps a proto message for convenience
11
+type message struct {
12
+ pb PBMessage
13
+}
14
+
15
+func newMessageFromProto(pb PBMessage) *message {
16
+ return &message{pb: pb}
17
+}
18
+
19
+func newMessage() *message {
20
+ return new(message)
21
+}
22
+
23
+func (m *message) AppendWanted(k u.Key) {
24
+ m.pb.Wantlist = append(m.pb.Wantlist, string(k))
25
+}
26
+
27
+func (m *message) AppendBlock(b *blocks.Block) {
28
+ m.pb.Blocks = append(m.pb.Blocks, b.Data)
29
+}
30
+
31
+func (m *message) ToProto() *PBMessage {
32
+ cp := m.pb
33
+ return &cp
34
+}
35
+
36
+func (m *message) ToSwarm(p *peer.Peer) *swarm.Message {
37
+ return swarm.NewMessage(p, m.ToProto())
38
+}
bitswap/message_test.go
new
+75
@@ -0,0 +1,75 @@
1
+package bitswap
2
+
3
+import (
4
+ "bytes"
5
+ "testing"
6
+
7
+ blocks "github.com/jbenet/go-ipfs/blocks"
8
+ u "github.com/jbenet/go-ipfs/util"
9
+)
10
+
11
+func TestAppendWanted(t *testing.T) {
12
+ const str = "foo"
13
+ m := newMessage()
14
+ m.AppendWanted(u.Key(str))
15
+
16
+ if !contains(m.ToProto().GetWantlist(), str) {
17
+ t.Fail()
18
+ }
19
+}
20
+
21
+func TestNewMessageFromProto(t *testing.T) {
22
+ const str = "a_key"
23
+ protoMessage := new(PBMessage)
24
+ protoMessage.Wantlist = []string{string(str)}
25
+ if !contains(protoMessage.Wantlist, str) {
26
+ t.Fail()
27
+ }
28
+ m := newMessageFromProto(*protoMessage)
29
+ if !contains(m.ToProto().GetWantlist(), str) {
30
+ t.Fail()
31
+ }
32
+}
33
+
34
+func TestAppendBlock(t *testing.T) {
35
+
36
+ strs := make([]string, 2)
37
+ strs = append(strs, "Celeritas")
38
+ strs = append(strs, "Incendia")
39
+
40
+ m := newMessage()
41
+ for _, str := range strs {
42
+ block, err := blocks.NewBlock([]byte(str))
43
+ if err != nil {
44
+ t.Fail()
45
+ }
46
+ m.AppendBlock(block)
47
+ }
48
+
49
+ // assert strings are in proto message
50
+ for _, blockbytes := range m.ToProto().GetBlocks() {
51
+ s := bytes.NewBuffer(blockbytes).String()
52
+ if !contains(strs, s) {
53
+ t.Fail()
54
+ }
55
+ }
56
+}
57
+
58
+func TestCopyProtoByValue(t *testing.T) {
59
+ const str = "foo"
60
+ m := newMessage()
61
+ protoBeforeAppend := m.ToProto()
62
+ m.AppendWanted(u.Key(str))
63
+ if contains(protoBeforeAppend.GetWantlist(), str) {
64
+ t.Fail()
65
+ }
66
+}
67
+
68
+func contains(s []string, x string) bool {
69
+ for _, a := range s {
70
+ if a == x {
71
+ return true
72
+ }
73
+ }
74
+ return false
75
+}