Merge pull request #20 from stevvooe/pump-read-block

ttrpc: refactor channel to take a conn

channel.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"encoding/binary"
 	"io"
+	"net"
 	"sync"
 
 	"github.com/pkg/errors"
@@ -60,16 +61,18 @@ func writeMessageHeader(w io.Writer, p []byte, mh messageHeader) error {
 var buffers sync.Pool
 
 type channel struct {
+	conn  net.Conn
 	bw    *bufio.Writer
 	br    *bufio.Reader
 	hrbuf [messageHeaderLength]byte // avoid alloc when reading header
 	hwbuf [messageHeaderLength]byte
 }
 
-func newChannel(w io.Writer, r io.Reader) *channel {
+func newChannel(conn net.Conn) *channel {
 	return &channel{
-		bw: bufio.NewWriter(w),
-		br: bufio.NewReader(r),
+		conn: conn,
+		bw:   bufio.NewWriter(conn),
+		br:   bufio.NewReader(conn),
 	}
 }
 

channel_test.go

@@ -1,10 +1,10 @@
 package ttrpc
 
 import (
-	"bufio"
 	"bytes"
 	"context"
 	"io"
+	"net"
 	"reflect"
 	"testing"
 
@@ -16,28 +16,29 @@ import (
 func TestReadWriteMessage(t *testing.T) {
 	var (
 		ctx      = context.Background()
-		buffer   bytes.Buffer
-		w        = bufio.NewWriter(&buffer)
-		ch       = newChannel(w, nil)
+		w, r     = net.Pipe()
+		ch       = newChannel(w)
+		rch      = newChannel(r)
 		messages = [][]byte{
 			[]byte("hello"),
 			[]byte("this is a test"),
 			[]byte("of message framing"),
 		}
-	)
-
-	for i, msg := range messages {
-		if err := ch.send(ctx, uint32(i), 1, msg); err != nil {
-			t.Fatal(err)
-		}
-	}
-
-	var (
 		received [][]byte
-		r        = bufio.NewReader(bytes.NewReader(buffer.Bytes()))
-		rch      = newChannel(nil, r)
+		errs     = make(chan error, 1)
 	)
 
+	go func() {
+		for i, msg := range messages {
+			if err := ch.send(ctx, uint32(i), 1, msg); err != nil {
+				errs <- err
+				return
+			}
+		}
+
+		w.Close()
+	}()
+
 	for {
 		_, p, err := rch.recv(ctx)
 		if err != nil {
@@ -48,31 +49,44 @@ func TestReadWriteMessage(t *testing.T) {
 			break
 		}
 		received = append(received, p)
+
+		// make sure we don't have send errors
+		select {
+		case err := <-errs:
+			if err != nil {
+				t.Fatal(err)
+			}
+		default:
+		}
 	}
 
 	if !reflect.DeepEqual(received, messages) {
 		t.Fatalf("didn't received expected set of messages: %v != %v", received, messages)
 	}
+
+	select {
+	case err := <-errs:
+		if err != nil {
+			t.Fatal(err)
+		}
+	default:
+	}
 }
 
 func TestMessageOversize(t *testing.T) {
 	var (
-		ctx    = context.Background()
-		buffer bytes.Buffer
-		w      = bufio.NewWriter(&buffer)
-		ch     = newChannel(w, nil)
-		msg    = bytes.Repeat([]byte("a message of massive length"), 512<<10)
+		ctx      = context.Background()
+		w, r     = net.Pipe()
+		wch, rch = newChannel(w), newChannel(r)
+		msg      = bytes.Repeat([]byte("a message of massive length"), 512<<10)
+		errs     = make(chan error, 1)
 	)
 
-	if err := ch.send(ctx, 1, 1, msg); err != nil {
-		t.Fatal(err)
-	}
-
-	// now, read it off the channel with a small buffer
-	var (
-		r   = bufio.NewReader(bytes.NewReader(buffer.Bytes()))
-		rch = newChannel(nil, r)
-	)
+	go func() {
+		if err := wch.send(ctx, 1, 1, msg); err != nil {
+			errs <- err
+		}
+	}()
 
 	_, _, err := rch.recv(ctx)
 	if err == nil {
@@ -87,4 +101,12 @@ func TestMessageOversize(t *testing.T) {
 	if status.Code() != codes.ResourceExhausted {
 		t.Fatalf("expected grpc status code: %v != %v", status.Code(), codes.ResourceExhausted)
 	}
+
+	select {
+	case err := <-errs:
+		if err != nil {
+			t.Fatal(err)
+		}
+	default:
+	}
 }

client.go

@@ -35,7 +35,7 @@ func NewClient(conn net.Conn) *Client {
 	c := &Client{
 		codec:   codec{},
 		conn:    conn,
-		channel: newChannel(conn, conn),
+		channel: newChannel(conn),
 		calls:   make(chan *callRequest),
 		closed:  make(chan struct{}),
 		done:    make(chan struct{}),

server.go

@@ -281,7 +281,7 @@ func (c *serverConn) run(sctx context.Context) {
 	)
 
 	var (
-		ch          = newChannel(c.conn, c.conn)
+		ch          = newChannel(c.conn)
 		ctx, cancel = context.WithCancel(sctx)
 		active      int
 		state       connState = connStateIdle