*: move frame assembly from reader to stream

Previously, Reader assembled wire frames into complete packets before
handing them to the manager. This change makes Reader return individual
frames (ReadFrame), and the stream handles frame assembly itself
(HandleFrame). The manager now enforces global frame ID monotonicity and
other validation that are beyond a stream's scope.

This is groundwork for stream multiplexing, where frames from different
streams will be interleaved on the wire and must be routed to the correct
stream before assembly.

Author: shubhamdhama

drpcconn/conn_test.go

@@ -43,9 +43,9 @@ func TestConn_InvokeFlushesSendClose(t *testing.T) {
 		wr := drpcwire.NewWriter(ps, 64)
 		rd := drpcwire.NewReader(ps)
 
-		_, _ = rd.ReadPacket()    // Invoke
-		_, _ = rd.ReadPacket()    // Message
-		pkt, _ := rd.ReadPacket() // CloseSend
+		_, _ = rd.ReadFrame()    // Invoke
+		_, _ = rd.ReadFrame()    // Message
+		pkt, _ := rd.ReadFrame() // CloseSend
 
 		_ = wr.WritePacket(drpcwire.Packet{
 			Data: []byte("qux"),
@@ -54,8 +54,8 @@ func TestConn_InvokeFlushesSendClose(t *testing.T) {
 		})
 		_ = wr.Flush()
 
-		_, _ = rd.ReadPacket() // Close
-		<-invokeDone           // wait for invoke to return
+		_, _ = rd.ReadFrame() // Close
+		<-invokeDone          // wait for invoke to return
 
 		// ensure that any later packets are dropped by writing one
 		// before closing the transport.
@@ -98,7 +98,7 @@ func TestConn_InvokeSendsGrpcAndDrpcMetadata(t *testing.T) {
 		wr := drpcwire.NewWriter(ps, 64)
 		rd := drpcwire.NewReader(ps)
 
-		md, err := rd.ReadPacket() // Metadata
+		md, err := rd.ReadFrame() // Metadata
 		assert.NoError(t, err)
 		assert.Equal(t, md.Kind, drpcwire.KindInvokeMetadata)
 		metadata, err := drpcmetadata.Decode(md.Data)
@@ -110,9 +110,9 @@ func TestConn_InvokeSendsGrpcAndDrpcMetadata(t *testing.T) {
 			"common-key":           "common-value2",
 		})
 
-		_, _ = rd.ReadPacket()    // Invoke
-		_, _ = rd.ReadPacket()    // Message
-		pkt, _ := rd.ReadPacket() // CloseSend
+		_, _ = rd.ReadFrame()    // Invoke
+		_, _ = rd.ReadFrame()    // Message
+		pkt, _ := rd.ReadFrame() // CloseSend
 
 		_ = wr.WritePacket(drpcwire.Packet{
 			Data: []byte("qux"),
@@ -121,7 +121,7 @@ func TestConn_InvokeSendsGrpcAndDrpcMetadata(t *testing.T) {
 		})
 		_ = wr.Flush()
 
-		_, _ = rd.ReadPacket() // Close
+		_, _ = rd.ReadFrame() // Close
 	})
 
 	conn := New(pc)
@@ -154,7 +154,7 @@ func TestConn_NewStreamSendsGrpcAndDrpcMetadata(t *testing.T) {
 	ctx.Run(func(ctx context.Context) {
 		rd := drpcwire.NewReader(ps)
 
-		md, err := rd.ReadPacket() // Metadata
+		md, err := rd.ReadFrame() // Metadata
 		assert.NoError(t, err)
 		assert.Equal(t, md.Kind, drpcwire.KindInvokeMetadata)
 		metadata, err := drpcmetadata.Decode(md.Data)
@@ -164,8 +164,8 @@ func TestConn_NewStreamSendsGrpcAndDrpcMetadata(t *testing.T) {
 			"drpc-key": "drpc-value",
 		})
 
-		_, _ = rd.ReadPacket() // Invoke
-		_, _ = rd.ReadPacket() // CloseSend
+		_, _ = rd.ReadFrame() // Invoke
+		_, _ = rd.ReadFrame() // CloseSend
 	})
 
 	conn := New(pc)

drpcmanager/manager.go

@@ -72,6 +72,9 @@ type Manager struct {
 	rd   *drpcwire.Reader
 	opts Options
 
+	lastFrameID   drpcwire.ID
+	lastFrameKind drpcwire.Kind
+
 	sem     drpcsignal.Chan      // held by the active stream
 	sbuf    streamBuffer         // largest stream id created
 	pkts    chan drpcwire.Packet // channel for invoke packets
@@ -213,27 +216,15 @@ func (m *Manager) terminate(err error) {
 // manage reader
 //
 
-// manageReader is always reading a packet and dispatching it to the appropriate
-// stream or queue. It sets the read signal when it exits so that one can wait
-// to ensure that no one is reading on the reader. It sets the term signal if
-// there is any error reading packets.
+// manageReader reads the frame and dispatches them to the appropriate stream or
+// queue. It sets the read signal when it exits so that one can wait to ensure
+// that no one is reading on the reader. It sets the term signal if there is any
+// error reading frames.
 func (m *Manager) manageReader() {
 	defer m.sigs.read.Set(nil)
 
-	var pkt drpcwire.Packet
-	var err error
-	var run int
-
 	for !m.sigs.term.IsSet() {
-		// if we have a run of "small" packets, drop the buffer to release
-		// memory so that a burst of large packets does not cause eternally
-		// large heap usage.
-		if run > 10 {
-			pkt.Data = nil
-			run = 0
-		}
-
-		pkt, err = m.rd.ReadPacketUsing(pkt.Data[:0])
+		incomingFrame, err := m.rd.ReadFrame()
 		if err != nil {
 			if isConnectionReset(err) {
 				err = drpc.ClosedError.Wrap(err)
@@ -242,36 +233,36 @@ func (m *Manager) manageReader() {
 			return
 		}
 
-		if len(pkt.Data) < cap(pkt.Data)/4 {
-			run++
-		} else {
-			run = 0
-		}
+		m.log("READ", incomingFrame.String)
 
-		m.log("READ", pkt.String)
+		if incomingFrame.ID.Less(m.lastFrameID) {
+			m.terminate(managerClosed.Wrap(drpc.ProtocolError.New("id monotonicity violation")))
+			return
+		}
 
 		switch curr := m.sbuf.Get(); {
-		// if the packet is for the current stream, deliver it.
-		case curr != nil && pkt.ID.Stream == curr.ID():
-			if err := curr.HandlePacket(pkt); err != nil {
+		// If the frame is for the current stream, deliver it.
+		case curr != nil && incomingFrame.ID.Stream == curr.ID():
+			if err := curr.HandleFrame(incomingFrame); err != nil {
 				m.terminate(managerClosed.Wrap(err))
 				return
 			}
 
-		// if an old message has been sent, just ignore it.
-		case curr != nil && pkt.ID.Stream < curr.ID():
+		// If a frame arrives for an old stream, just ignore it.
+		case curr != nil && incomingFrame.ID.Stream < curr.ID():
 
-		// if any invoke sequence is being sent, close any old unterminated
-		// stream and forward it to be handled.
-		case pkt.Kind == drpcwire.KindInvoke || pkt.Kind == drpcwire.KindInvokeMetadata:
+		// If an invoke sequence is being sent for a new stream, close any
+		// old unterminated stream and forward it to be handled.
+		case incomingFrame.Kind == drpcwire.KindInvoke || incomingFrame.Kind == drpcwire.KindInvokeMetadata:
 			if curr != nil && !curr.IsTerminated() {
 				curr.Cancel(context.Canceled)
 			}
 
+			pkt := drpcwire.Packet{ID: incomingFrame.ID, Kind: incomingFrame.Kind, Data: incomingFrame.Data}
 			select {
 			case m.pkts <- pkt:
 				// Wait for NewServerStream to finish stream creation (including
-				// sbuf.Set) before reading the next packet. This guarantees curr
+				// sbuf.Set) before reading the next frame. This guarantees curr
 				// is set for subsequent non-invoke packets.
 				m.pdone.Recv()
 
@@ -280,15 +271,21 @@ func (m *Manager) manageReader() {
 			}
 
 		default:
-			// A non-invoke packet arrived for a stream that doesn't exist yet
-			// (curr is nil or pkt.ID.Stream > curr.ID). The first packet of a
-			// new stream must be KindInvoke or KindInvokeMetadata.
-			m.terminate(managerClosed.Wrap(drpc.ProtocolError.New(
-				"first packet of a new stream must be Invoke, got %v (ID:%v)",
-				pkt.Kind,
-				pkt.ID)))
+			// A non-invoke frame arrived for a stream that doesn't exist yet
+			// (curr is nil or incomingFrame.ID.Stream > curr.ID). The first
+			// frame of a new stream must be KindInvoke or KindInvokeMetadata.
+			m.terminate(managerClosed.New(
+				"first frame of a new stream must be Invoke, got %v (ID:%v)",
+				incomingFrame.Kind,
+				incomingFrame.ID))
 			return
 		}
+
+		m.lastFrameKind = incomingFrame.Kind
+		m.lastFrameID = incomingFrame.ID
+		if incomingFrame.Done {
+			m.lastFrameID.Message += 1
+		}
 	}
 }
 

drpcmanager/manager_test.go

@@ -366,7 +366,7 @@ func TestManageReader_OldStreamFramesIgnored(t *testing.T) {
 
 // The first frame for a new stream must be KindInvoke or KindInvokeMetadata.
 // A non-invoke kind causes a protocol error.
-func Disabled_TestManageReader_FirstFrameMustBeInvoke(t *testing.T) {
+func TestManageReader_FirstFrameMustBeInvoke(t *testing.T) {
 	for _, kind := range []drpcwire.Kind{
 		drpcwire.KindMessage,
 		drpcwire.KindCancel,

drpcstream/stream.go

@@ -23,7 +23,7 @@ import (
 
 // Options controls configuration settings for a stream.
 type Options struct {
-	// SplitSize controls the default size we split packets into frames.
+	// SplitSize controls the default size we split data packets into frames.
 	SplitSize int
 
 	// ManualFlush controls if the stream will automatically flush after every
@@ -52,6 +52,11 @@ type Stream struct {
 	read  inspectMutex
 	flush sync.Once
 
+	assembling    bool
+	pktBuf        []byte
+	pktKind       drpcwire.Kind
+	nextMessageID uint64
+
 	id   drpcwire.ID
 	wr   *drpcwire.Writer
 	pbuf packetBuffer
@@ -98,6 +103,9 @@ func NewWithOptions(ctx context.Context, sid uint64, wr *drpcwire.Writer, opts O
 		fin:  drpcopts.GetStreamFin(&opts.Internal),
 		task: task,
 
+		// TODO: add a test case.
+		nextMessageID: 1,
+
 		id: drpcwire.ID{Stream: sid},
 		wr: wr.Reset(),
 	}
@@ -211,18 +219,58 @@ func (s *Stream) IsFinished() bool { return s.sigs.fin.IsSet() }
 func (s *Stream) SetManualFlush(mf bool) { s.opts.ManualFlush = mf }
 
 //
-// packet handler
+// frame handler
 //
 
-// HandlePacket advances the stream state machine by inspecting the packet. It
-// returns any major errors that should terminate the transport the stream is
-// operating on as well as a boolean indicating if the stream expects more
-// packets.
-func (s *Stream) HandlePacket(pkt drpcwire.Packet) (err error) {
-	if pkt.ID.Stream != s.id.Stream {
+// HandleFrame processes an incoming frame, assembling multi-frame packets
+// and dispatching complete packets to the stream state machine.
+func (s *Stream) HandleFrame(fr drpcwire.Frame) (err error) {
+	if fr.ID.Stream != s.ID() {
+		return drpc.ProtocolError.New("frame doesn't belong to this stream (fr: %v)", fr.ID)
+	}
+
+	if fr.ID.Message < s.nextMessageID {
+		return drpc.ProtocolError.New(
+			"id monotonicity violation: frame %v has message ID less than expected %v", fr.ID, s.nextMessageID)
+	} else if fr.ID.Message > s.nextMessageID || !s.assembling {
+		s.pktBuf = s.pktBuf[:0]
+		s.assembling = true
+		s.nextMessageID = fr.ID.Message
+	} else if fr.Kind != s.pktKind {
+		return drpc.ProtocolError.New("frame kind change within packet: got %v, expected %v", fr.Kind, s.pktKind)
+	}
+
+	// TODO(shubham): add buf reuse
+	s.pktBuf = append(s.pktBuf, fr.Data...)
+
+	s.pktKind = fr.Kind
+
+	if s.opts.MaximumBufferSize > 0 && len(s.pktBuf) > s.opts.MaximumBufferSize {
+		return drpc.ProtocolError.New("data overflow (len:%d)", len(s.pktBuf))
+	}
+
+	if !fr.Done {
 		return nil
 	}
 
+	s.assembling = false
+	s.nextMessageID = fr.ID.Message + 1
+
+	err = s.handlePacket(drpcwire.Packet{
+		ID:   fr.ID,
+		Kind: fr.Kind,
+		Data: s.pktBuf,
+	})
+
+	// TODO(shubham): add buf reuse
+	s.pktBuf = nil
+	return err
+}
+
+// handlePacket advances the stream state machine by inspecting the packet. It
+// returns any major errors that should terminate the transport the stream is
+// operating on.
+func (s *Stream) handlePacket(pkt drpcwire.Packet) (err error) {
 	drpcopts.GetStreamStats(&s.opts.Internal).AddRead(uint64(len(pkt.Data)))
 
 	if s.sigs.term.IsSet() {
@@ -240,7 +288,7 @@ func (s *Stream) HandlePacket(pkt drpcwire.Packet) (err error) {
 	defer s.mu.Unlock()
 
 	switch pkt.Kind {
-	case drpcwire.KindInvoke:
+	case drpcwire.KindInvoke, drpcwire.KindInvokeMetadata:
 		err := drpc.ProtocolError.New("invoke on existing stream")
 		s.terminate(err)
 		return err
@@ -375,9 +423,13 @@ func (s *Stream) RawWrite(kind drpcwire.Kind, data []byte) (err error) {
 
 // rawWriteLocked does the body of RawWrite assuming the caller is holding the
 // appropriate locks.
+// TODO(shubham): can we merge this with sendPacketLocked?
 func (s *Stream) rawWriteLocked(kind drpcwire.Kind, data []byte) (err error) {
 	fr := s.newFrameLocked(kind)
-	n := s.opts.SplitSize
+	n := 0
+	if kind == drpcwire.KindMessage {
+		n = s.opts.SplitSize
+	}
 
 	for {
 		switch {