updated msgio (varints)
Juan Batiz-Benet committed
Dec 29, 2014 at 07:40 UTC
4309d15edb29b0f524a0d58bc1335fb73c66116d
4 files changed
+269
-9
Godeps/Godeps.json
+1
-1
@@ -118,7 +118,7 @@
118
},
119
{
120
"ImportPath": "github.com/jbenet/go-msgio",
121
- "Rev": "5e7289d3a0cd046a5bee30b187cc844c31f54dce"
121
+ "Rev": "dbae89193876910c736b2ce1291fa8bbcf299d77"
122
},
123
{
124
"ImportPath": "github.com/jbenet/go-multiaddr",
Godeps/_workspace/src/github.com/jbenet/go-msgio/msgio_test.go
+20
-8
@@ -15,6 +15,24 @@ func TestReadWrite(t *testing.T) {
15
buf := bytes.NewBuffer(nil)
16
writer := NewWriter(buf)
17
reader := NewReader(buf)
18
+ SubtestReadWrite(t, writer, reader)
19
+}
20
+
21
+func TestReadWriteMsg(t *testing.T) {
22
+ buf := bytes.NewBuffer(nil)
23
+ writer := NewWriter(buf)
24
+ reader := NewReader(buf)
25
+ SubtestReadWriteMsg(t, writer, reader)
26
+}
27
+
28
+func TestReadWriteMsgSync(t *testing.T) {
29
+ buf := bytes.NewBuffer(nil)
30
+ writer := NewWriter(buf)
31
+ reader := NewReader(buf)
32
+ SubtestReadWriteMsgSync(t, writer, reader)
33
+}
34
+
35
+func SubtestReadWrite(t *testing.T, writer WriteCloser, reader ReadCloser) {
36
msgs := [1000][]byte{}
37
38
r := rand.New(rand.NewSource(time.Now().UnixNano()))
@@ -58,10 +76,7 @@ func TestReadWrite(t *testing.T) {
76
}
77
}
78
61
-func TestReadWriteMsg(t *testing.T) {
62
- buf := bytes.NewBuffer(nil)
63
- writer := NewWriter(buf)
64
- reader := NewReader(buf)
79
+func SubtestReadWriteMsg(t *testing.T, writer WriteCloser, reader ReadCloser) {
80
msgs := [1000][]byte{}
81
82
r := rand.New(rand.NewSource(time.Now().UnixNano()))
@@ -100,10 +115,7 @@ func TestReadWriteMsg(t *testing.T) {
115
}
116
}
117
103
-func TestReadWriteMsgSync(t *testing.T) {
104
- buf := bytes.NewBuffer(nil)
105
- writer := NewWriter(buf)
106
- reader := NewReader(buf)
118
+func SubtestReadWriteMsgSync(t *testing.T, writer WriteCloser, reader ReadCloser) {
119
msgs := [1000][]byte{}
120
121
r := rand.New(rand.NewSource(time.Now().UnixNano()))
Godeps/_workspace/src/github.com/jbenet/go-msgio/varint.go
new
+182
@@ -0,0 +1,182 @@
1
+package msgio
2
+
3
+import (
4
+ "encoding/binary"
5
+ "io"
6
+ "sync"
7
+
8
+ mpool "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-msgio/mpool"
9
+)
10
+
11
+// varintWriter is the underlying type that implements the Writer interface.
12
+type varintWriter struct {
13
+ W io.Writer
14
+
15
+ lbuf []byte // for encoding varints
16
+ lock sync.Locker // for threadsafe writes
17
+}
18
+
19
+// NewVarintWriter wraps an io.Writer with a varint msgio framed writer.
20
+// The msgio.Writer will write the length prefix of every message written
21
+// as a varint, using https://golang.org/pkg/encoding/binary/#PutUvarint
22
+func NewVarintWriter(w io.Writer) WriteCloser {
23
+ return &varintWriter{
24
+ W: w,
25
+ lbuf: make([]byte, binary.MaxVarintLen64),
26
+ lock: new(sync.Mutex),
27
+ }
28
+}
29
+
30
+func (s *varintWriter) Write(msg []byte) (int, error) {
31
+ err := s.WriteMsg(msg)
32
+ if err != nil {
33
+ return 0, err
34
+ }
35
+ return len(msg), nil
36
+}
37
+
38
+func (s *varintWriter) WriteMsg(msg []byte) error {
39
+ s.lock.Lock()
40
+ defer s.lock.Unlock()
41
+
42
+ length := uint64(len(msg))
43
+ n := binary.PutUvarint(s.lbuf, length)
44
+ if _, err := s.W.Write(s.lbuf[:n]); err != nil {
45
+ return err
46
+ }
47
+ _, err := s.W.Write(msg)
48
+ return err
49
+}
50
+
51
+func (s *varintWriter) Close() error {
52
+ s.lock.Lock()
53
+ defer s.lock.Unlock()
54
+
55
+ if c, ok := s.W.(io.Closer); ok {
56
+ return c.Close()
57
+ }
58
+ return nil
59
+}
60
+
61
+// varintReader is the underlying type that implements the Reader interface.
62
+type varintReader struct {
63
+ R io.Reader
64
+ br io.ByteReader // for reading varints.
65
+
66
+ lbuf []byte
67
+ next int
68
+ pool *mpool.Pool
69
+ lock sync.Locker
70
+}
71
+
72
+// NewVarintReader wraps an io.Reader with a varint msgio framed reader.
73
+// The msgio.Reader will read whole messages at a time (using the length).
74
+// Varints read according to https://golang.org/pkg/encoding/binary/#ReadUvarint
75
+// Assumes an equivalent writer on the other side.
76
+func NewVarintReader(r io.Reader) ReadCloser {
77
+ return NewVarintReaderWithPool(r, &mpool.ByteSlicePool)
78
+}
79
+
80
+// NewVarintReaderWithPool wraps an io.Reader with a varint msgio framed reader.
81
+// The msgio.Reader will read whole messages at a time (using the length).
82
+// Varints read according to https://golang.org/pkg/encoding/binary/#ReadUvarint
83
+// Assumes an equivalent writer on the other side. It uses a given mpool.Pool
84
+func NewVarintReaderWithPool(r io.Reader, p *mpool.Pool) ReadCloser {
85
+ if p == nil {
86
+ panic("nil pool")
87
+ }
88
+ return &varintReader{
89
+ R: r,
90
+ br: &simpleByteReader{R: r},
91
+ lbuf: make([]byte, binary.MaxVarintLen64),
92
+ next: -1,
93
+ pool: p,
94
+ lock: new(sync.Mutex),
95
+ }
96
+}
97
+
98
+// NextMsgLen reads the length of the next msg into s.lbuf, and returns it.
99
+// WARNING: like Read, NextMsgLen is destructive. It reads from the internal
100
+// reader.
101
+func (s *varintReader) NextMsgLen() (int, error) {
102
+ s.lock.Lock()
103
+ defer s.lock.Unlock()
104
+ return s.nextMsgLen()
105
+}
106
+
107
+func (s *varintReader) nextMsgLen() (int, error) {
108
+ if s.next == -1 {
109
+ length, err := binary.ReadUvarint(s.br)
110
+ if err != nil {
111
+ return 0, err
112
+ }
113
+ s.next = int(length)
114
+ }
115
+ return s.next, nil
116
+}
117
+
118
+func (s *varintReader) Read(msg []byte) (int, error) {
119
+ s.lock.Lock()
120
+ defer s.lock.Unlock()
121
+
122
+ length, err := s.nextMsgLen()
123
+ if err != nil {
124
+ return 0, err
125
+ }
126
+
127
+ if length > len(msg) {
128
+ return 0, io.ErrShortBuffer
129
+ }
130
+ _, err = io.ReadFull(s.R, msg[:length])
131
+ s.next = -1 // signal we've consumed this msg
132
+ return length, err
133
+}
134
+
135
+func (s *varintReader) ReadMsg() ([]byte, error) {
136
+ s.lock.Lock()
137
+ defer s.lock.Unlock()
138
+
139
+ length, err := s.nextMsgLen()
140
+ if err != nil {
141
+ return nil, err
142
+ }
143
+
144
+ msgb := s.pool.Get(uint32(length))
145
+ if msgb == nil {
146
+ return nil, io.ErrShortBuffer
147
+ }
148
+ msg := msgb.([]byte)[:length]
149
+ _, err = io.ReadFull(s.R, msg)
150
+ s.next = -1 // signal we've consumed this msg
151
+ return msg, err
152
+}
153
+
154
+func (s *varintReader) ReleaseMsg(msg []byte) {
155
+ s.pool.Put(uint32(cap(msg)), msg)
156
+}
157
+
158
+func (s *varintReader) Close() error {
159
+ s.lock.Lock()
160
+ defer s.lock.Unlock()
161
+
162
+ if c, ok := s.R.(io.Closer); ok {
163
+ return c.Close()
164
+ }
165
+ return nil
166
+}
167
+
168
+type simpleByteReader struct {
169
+ R io.Reader
170
+ buf []byte
171
+}
172
+
173
+func (r *simpleByteReader) ReadByte() (c byte, err error) {
174
+ if r.buf == nil {
175
+ r.buf = make([]byte, 1)
176
+ }
177
+
178
+ if _, err := io.ReadFull(r.R, r.buf); err != nil {
179
+ return 0, err
180
+ }
181
+ return r.buf[0], nil
182
+}
Godeps/_workspace/src/github.com/jbenet/go-msgio/varint_test.go
new
+66
@@ -0,0 +1,66 @@
1
+package msgio
2
+
3
+import (
4
+ "bytes"
5
+ "encoding/binary"
6
+ "testing"
7
+)
8
+
9
+func TestVarintReadWrite(t *testing.T) {
10
+ buf := bytes.NewBuffer(nil)
11
+ writer := NewVarintWriter(buf)
12
+ reader := NewVarintReader(buf)
13
+ SubtestReadWrite(t, writer, reader)
14
+}
15
+
16
+func TestVarintReadWriteMsg(t *testing.T) {
17
+ buf := bytes.NewBuffer(nil)
18
+ writer := NewVarintWriter(buf)
19
+ reader := NewVarintReader(buf)
20
+ SubtestReadWriteMsg(t, writer, reader)
21
+}
22
+
23
+func TestVarintReadWriteMsgSync(t *testing.T) {
24
+ buf := bytes.NewBuffer(nil)
25
+ writer := NewVarintWriter(buf)
26
+ reader := NewVarintReader(buf)
27
+ SubtestReadWriteMsgSync(t, writer, reader)
28
+}
29
+
30
+func TestVarintWrite(t *testing.T) {
31
+ SubtestVarintWrite(t, []byte("hello world"))
32
+ SubtestVarintWrite(t, []byte("hello world hello world hello world"))
33
+ SubtestVarintWrite(t, make([]byte, 1<<20))
34
+ SubtestVarintWrite(t, []byte(""))
35
+}
36
+
37
+func SubtestVarintWrite(t *testing.T, msg []byte) {
38
+ buf := bytes.NewBuffer(nil)
39
+ writer := NewVarintWriter(buf)
40
+
41
+ if err := writer.WriteMsg(msg); err != nil {
42
+ t.Fatal(err)
43
+ }
44
+
45
+ bb := buf.Bytes()
46
+
47
+ sbr := simpleByteReader{R: buf}
48
+ length, err := binary.ReadUvarint(&sbr)
49
+ if err != nil {
50
+ t.Fatal(err)
51
+ }
52
+
53
+ t.Logf("checking varint is %d", len(msg))
54
+ if int(length) != len(msg) {
55
+ t.Fatalf("incorrect varint: %d != %d", length, len(msg))
56
+ }
57
+
58
+ lbuf := make([]byte, binary.MaxVarintLen64)
59
+ n := binary.PutUvarint(lbuf, length)
60
+
61
+ bblen := int(length) + n
62
+ t.Logf("checking wrote (%d + %d) bytes", length, n)
63
+ if len(bb) != bblen {
64
+ t.Fatalf("wrote incorrect number of bytes: %d != %d", len(bb), bblen)
65
+ }
66
+}