util: cleaner ByteChanReader.Read

util/util.go

@@ -60,40 +60,28 @@ func NewByteChanReader(in chan []byte) io.Reader {
 	return &byteChanReader{in: in}
 }
 
-func (bcr *byteChanReader) Read(b []byte) (int, error) {
-	if len(bcr.buf) == 0 {
-		data, ok := <-bcr.in
-		if !ok {
-			return 0, io.EOF
-		}
-		bcr.buf = data
-	}
-
-	if len(bcr.buf) >= len(b) {
-		copy(b, bcr.buf)
-		bcr.buf = bcr.buf[len(b):]
-		return len(b), nil
+func (bcr *byteChanReader) Read(output []byte) (int, error) {
+	remain := output
+	remainLen := len(output)
+	outputLen := 0
+	more := false
+	next := bcr.buf
+
+	for {
+		n := copy(remain, next)
+		remainLen -= n
+		outputLen += n
+		if remainLen == 0 {
+			bcr.buf = next[n:]
+			return outputLen, nil
 		}
 
-	copy(b, bcr.buf)
-	b = b[len(bcr.buf):]
-	totread := len(bcr.buf)
-
-	for data := range bcr.in {
-		if len(data) > len(b) {
-			totread += len(b)
-			copy(b, data[:len(b)])
-			bcr.buf = data[len(b):]
-			return totread, nil
-		}
-		copy(b, data)
-		totread += len(data)
-		b = b[len(data):]
-		if len(b) == 0 {
-			return totread, nil
+		remain = remain[n:]
+		next, more = <-bcr.in
+		if !more {
+			return outputLen, io.EOF
 		}
 	}
-	return totread, io.EOF
 }
 
 type randGen struct {

util/util_test.go

@@ -2,6 +2,7 @@ package util
 
 import (
 	"bytes"
+	"fmt"
 	"io/ioutil"
 	"math/rand"
 	"testing"
@@ -30,22 +31,20 @@ func TestKey(t *testing.T) {
 }
 
 func TestByteChanReader(t *testing.T) {
-	data := make([]byte, 1024*1024)
- r := NewTimeSeededRand()
-	r.Read(data)
+
+	var data bytes.Buffer
 	dch := make(chan []byte, 8)
+	randr := NewTimeSeededRand()
 
 	go func() {
-		beg := 0
-		for i := 0; i < len(data); {
-			i += rand.Intn(100) + 1
-			if i > len(data) {
-				i = len(data)
-			}
-			dch <- data[beg:i]
-			beg = i
+		defer close(dch)
+		for i := 0; i < rand.Intn(100)+100; i++ {
+			chunk := make([]byte, rand.Intn(100000)+10)
+			randr.Read(chunk)
+			data.Write(chunk)
+			fmt.Printf("chunk: %6.d %v\n", len(chunk), chunk[:10])
+			dch <- chunk
 		}
-		close(dch)
 	}()
 
 	read := NewByteChanReader(dch)
@@ -54,7 +53,8 @@ func TestByteChanReader(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	if !bytes.Equal(out, data) {
+	// fmt.Printf("lens: %d == %d\n", len(out), len(data.Bytes()))
+	if !bytes.Equal(out, data.Bytes()) {
 		t.Fatal("Reader failed to stream correct bytes")
 	}
 }