add transport logic to mocknet
License: MIT Signed-off-by: Karthik Bala <karthikbala444@gmail.com>
Karthik Bala committed
Jul 6, 2015 at 15:10 UTC
0e597b3ae6e2895c1ba92269cfeb58b657dbe29f
7 files changed
+312
-15
exchange/bitswap/testutils.go
+2
-2
@@ -46,7 +46,7 @@ func (g *SessionGenerator) Next() Instance {
46
if err != nil {
47
panic("FIXME") // TODO change signature
48
}
49
- return session(g.ctx, g.net, p)
49
+ return Session(g.ctx, g.net, p)
50
}
51
52
func (g *SessionGenerator) Instances(n int) []Instance {
@@ -85,7 +85,7 @@ func (i *Instance) SetBlockstoreLatency(t time.Duration) time.Duration {
85
// NB: It's easy make mistakes by providing the same peer ID to two different
86
// sessions. To safeguard, use the SessionGenerator to generate sessions. It's
87
// just a much better idea.
88
-func session(ctx context.Context, net tn.Network, p testutil.Identity) Instance {
88
+func Session(ctx context.Context, net tn.Network, p testutil.Identity) Instance {
89
bsdelay := delay.Fixed(0)
90
const writeCacheElems = 100
91
p2p/net/mock/interface.go
+4
-4
@@ -7,13 +7,12 @@
7
package mocknet
8
9
import (
10
- "io"
11
- "time"
12
-
10
ic "github.com/ipfs/go-ipfs/p2p/crypto"
11
host "github.com/ipfs/go-ipfs/p2p/host"
12
inet "github.com/ipfs/go-ipfs/p2p/net"
13
peer "github.com/ipfs/go-ipfs/p2p/peer"
14
+ "io"
15
+ "time"
16
17
ma "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr"
18
)
@@ -59,13 +58,14 @@ type Mocknet interface {
58
ConnectNets(inet.Network, inet.Network) (inet.Conn, error)
59
DisconnectPeers(peer.ID, peer.ID) error
60
DisconnectNets(inet.Network, inet.Network) error
61
+ LinkAll() error
62
}
63
64
// LinkOptions are used to change aspects of the links.
65
// Sorry but they dont work yet :(
66
type LinkOptions struct {
67
Latency time.Duration
68
- Bandwidth int // in bytes-per-second
68
+ Bandwidth float64 // in bytes-per-second
69
// we can make these values distributions down the road.
70
}
71
p2p/net/mock/mock_link.go
+21
-7
@@ -1,8 +1,10 @@
1
package mocknet
2
3
import (
4
+ // "fmt"
5
"io"
6
"sync"
7
+ "time"
8
9
inet "github.com/ipfs/go-ipfs/p2p/net"
10
peer "github.com/ipfs/go-ipfs/p2p/peer"
@@ -11,17 +13,20 @@ import (
13
// link implements mocknet.Link
14
// and, for simplicity, inet.Conn
15
type link struct {
14
- mock *mocknet
15
- nets []*peernet
16
- opts LinkOptions
17
-
16
+ mock *mocknet
17
+ nets []*peernet
18
+ opts LinkOptions
19
+ ratelimiter *ratelimiter
20
// this could have addresses on both sides.
21
22
sync.RWMutex
23
}
24
25
func newLink(mn *mocknet, opts LinkOptions) *link {
24
- return &link{mock: mn, opts: opts}
26
+ l := &link{mock: mn,
27
+ opts: opts,
28
+ ratelimiter: NewRatelimiter(opts.Bandwidth)}
29
+ return l
30
}
31
32
func (l *link) newConnPair(dialer *peernet) (*conn, *conn) {
@@ -57,8 +62,8 @@ func (l *link) newStreamPair() (*stream, *stream) {
62
r1, w1 := io.Pipe()
63
r2, w2 := io.Pipe()
64
60
- s1 := &stream{Reader: r1, Writer: w2}
61
- s2 := &stream{Reader: r2, Writer: w1}
65
+ s1 := NewStream(w2, r1)
66
+ s2 := NewStream(w1, r2)
67
return s1, s2
68
}
69
@@ -86,8 +91,17 @@ func (l *link) Peers() []peer.ID {
91
92
func (l *link) SetOptions(o LinkOptions) {
93
l.opts = o
94
+ l.ratelimiter.UpdateBandwidth(l.opts.Bandwidth)
95
}
96
97
func (l *link) Options() LinkOptions {
98
return l.opts
99
}
100
+
101
+func (l *link) GetLatency() time.Duration {
102
+ return l.opts.Latency
103
+}
104
+
105
+func (l *link) RateLimit(dataSize int) time.Duration {
106
+ return l.ratelimiter.Limit(dataSize)
107
+}
p2p/net/mock/mock_notif_test.go
+1
-1
@@ -63,7 +63,7 @@ func TestNotifications(t *testing.T) {
63
}
64
}
65
if !found {
66
- t.Error("connection not found")
66
+ t.Error("connection not found", c1, len(expect), len(actual))
67
}
68
}
69
p2p/net/mock/mock_stream.go
+114
-1
@@ -1,7 +1,11 @@
1
package mocknet
2
3
import (
4
+ "bytes"
5
"io"
6
+ "time"
7
+
8
+ process "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/goprocess"
9
10
inet "github.com/ipfs/go-ipfs/p2p/net"
11
)
@@ -10,10 +14,51 @@ import (
14
type stream struct {
15
io.Reader
16
io.Writer
13
- conn *conn
17
+ conn *conn
18
+ toDeliver chan *transportObject
19
+ proc process.Process
20
+}
21
+
22
+type transportObject struct {
23
+ msg []byte
24
+ arrivalTime time.Time
25
+}
26
+
27
+func NewStream(w io.Writer, r io.Reader) *stream {
28
+ s := &stream{
29
+ Reader: r,
30
+ Writer: w,
31
+ toDeliver: make(chan *transportObject),
32
+ }
33
+
34
+ s.proc = process.WithTeardown(s.teardown)
35
+ s.proc.Go(s.transport)
36
+ return s
37
+}
38
+
39
+// How to handle errors with writes?
40
+func (s *stream) Write(p []byte) (n int, err error) {
41
+ l := s.conn.link
42
+ delay := l.GetLatency() + l.RateLimit(len(p))
43
+ t := time.Now().Add(delay)
44
+ select {
45
+ case <-s.proc.Closing(): // bail out if we're closing.
46
+ return 0, io.ErrClosedPipe
47
+ case s.toDeliver <- &transportObject{msg: p, arrivalTime: t}:
48
+ }
49
+ return len(p), nil
50
}
51
52
func (s *stream) Close() error {
53
+ return s.proc.Close()
54
+}
55
+
56
+// teardown shuts down the stream. it is called by s.proc.Close()
57
+// after all the children of this s.proc (i.e. transport's proc)
58
+// are done.
59
+func (s *stream) teardown() error {
60
+ // at this point, no streams are writing.
61
+
62
s.conn.removeStream(s)
63
if r, ok := (s.Reader).(io.Closer); ok {
64
r.Close()
@@ -30,3 +75,71 @@ func (s *stream) Close() error {
75
func (s *stream) Conn() inet.Conn {
76
return s.conn
77
}
78
+
79
+// transport will grab message arrival times, wait until that time, and
80
+// then write the message out when it is scheduled to arrive
81
+func (s *stream) transport(proc process.Process) {
82
+ bufsize := 256
83
+ buf := new(bytes.Buffer)
84
+ ticker := time.NewTicker(time.Millisecond * 4)
85
+
86
+ // writeBuf writes the contents of buf through to the s.Writer.
87
+ // done only when arrival time makes sense.
88
+ drainBuf := func() {
89
+ if buf.Len() > 0 {
90
+ _, err := s.Writer.Write(buf.Bytes())
91
+ if err != nil {
92
+ return
93
+ }
94
+ buf.Reset()
95
+ }
96
+ }
97
+
98
+ // deliverOrWait is a helper func that processes
99
+ // an incoming packet. it waits until the arrival time,
100
+ // and then writes things out.
101
+ deliverOrWait := func(o *transportObject) {
102
+ buffered := len(o.msg) + buf.Len()
103
+
104
+ now := time.Now()
105
+ if now.Before(o.arrivalTime) {
106
+ if buffered < bufsize {
107
+ buf.Write(o.msg)
108
+ return
109
+ }
110
+
111
+ // we do not buffer + return here, instead hanging the
112
+ // call (i.e. not accepting any more transportObjects)
113
+ // so that we apply back-pressure to the sender.
114
+ // this sleep should wake up same time as ticker.
115
+ time.Sleep(o.arrivalTime.Sub(now))
116
+ }
117
+
118
+ // ok, we waited our due time. now rite the buf + msg.
119
+
120
+ // drainBuf first, before we write this message.
121
+ drainBuf()
122
+
123
+ // write this message.
124
+ _, err := s.Writer.Write(o.msg)
125
+ if err != nil {
126
+ log.Error("mock_stream", err)
127
+ }
128
+ }
129
+
130
+ for {
131
+ select {
132
+ case <-proc.Closing():
133
+ return // bail out of here.
134
+
135
+ case o, ok := <-s.toDeliver:
136
+ if !ok {
137
+ return
138
+ }
139
+ deliverOrWait(o)
140
+
141
+ case <-ticker.C: // ok, due to write it out.
142
+ drainBuf()
143
+ }
144
+ }
145
+}
p2p/net/mock/mock_test.go
+101
@@ -3,9 +3,11 @@ package mocknet
3
import (
4
"bytes"
5
"io"
6
+ "math"
7
"math/rand"
8
"sync"
9
"testing"
10
+ "time"
11
12
inet "github.com/ipfs/go-ipfs/p2p/net"
13
peer "github.com/ipfs/go-ipfs/p2p/peer"
@@ -478,3 +480,102 @@ func TestAdding(t *testing.T) {
480
}
481
482
}
483
+
484
+func TestRateLimiting(t *testing.T) {
485
+ rl := NewRatelimiter(10)
486
+
487
+ if !within(rl.Limit(10), time.Duration(float32(time.Second)), time.Millisecond/10) {
488
+ t.Fail()
489
+ }
490
+ if !within(rl.Limit(10), time.Duration(float32(time.Second*2)), time.Millisecond) {
491
+ t.Fail()
492
+ }
493
+ if !within(rl.Limit(10), time.Duration(float32(time.Second*3)), time.Millisecond) {
494
+ t.Fail()
495
+ }
496
+
497
+ if within(rl.Limit(10), time.Duration(float32(time.Second*3)), time.Millisecond) {
498
+ t.Fail()
499
+ }
500
+
501
+ rl.UpdateBandwidth(50)
502
+ if !within(rl.Limit(75), time.Duration(float32(time.Second)*1.5), time.Millisecond/10) {
503
+ t.Fail()
504
+ }
505
+
506
+ if within(rl.Limit(75), time.Duration(float32(time.Second)*1.5), time.Millisecond/10) {
507
+ t.Fail()
508
+ }
509
+
510
+ rl.UpdateBandwidth(100)
511
+ if !within(rl.Limit(1), time.Duration(time.Millisecond*10), time.Millisecond/10) {
512
+ t.Fail()
513
+ }
514
+
515
+ if within(rl.Limit(1), time.Duration(time.Millisecond*10), time.Millisecond/10) {
516
+ t.Fail()
517
+ }
518
+}
519
+
520
+func within(t1 time.Duration, t2 time.Duration, tolerance time.Duration) bool {
521
+ return math.Abs(float64(t1)-float64(t2)) < float64(tolerance)
522
+}
523
+
524
+func TestLimitedStreams(t *testing.T) {
525
+ mn, err := FullMeshConnected(context.Background(), 2)
526
+ if err != nil {
527
+ t.Fatal(err)
528
+ }
529
+
530
+ var wg sync.WaitGroup
531
+ messages := 4
532
+ messageSize := 500
533
+ handler := func(s inet.Stream) {
534
+ b := make([]byte, messageSize)
535
+ for i := 0; i < messages; i++ {
536
+ if _, err := io.ReadFull(s, b); err != nil {
537
+ log.Fatal(err)
538
+ }
539
+ if !bytes.Equal(b[:4], []byte("ping")) {
540
+ log.Fatal("bytes mismatch")
541
+ }
542
+ wg.Done()
543
+ }
544
+ s.Close()
545
+ }
546
+
547
+ hosts := mn.Hosts()
548
+ for _, h := range mn.Hosts() {
549
+ h.SetStreamHandler(protocol.TestingID, handler)
550
+ }
551
+
552
+ peers := mn.Peers()
553
+ links := mn.LinksBetweenPeers(peers[0], peers[1])
554
+ // 1000 byte per second bandwidth
555
+ bps := float64(1000)
556
+ opts := links[0].Options()
557
+ opts.Bandwidth = bps
558
+ for _, link := range links {
559
+ link.SetOptions(opts)
560
+ }
561
+
562
+ s, err := hosts[0].NewStream(protocol.TestingID, hosts[1].ID())
563
+ if err != nil {
564
+ t.Fatal(err)
565
+ }
566
+
567
+ filler := make([]byte, messageSize-4)
568
+ data := append([]byte("ping"), filler...)
569
+ before := time.Now()
570
+ for i := 0; i < messages; i++ {
571
+ wg.Add(1)
572
+ if _, err := s.Write(data); err != nil {
573
+ panic(err)
574
+ }
575
+ }
576
+
577
+ wg.Wait()
578
+ if !within(time.Since(before), time.Duration(time.Second*2), time.Second/3) {
579
+ t.Fatal("Expected 2ish seconds but got ", time.Since(before))
580
+ }
581
+}
p2p/net/mock/ratelimiter.go
new
+69
@@ -0,0 +1,69 @@
1
+package mocknet
2
+
3
+import (
4
+ "time"
5
+)
6
+
7
+// A ratelimiter is used by a link to determine how long to wait before sending
8
+// data given a bandwidth cap.
9
+type ratelimiter struct {
10
+ bandwidth float64 // bytes per nanosecond
11
+ allowance float64 // in bytes
12
+ maxAllowance float64 // in bytes
13
+ lastUpdate time.Time // when allowance was updated last
14
+ count int // number of times rate limiting was applied
15
+ duration time.Duration // total delay introduced due to rate limiting
16
+}
17
+
18
+// Creates a new ratelimiter with bandwidth (in bytes/sec)
19
+func NewRatelimiter(bandwidth float64) *ratelimiter {
20
+ // convert bandwidth to bytes per nanosecond
21
+ b := bandwidth / float64(time.Second)
22
+ return &ratelimiter{
23
+ bandwidth: b,
24
+ allowance: 0,
25
+ maxAllowance: bandwidth,
26
+ lastUpdate: time.Now(),
27
+ }
28
+}
29
+
30
+// Changes bandwidth of a ratelimiter and resets its allowance
31
+func (r *ratelimiter) UpdateBandwidth(bandwidth float64) {
32
+ // Convert bandwidth from bytes/second to bytes/nanosecond
33
+ b := bandwidth / float64(time.Second)
34
+ r.bandwidth = b
35
+ // Reset allowance
36
+ r.allowance = 0
37
+ r.maxAllowance = bandwidth
38
+ r.lastUpdate = time.Now()
39
+}
40
+
41
+// Returns how long to wait before sending data with length 'dataSize' bytes
42
+func (r *ratelimiter) Limit(dataSize int) time.Duration {
43
+ // update time
44
+ var duration time.Duration = time.Duration(0)
45
+ if r.bandwidth == 0 {
46
+ return duration
47
+ }
48
+ current := time.Now()
49
+ elapsedTime := current.Sub(r.lastUpdate)
50
+ r.lastUpdate = current
51
+
52
+ allowance := r.allowance + float64(elapsedTime)*r.bandwidth
53
+ // allowance can't exceed bandwidth
54
+ if allowance > r.maxAllowance {
55
+ allowance = r.maxAllowance
56
+ }
57
+
58
+ allowance -= float64(dataSize)
59
+ if allowance < 0 {
60
+ // sleep until allowance is back to 0
61
+ duration = time.Duration(-allowance / r.bandwidth)
62
+ // rate limiting was applied, record stats
63
+ r.count++
64
+ r.duration += duration
65
+ }
66
+
67
+ r.allowance = allowance
68
+ return duration
69
+}