mgrpc: initial implementation of message channel

A varint prefixed framing format is defined for transport messages for
mgrpc. We may make changes to include a more structured frame, but this
initial work can be used to validate the concepts.

Signed-off-by: Stephen J Day <stephen.day@docker.com>

channel.go

@@ -0,0 +1,101 @@
+package mgrpc
+
+import (
+	"bufio"
+	"context"
+	"encoding/binary"
+	"io"
+	"net"
+
+	"github.com/containerd/containerd/log"
+	"github.com/gogo/protobuf/proto"
+	"github.com/pkg/errors"
+)
+
+type channel struct {
+	conn net.Conn
+	bw   *bufio.Writer
+	br   *bufio.Reader
+}
+
+func newChannel(conn net.Conn) *channel {
+	return &channel{
+		conn: conn,
+		bw:   bufio.NewWriter(conn),
+		br:   bufio.NewReader(conn),
+	}
+}
+
+func (ch *channel) recv(ctx context.Context, msg interface{}) error {
+	defer log.G(ctx).WithField("msg", msg).Info("recv")
+	var p [maxMessageSize]byte
+	n, err := readmsg(ch.br, p[:])
+	if err != nil {
+		return err
+	}
+
+	switch msg := msg.(type) {
+	case proto.Message:
+		return proto.Unmarshal(p[:n], msg)
+	default:
+		return errors.Errorf("unnsupported type in channel: %#v", msg)
+	}
+}
+
+func (ch *channel) send(ctx context.Context, msg interface{}) error {
+	log.G(ctx).WithField("msg", msg).Info("send")
+	var p []byte
+	switch msg := msg.(type) {
+	case proto.Message:
+		var err error
+		// TODO(stevvooe): trickiest allocation of the bunch. This will be hard
+		// to get rid of without using `MarshalTo` directly.
+		p, err = proto.Marshal(msg)
+		if err != nil {
+			return err
+		}
+	default:
+		return errors.Errorf("unsupported type recv from channel: %#v", msg)
+	}
+
+	return writemsg(ch.bw, p)
+}
+
+func readmsg(r *bufio.Reader, p []byte) (int, error) {
+	mlen, err := binary.ReadVarint(r)
+	if err != nil {
+		return 0, errors.Wrapf(err, "failed reading message size")
+	}
+
+	if mlen > int64(len(p)) {
+		return 0, errors.Wrapf(io.ErrShortBuffer, "message length %v over buffer size %v", mlen, len(p))
+	}
+
+	nn, err := io.ReadFull(r, p[:mlen])
+	if err != nil {
+		return 0, errors.Wrapf(err, "failed reading message size")
+	}
+
+	if int64(nn) != mlen {
+		return 0, errors.Errorf("mismatched read against message length %v != %v", nn, mlen)
+	}
+
+	return int(mlen), nil
+}
+
+func writemsg(w *bufio.Writer, p []byte) error {
+	var (
+		mlenp [binary.MaxVarintLen64]byte
+		n     = binary.PutVarint(mlenp[:], int64(len(p)))
+	)
+
+	if _, err := w.Write(mlenp[:n]); err != nil {
+		return errors.Wrapf(err, "failed writing message header")
+	}
+
+	if _, err := w.Write(p); err != nil {
+		return errors.Wrapf(err, "failed writing message")
+	}
+
+	return w.Flush()
+}

channel_test.go

@@ -0,0 +1,115 @@
+package mgrpc
+
+import (
+	"bufio"
+	"bytes"
+	"io"
+	"reflect"
+	"testing"
+
+	"github.com/pkg/errors"
+)
+
+func TestReadWriteMessage(t *testing.T) {
+	var (
+		channel  bytes.Buffer
+		w        = bufio.NewWriter(&channel)
+		messages = [][]byte{
+			[]byte("hello"),
+			[]byte("this is a test"),
+			[]byte("of message framing"),
+		}
+	)
+
+	for _, msg := range messages {
+		if err := writemsg(w, msg); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	var (
+		received [][]byte
+		r        = bufio.NewReader(bytes.NewReader(channel.Bytes()))
+	)
+
+	for {
+		var p [4096]byte
+		n, err := readmsg(r, p[:])
+		if err != nil {
+			if errors.Cause(err) != io.EOF {
+				t.Fatal(err)
+			}
+
+			break
+		}
+		received = append(received, p[:n])
+	}
+
+	if !reflect.DeepEqual(received, messages) {
+		t.Fatal("didn't received expected set of messages: %v != %v", received, messages)
+	}
+}
+
+func TestSmallBuffer(t *testing.T) {
+	var (
+		channel bytes.Buffer
+		w       = bufio.NewWriter(&channel)
+		msg     = []byte("a message of massive length")
+	)
+
+	if err := writemsg(w, msg); err != nil {
+		t.Fatal(err)
+	}
+
+	// now, read it off the channel with a small buffer
+	var (
+		p = make([]byte, len(msg)-1)
+		r = bufio.NewReader(bytes.NewReader(channel.Bytes()))
+	)
+	_, err := readmsg(r, p[:])
+	if err == nil {
+		t.Fatalf("error expected reading with small buffer")
+	}
+
+	if errors.Cause(err) != io.ErrShortBuffer {
+		t.Fatalf("errors.Cause(err) should equal io.ErrShortBuffer: %v != %v", err, io.ErrShortBuffer)
+	}
+}
+
+func BenchmarkReadWrite(b *testing.B) {
+	b.StopTimer()
+	var (
+		messages = [][]byte{
+			[]byte("hello"),
+			[]byte("this is a test"),
+			[]byte("of message framing"),
+		}
+		total   int64
+		channel bytes.Buffer
+		w       = bufio.NewWriter(&channel)
+		p       [4096]byte
+	)
+
+	b.ReportAllocs()
+	b.StartTimer()
+	for i := 0; i < b.N; i++ {
+		msg := messages[i%len(messages)]
+		if err := writemsg(w, msg); err != nil {
+			b.Fatal(err)
+		}
+		total += int64(len(msg))
+	}
+	b.SetBytes(total)
+
+	r := bufio.NewReader(bytes.NewReader(channel.Bytes()))
+	for i := 0; i < b.N; i++ {
+		_, err := readmsg(r, p[:])
+		if err != nil {
+			if errors.Cause(err) != io.EOF {
+				b.Fatal(err)
+			}
+
+			break
+		}
+	}
+}