util: cleaner ByteChanReader.Read
Juan Batiz-Benet committed
Oct 30, 2014 at 04:14 UTC
adec8dac127da62d471bf46a7715d07a3c6936ee
2 files changed
+31
-43
util/util.go
+18
-30
@@ -60,40 +60,28 @@ func NewByteChanReader(in chan []byte) io.Reader {
60
return &byteChanReader{in: in}
61
}
62
63
-func (bcr *byteChanReader) Read(b []byte) (int, error) {
64
- if len(bcr.buf) == 0 {
65
- data, ok := <-bcr.in
66
- if !ok {
67
- return 0, io.EOF
63
+func (bcr *byteChanReader) Read(output []byte) (int, error) {
64
+ remain := output
65
+ remainLen := len(output)
66
+ outputLen := 0
67
+ more := false
68
+ next := bcr.buf
69
+
70
+ for {
71
+ n := copy(remain, next)
72
+ remainLen -= n
73
+ outputLen += n
74
+ if remainLen == 0 {
75
+ bcr.buf = next[n:]
76
+ return outputLen, nil
77
}
69
- bcr.buf = data
70
- }
71
-
72
- if len(bcr.buf) >= len(b) {
73
- copy(b, bcr.buf)
74
- bcr.buf = bcr.buf[len(b):]
75
- return len(b), nil
76
- }
78
78
- copy(b, bcr.buf)
79
- b = b[len(bcr.buf):]
80
- totread := len(bcr.buf)
81
-
82
- for data := range bcr.in {
83
- if len(data) > len(b) {
84
- totread += len(b)
85
- copy(b, data[:len(b)])
86
- bcr.buf = data[len(b):]
87
- return totread, nil
88
- }
89
- copy(b, data)
90
- totread += len(data)
91
- b = b[len(data):]
92
- if len(b) == 0 {
93
- return totread, nil
79
+ remain = remain[n:]
80
+ next, more = <-bcr.in
81
+ if !more {
82
+ return outputLen, io.EOF
83
}
84
}
96
- return totread, io.EOF
85
}
86
87
type randGen struct {
util/util_test.go
+13
-13
@@ -2,6 +2,7 @@ package util
2
3
import (
4
"bytes"
5
+ "fmt"
6
"io/ioutil"
7
"math/rand"
8
"testing"
@@ -30,22 +31,20 @@ func TestKey(t *testing.T) {
31
}
32
33
func TestByteChanReader(t *testing.T) {
33
- data := make([]byte, 1024*1024)
34
- r := NewTimeSeededRand()
35
- r.Read(data)
34
+
35
+ var data bytes.Buffer
36
dch := make(chan []byte, 8)
37
+ randr := NewTimeSeededRand()
38
39
go func() {
39
- beg := 0
40
- for i := 0; i < len(data); {
41
- i += rand.Intn(100) + 1
42
- if i > len(data) {
43
- i = len(data)
44
- }
45
- dch <- data[beg:i]
46
- beg = i
40
+ defer close(dch)
41
+ for i := 0; i < rand.Intn(100)+100; i++ {
42
+ chunk := make([]byte, rand.Intn(100000)+10)
43
+ randr.Read(chunk)
44
+ data.Write(chunk)
45
+ fmt.Printf("chunk: %6.d %v\n", len(chunk), chunk[:10])
46
+ dch <- chunk
47
}
48
- close(dch)
48
}()
49
50
read := NewByteChanReader(dch)
@@ -54,7 +53,8 @@ func TestByteChanReader(t *testing.T) {
53
t.Fatal(err)
54
}
55
57
- if !bytes.Equal(out, data) {
56
+ // fmt.Printf("lens: %d == %d\n", len(out), len(data.Bytes()))
57
+ if !bytes.Equal(out, data.Bytes()) {
58
t.Fatal("Reader failed to stream correct bytes")
59
}
60
}