*: 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.

Re-enables TestManageReader_FirstFrameMustBeInvoke and
TestManageReader_InvokeOnExistingStream, which were disabled in the
previous commit because the old code did not enforce these conditions.

This change also includes wire-protocol.md which can be considered as spec.

Author: shubhamdhama

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 packets.
 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,33 +233,35 @@ 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
+		}
 
 	again:
 		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():
+		case curr != nil && incomingFrame.ID.Stream < curr.ID():
+
+		// Either we are receiving first frame or a frame greater than the
+		// existing one in which case it should be an Invoke frame.
 
-		// if any invoke sequence is being sent, close any old unterminated
+		// 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:
+		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:
 				m.pdone.Recv()
@@ -277,19 +270,32 @@ func (m *Manager) manageReader() {
 				return
 			}
 
-		// a non-invoke packet should be delivered to some stream so we wait for
-		// a new stream to be created and try again. like an invoke, we
-		// implicitly close any previous stream.
+		// A non-invoke packet is delivered while the stream is being created
+		// with an invoke so we wait for a new stream to be created and try
+		// again. Like an invoke, we implicitly close any previous stream.
 		default:
+			if m.lastFrameKind != drpcwire.KindInvoke {
+				m.terminate(managerClosed.Wrap(
+					drpc.ProtocolError.New("first message of a stream cannot be of kind non-Invoke")))
+				return
+			}
+
 			if curr != nil && !curr.IsTerminated() {
 				curr.Cancel(context.Canceled)
 			}
 
 			if !m.sbuf.Wait(curr.ID()) {
 				return
 			}
+			// Note: it is safe to not update lastFrameKind and lastFrameID, as we are merely retrying
 			goto again
 		}
+
+		m.lastFrameKind = incomingFrame.Kind
+		m.lastFrameID = incomingFrame.ID
+		if incomingFrame.Done {
+			m.lastFrameID.Message += 1
+		}
 	}
 }
 

drpcmanager/manager_test.go

@@ -365,7 +365,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,
@@ -552,7 +552,7 @@ func TestManageReader_MultiFrameWithSkippedMessageID(t *testing.T) {
 
 // A second invoke for the same stream ID is rejected — the stream treats
 // it as a protocol error, terminating the manager.
-func Disabled_TestManageReader_InvokeOnExistingStream(t *testing.T) {
+func TestManageReader_InvokeOnExistingStream(t *testing.T) {
 	ctx := drpctest.NewTracker(t)
 	defer ctx.Close()
 

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(),
 	}
@@ -214,15 +222,54 @@ func (s *Stream) SetManualFlush(mf bool) { s.opts.ManualFlush = mf }
 // packet handler
 //
 
+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 (current frame:%v last frame:%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("packet kind change (fr:%v pkt:%v)", fr.Kind, s.pktKind)
+	}
+
+	// TODO(shubham): add buf reuse
+	s.pktBuf = append(s.pktBuf, fr.Data...)
+
+	s.pktKind = fr.Kind
+
+	// TODO:
+	// case len(pkt.Data) > r.opts.MaximumBufferSize:
+	// 	return Packet{}, drpc.ProtocolError.New("data overflow (len:%v)", len(pkt.Data))
+
+	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 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 {
-		return nil
-	}
-
+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 +287,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 +422,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 {