added mux
Juan Batiz-Benet committed
Sep 10, 2014 at 22:01 UTC
161f8158a8ae0f62625ba9dcd5fd6bf3aaf66566
5 files changed
+496
net/mux/Makefile
new
+8
@@ -0,0 +1,8 @@
1
+
2
+all: mux.pb.go
3
+
4
+mux.pb.go: mux.proto
5
+ protoc --gogo_out=. --proto_path=../../../../../:/usr/local/opt/protobuf/include:. $<
6
+
7
+clean:
8
+ rm mux.pb.go
net/mux/mux.go
new
+167
@@ -0,0 +1,167 @@
1
+package mux
2
+
3
+import (
4
+ "errors"
5
+
6
+ msg "github.com/jbenet/go-ipfs/net/message"
7
+ u "github.com/jbenet/go-ipfs/util"
8
+
9
+ context "code.google.com/p/go.net/context"
10
+ proto "code.google.com/p/goprotobuf/proto"
11
+)
12
+
13
+// Protocol objects produce + consume raw data. They are added to the Muxer
14
+// with a ProtocolID, which is added to outgoing payloads. Muxer properly
15
+// encapsulates and decapsulates when interfacing with its Protocols. The
16
+// Protocols do not encounter their ProtocolID.
17
+type Protocol interface {
18
+ GetPipe() *msg.Pipe
19
+}
20
+
21
+// ProtocolMap maps ProtocolIDs to Protocols.
22
+type ProtocolMap map[ProtocolID]Protocol
23
+
24
+// Muxer is a simple multiplexor that reads + writes to Incoming and Outgoing
25
+// channels. It multiplexes various protocols, wrapping and unwrapping data
26
+// with a ProtocolID.
27
+type Muxer struct {
28
+ // Protocols are the multiplexed services.
29
+ Protocols ProtocolMap
30
+
31
+ // cancel is the function to stop the Muxer
32
+ cancel context.CancelFunc
33
+
34
+ *msg.Pipe
35
+}
36
+
37
+// GetPipe implements the Protocol interface
38
+func (m *Muxer) GetPipe() *msg.Pipe {
39
+ return m.Pipe
40
+}
41
+
42
+// Start kicks off the Muxer goroutines.
43
+func (m *Muxer) Start(ctx context.Context) error {
44
+ if m.cancel != nil {
45
+ return errors.New("Muxer already started.")
46
+ }
47
+
48
+ // make a cancellable context.
49
+ ctx, m.cancel = context.WithCancel(ctx)
50
+
51
+ go m.handleIncomingMessages(ctx)
52
+ for pid, proto := range m.Protocols {
53
+ go m.handleOutgoingMessages(ctx, pid, proto)
54
+ }
55
+
56
+ return nil
57
+}
58
+
59
+// Stop stops muxer activity.
60
+func (m *Muxer) Stop() {
61
+ m.cancel()
62
+ m.cancel = context.CancelFunc(nil)
63
+}
64
+
65
+// AddProtocol adds a Protocol with given ProtocolID to the Muxer.
66
+func (m *Muxer) AddProtocol(p Protocol, pid ProtocolID) error {
67
+ if _, found := m.Protocols[pid]; found {
68
+ return errors.New("Another protocol already using this ProtocolID")
69
+ }
70
+
71
+ m.Protocols[pid] = p
72
+ return nil
73
+}
74
+
75
+// handleIncoming consumes the messages on the m.Incoming channel and
76
+// routes them appropriately (to the protocols).
77
+func (m *Muxer) handleIncomingMessages(ctx context.Context) {
78
+ for {
79
+ select {
80
+ case msg := <-m.Incoming:
81
+ go m.handleIncomingMessage(ctx, msg)
82
+
83
+ case <-ctx.Done():
84
+ close(m.Incoming)
85
+ close(m.Outgoing)
86
+ return
87
+ }
88
+ }
89
+}
90
+
91
+// handleIncomingMessage routes message to the appropriate protocol.
92
+func (m *Muxer) handleIncomingMessage(ctx context.Context, m1 *msg.Message) {
93
+
94
+ data, pid, err := unwrapData(m1.Data)
95
+ if err != nil {
96
+ u.PErr("muxer de-serializing error: %v\n", err)
97
+ return
98
+ }
99
+
100
+ m2 := &msg.Message{Peer: m1.Peer, Data: data}
101
+ proto, found := m.Protocols[pid]
102
+ if !found {
103
+ u.PErr("muxer unknown protocol %v\n", pid)
104
+ return
105
+ }
106
+
107
+ select {
108
+ case proto.GetPipe().Incoming <- m2:
109
+ case <-ctx.Done():
110
+ u.PErr("%v\n", ctx.Err())
111
+ return
112
+ }
113
+}
114
+
115
+// handleOutgoingMessages consumes the messages on the proto.Outgoing channel,
116
+// wraps them and sends them out.
117
+func (m *Muxer) handleOutgoingMessages(ctx context.Context, pid ProtocolID, proto Protocol) {
118
+ for {
119
+ select {
120
+ case msg := <-proto.GetPipe().Outgoing:
121
+ go m.handleOutgoingMessage(ctx, pid, msg)
122
+
123
+ case <-ctx.Done():
124
+ return
125
+ }
126
+ }
127
+}
128
+
129
+// handleOutgoingMessage wraps out a message and sends it out the
130
+func (m *Muxer) handleOutgoingMessage(ctx context.Context, pid ProtocolID, m1 *msg.Message) {
131
+ data, err := wrapData(m1.Data, pid)
132
+ if err != nil {
133
+ u.PErr("muxer serializing error: %v\n", err)
134
+ return
135
+ }
136
+
137
+ m2 := &msg.Message{Peer: m1.Peer, Data: data}
138
+ select {
139
+ case m.GetPipe().Outgoing <- m2:
140
+ case <-ctx.Done():
141
+ return
142
+ }
143
+}
144
+
145
+func wrapData(data []byte, pid ProtocolID) ([]byte, error) {
146
+ // Marshal
147
+ pbm := new(PBProtocolMessage)
148
+ pbm.ProtocolID = &pid
149
+ pbm.Data = data
150
+ b, err := proto.Marshal(pbm)
151
+ if err != nil {
152
+ return nil, err
153
+ }
154
+
155
+ return b, nil
156
+}
157
+
158
+func unwrapData(data []byte) ([]byte, ProtocolID, error) {
159
+ // Unmarshal
160
+ pbm := new(PBProtocolMessage)
161
+ err := proto.Unmarshal(data, pbm)
162
+ if err != nil {
163
+ return nil, 0, err
164
+ }
165
+
166
+ return pbm.GetData(), pbm.GetProtocolID(), nil
167
+}
net/mux/mux.pb.go
new
+90
@@ -0,0 +1,90 @@
1
+// Code generated by protoc-gen-gogo.
2
+// source: mux.proto
3
+// DO NOT EDIT!
4
+
5
+/*
6
+Package mux is a generated protocol buffer package.
7
+
8
+It is generated from these files:
9
+ mux.proto
10
+
11
+It has these top-level messages:
12
+ PBProtocolMessage
13
+*/
14
+package mux
15
+
16
+import proto "code.google.com/p/gogoprotobuf/proto"
17
+import json "encoding/json"
18
+import math "math"
19
+
20
+// Reference proto, json, and math imports to suppress error if they are not otherwise used.
21
+var _ = proto.Marshal
22
+var _ = &json.SyntaxError{}
23
+var _ = math.Inf
24
+
25
+type ProtocolID int32
26
+
27
+const (
28
+ ProtocolID_Test ProtocolID = 0
29
+ ProtocolID_Identify ProtocolID = 1
30
+ ProtocolID_Routing ProtocolID = 2
31
+ ProtocolID_Exchange ProtocolID = 3
32
+)
33
+
34
+var ProtocolID_name = map[int32]string{
35
+ 0: "Test",
36
+ 1: "Identify",
37
+ 2: "Routing",
38
+ 3: "Exchange",
39
+}
40
+var ProtocolID_value = map[string]int32{
41
+ "Test": 0,
42
+ "Identify": 1,
43
+ "Routing": 2,
44
+ "Exchange": 3,
45
+}
46
+
47
+func (x ProtocolID) Enum() *ProtocolID {
48
+ p := new(ProtocolID)
49
+ *p = x
50
+ return p
51
+}
52
+func (x ProtocolID) String() string {
53
+ return proto.EnumName(ProtocolID_name, int32(x))
54
+}
55
+func (x *ProtocolID) UnmarshalJSON(data []byte) error {
56
+ value, err := proto.UnmarshalJSONEnum(ProtocolID_value, data, "ProtocolID")
57
+ if err != nil {
58
+ return err
59
+ }
60
+ *x = ProtocolID(value)
61
+ return nil
62
+}
63
+
64
+type PBProtocolMessage struct {
65
+ ProtocolID *ProtocolID `protobuf:"varint,1,req,enum=mux.ProtocolID" json:"ProtocolID,omitempty"`
66
+ Data []byte `protobuf:"bytes,2,req" json:"Data,omitempty"`
67
+ XXX_unrecognized []byte `json:"-"`
68
+}
69
+
70
+func (m *PBProtocolMessage) Reset() { *m = PBProtocolMessage{} }
71
+func (m *PBProtocolMessage) String() string { return proto.CompactTextString(m) }
72
+func (*PBProtocolMessage) ProtoMessage() {}
73
+
74
+func (m *PBProtocolMessage) GetProtocolID() ProtocolID {
75
+ if m != nil && m.ProtocolID != nil {
76
+ return *m.ProtocolID
77
+ }
78
+ return ProtocolID_Test
79
+}
80
+
81
+func (m *PBProtocolMessage) GetData() []byte {
82
+ if m != nil {
83
+ return m.Data
84
+ }
85
+ return nil
86
+}
87
+
88
+func init() {
89
+ proto.RegisterEnum("mux.ProtocolID", ProtocolID_name, ProtocolID_value)
90
+}
net/mux/mux.proto
new
+13
@@ -0,0 +1,13 @@
1
+package mux;
2
+
3
+enum ProtocolID {
4
+ Test = 0;
5
+ Identify = 1; // setup
6
+ Routing = 2; // dht
7
+ Exchange = 3; // bitswap
8
+}
9
+
10
+message PBProtocolMessage {
11
+ required ProtocolID ProtocolID = 1;
12
+ required bytes Data = 2;
13
+}
net/mux/mux_test.go
new
+218
@@ -0,0 +1,218 @@
1
+package mux
2
+
3
+import (
4
+ "bytes"
5
+ "fmt"
6
+ "testing"
7
+ "time"
8
+
9
+ msg "github.com/jbenet/go-ipfs/net/message"
10
+ peer "github.com/jbenet/go-ipfs/peer"
11
+ u "github.com/jbenet/go-ipfs/util"
12
+ mh "github.com/jbenet/go-multihash"
13
+
14
+ context "code.google.com/p/go.net/context"
15
+)
16
+
17
+type TestProtocol struct {
18
+ *msg.Pipe
19
+}
20
+
21
+func (t *TestProtocol) GetPipe() *msg.Pipe {
22
+ return t.Pipe
23
+}
24
+
25
+func newPeer(t *testing.T, id string) *peer.Peer {
26
+ mh, err := mh.FromHexString(id)
27
+ if err != nil {
28
+ t.Error(err)
29
+ return nil
30
+ }
31
+
32
+ return &peer.Peer{ID: peer.ID(mh)}
33
+}
34
+
35
+func testMsg(t *testing.T, m *msg.Message, data []byte) {
36
+ if !bytes.Equal(data, m.Data) {
37
+ t.Errorf("Data does not match: %v != %v", data, m.Data)
38
+ }
39
+}
40
+
41
+func testWrappedMsg(t *testing.T, m *msg.Message, pid ProtocolID, data []byte) {
42
+ data2, pid2, err := unwrapData(m.Data)
43
+ if err != nil {
44
+ t.Error(err)
45
+ }
46
+
47
+ if pid != pid2 {
48
+ t.Errorf("ProtocolIDs do not match: %v != %v", pid, pid2)
49
+ }
50
+
51
+ if !bytes.Equal(data, data2) {
52
+ t.Errorf("Data does not match: %v != %v", data, data2)
53
+ }
54
+}
55
+
56
+func TestSimpleMuxer(t *testing.T) {
57
+
58
+ // setup
59
+ p1 := &TestProtocol{Pipe: msg.NewPipe(10)}
60
+ p2 := &TestProtocol{Pipe: msg.NewPipe(10)}
61
+ pid1 := ProtocolID_Test
62
+ pid2 := ProtocolID_Routing
63
+ mux1 := &Muxer{
64
+ Pipe: msg.NewPipe(10),
65
+ Protocols: ProtocolMap{
66
+ pid1: p1,
67
+ pid2: p2,
68
+ },
69
+ }
70
+ peer1 := newPeer(t, "11140beec7b5ea3f0fdbc95d0dd47f3c5bc275aaaaaa")
71
+ // peer2 := newPeer(t, "11140beec7b5ea3f0fdbc95d0dd47f3c5bc275bbbbbb")
72
+
73
+ // run muxer
74
+ ctx := context.Background()
75
+ mux1.Start(ctx)
76
+
77
+ // test outgoing p1
78
+ for _, s := range []string{"foo", "bar", "baz"} {
79
+ p1.Outgoing <- &msg.Message{Peer: peer1, Data: []byte(s)}
80
+ testWrappedMsg(t, <-mux1.Outgoing, pid1, []byte(s))
81
+ }
82
+
83
+ // test incoming p1
84
+ for _, s := range []string{"foo", "bar", "baz"} {
85
+ d, err := wrapData([]byte(s), pid1)
86
+ if err != nil {
87
+ t.Error(err)
88
+ }
89
+ mux1.Incoming <- &msg.Message{Peer: peer1, Data: d}
90
+ testMsg(t, <-p1.Incoming, []byte(s))
91
+ }
92
+
93
+ // test outgoing p2
94
+ for _, s := range []string{"foo", "bar", "baz"} {
95
+ p2.Outgoing <- &msg.Message{Peer: peer1, Data: []byte(s)}
96
+ testWrappedMsg(t, <-mux1.Outgoing, pid2, []byte(s))
97
+ }
98
+
99
+ // test incoming p2
100
+ for _, s := range []string{"foo", "bar", "baz"} {
101
+ d, err := wrapData([]byte(s), pid2)
102
+ if err != nil {
103
+ t.Error(err)
104
+ }
105
+ mux1.Incoming <- &msg.Message{Peer: peer1, Data: d}
106
+ testMsg(t, <-p2.Incoming, []byte(s))
107
+ }
108
+}
109
+
110
+func TestSimultMuxer(t *testing.T) {
111
+
112
+ // setup
113
+ p1 := &TestProtocol{Pipe: msg.NewPipe(10)}
114
+ p2 := &TestProtocol{Pipe: msg.NewPipe(10)}
115
+ pid1 := ProtocolID_Test
116
+ pid2 := ProtocolID_Identify
117
+ mux1 := &Muxer{
118
+ Pipe: msg.NewPipe(10),
119
+ Protocols: ProtocolMap{
120
+ pid1: p1,
121
+ pid2: p2,
122
+ },
123
+ }
124
+ peer1 := newPeer(t, "11140beec7b5ea3f0fdbc95d0dd47f3c5bc275aaaaaa")
125
+ // peer2 := newPeer(t, "11140beec7b5ea3f0fdbc95d0dd47f3c5bc275bbbbbb")
126
+
127
+ // run muxer
128
+ ctx, cancel := context.WithCancel(context.Background())
129
+ mux1.Start(ctx)
130
+
131
+ // counts
132
+ total := 10000
133
+ speed := time.Microsecond * 1
134
+ counts := [2][2][2]int{}
135
+
136
+ // run producers at every end sending incrementing messages
137
+ produceOut := func(pid ProtocolID, size int) {
138
+ limiter := time.Tick(speed)
139
+ for i := 0; i < size; i++ {
140
+ <-limiter
141
+ s := fmt.Sprintf("proto %v out %v", pid, i)
142
+ m := &msg.Message{Peer: peer1, Data: []byte(s)}
143
+ mux1.Protocols[pid].GetPipe().Outgoing <- m
144
+ counts[pid][0][0]++
145
+ u.DOut("sent %v\n", s)
146
+ }
147
+ }
148
+
149
+ produceIn := func(pid ProtocolID, size int) {
150
+ limiter := time.Tick(speed)
151
+ for i := 0; i < size; i++ {
152
+ <-limiter
153
+ s := fmt.Sprintf("proto %v in %v", pid, i)
154
+ d, err := wrapData([]byte(s), pid)
155
+ if err != nil {
156
+ t.Error(err)
157
+ }
158
+
159
+ m := &msg.Message{Peer: peer1, Data: d}
160
+ mux1.Incoming <- m
161
+ counts[pid][1][0]++
162
+ u.DOut("sent %v\n", s)
163
+ }
164
+ }
165
+
166
+ consumeOut := func() {
167
+ for {
168
+ select {
169
+ case m := <-mux1.Outgoing:
170
+ data, pid, err := unwrapData(m.Data)
171
+ if err != nil {
172
+ t.Error(err)
173
+ }
174
+
175
+ u.DOut("got %v\n", string(data))
176
+ counts[pid][1][1]++
177
+
178
+ case <-ctx.Done():
179
+ return
180
+ }
181
+ }
182
+ }
183
+
184
+ consumeIn := func(pid ProtocolID) {
185
+ for {
186
+ select {
187
+ case m := <-mux1.Protocols[pid].GetPipe().Incoming:
188
+ counts[pid][0][1]++
189
+ u.DOut("got %v\n", string(m.Data))
190
+ case <-ctx.Done():
191
+ return
192
+ }
193
+ }
194
+ }
195
+
196
+ go produceOut(pid1, total)
197
+ go produceOut(pid2, total)
198
+ go produceIn(pid1, total)
199
+ go produceIn(pid2, total)
200
+ go consumeOut()
201
+ go consumeIn(pid1)
202
+ go consumeIn(pid2)
203
+
204
+ limiter := time.Tick(speed)
205
+ for {
206
+ <-limiter
207
+ got := counts[0][0][0] + counts[0][0][1] +
208
+ counts[0][1][0] + counts[0][1][1] +
209
+ counts[1][0][0] + counts[1][0][1] +
210
+ counts[1][1][0] + counts[1][1][1]
211
+
212
+ if got == total*8 {
213
+ cancel()
214
+ return
215
+ }
216
+ }
217
+
218
+}