Skip to content

*: add opt-in stream multiplexing support#34

Open
shubhamdhama wants to merge 1 commit intocockroachdb:mainfrom
shubhamdhama:stream-mux-with-no-mux
Open

*: add opt-in stream multiplexing support#34
shubhamdhama wants to merge 1 commit intocockroachdb:mainfrom
shubhamdhama:stream-mux-with-no-mux

Conversation

@shubhamdhama
Copy link

@shubhamdhama shubhamdhama commented Mar 5, 2026

DRPC currently allows only one active stream per transport at a time. This commit adds multiplexing as an opt-in mode (Options.Mux) that enables concurrent streams over a single transport, while preserving the original sequential behavior as the default.

Mux and non-mux paths have fundamentally different concurrency models, so a single Manager with conditionals would add branching in hot paths. Instead, two separate types are used: Manager (non-mux, unchanged) and MuxManager (new, in manager_mux.go). MuxManager runs two goroutines: manageReader routes packets to streams via a streamRegistry, and manageWriter batches frames from a sharedWriteBuf into transport writes.

Elaborate details on separating MuxManager and existing Manager

The alternative was a single Manager with a mux bool field and conditionals throughout. This was rejected because the internal state is fundamentally different between modes. Manager uses a semaphore (drpcsignal.Chan) for single-stream-at-a-time, a streamBuffer for tracking the current stream, a shared *drpcwire.Writer for direct transport writes, and manageReader + manageStreams goroutines. MuxManager uses a streamRegistry (thread-safe map by stream ID), a sharedWriteBuf for batched writes, a sync.WaitGroup for per-stream goroutines, an atomic.Uint64 for client stream ID generation, and manageReader + manageWriter goroutines. Branching across all of these would make both paths harder to reason about.

Shared code stays at package level: the Options struct (with Mux bool added), managerClosed error class, isConnectionReset helper, and closedCh pre-closed channel.


The Stream now accepts a StreamWriter interface instead of *Writer directly, so both paths share the same Stream code. Non-mux uses *Writer (direct transport writes), mux uses muxWriter (serializes each packet atomically into the shared buffer to prevent frame interleaving).


Packet buffering is abstracted behind a packetStore interface with two implementations: syncPacketBuffer (blocking single-slot, non-mux) and queuePacketBuffer (non-blocking queue with sync.Pool recycling, mux). In non-mux mode, Put blocks the reader until the consumer calls Get + Done, providing natural backpressure since there's only one stream. In mux mode, Put appends and returns immediately so the reader stays unblocked to dispatch packets to other concurrent streams. RawRecv and MsgRecv branch on the mux flag for buffer lifecycle: non-mux copies data then calls Done() to unblock the reader, mux takes ownership and recycles after consumption.

packetStore details

Buffer ownership in mux mode: AcquirePacketBuf() gets a buffer from the pool, the reader fills it, HandlePacket calls Put (transfers ownership to the queue), the reader calls AcquirePacketBuf() again for a fresh buffer. On the consumer side, Get takes ownership from the queue, and the consumer calls Recycle (or RawRecv returns the buffer directly to the caller).

The pool uses a pktBuf wrapper struct instead of storing bare []byte values. Without the wrapper, every sync.Pool.Put would box the slice header into an interface, causing an allocation per Put. The wrapper is a pointer type, so pool operations don't allocate.

Close behavior differs by error: with a non-EOF error, queued buffers are returned to the pool immediately (the stream is broken, no one will read them). With io.EOF, queued buffers are preserved so readers can drain remaining messages before seeing EOF.


HandlePacket now processes KindMessage before checking the term signal. This is needed for mux mode where Put must always run to return pool buffers. In non-mux mode, Put on a closed syncPacketBuffer returns immediately, so there is no behavioral change. Though there is a brief blocking window if terminate() has set the term signal but hasn't called pbuf.Close() yet, but is safe.

Details

There is a race window worth noting: terminate() sets term before calling pbuf.Close(). If the reader goroutine enters HandlePacket with a KindMessage during that window, Put can enter syncPacketBuffer while it's not yet closed. Three cases:

  • Slot occupied (blocking at for pb.set && pb.err == nil): Close sets pb.err and broadcasts, waking Put, which exits immediately.
  • Slot empty, no consumer (Put places data, blocks at for pb.set || pb.held): Close sees pb.held == false, sets pb.set = false and pb.err, broadcasts. Put wakes and exits.
  • Slot empty, consumer active (Put places data, consumer calls Get then Done): Put completes normally, Close runs afterward.

No deadlock in any case. The third case means the message gets delivered during termination, which is fine. The goal is to close the stream, not to prevent an in-flight message from being consumed.


The reader's monotonicity check is relaxed from global to per-stream so interleaved frames from different streams are accepted. Non-mux never produces interleaved stream IDs, so this has no behavioral impact there.


Conn uses a streamManager interface satisfied by both Manager types, branching once in the constructor. In mux mode, Invoke allocates a per-call marshal buffer instead of reusing a shared one to support concurrent calls. On the server side, ServeOne branches once to either serveOneNonMux (sequential handleRPC) or serveOneMux (concurrent handleRPC with WaitGroup).

@shubhamdhama
Copy link
Author

> benchdiff --bazel --count 20 --benchtime 1000x ./pkg/sql/tests -r Sysbench/SQL/3node/oltp_read_write
old:  f6dfcdc update dependency
new:  830adcf let's experiment drpc mux
args: benchdiff "--bazel" "--count" "20" "--benchtime" "1000x" "./pkg/sql/tests" "-r" "Sysbench/SQL/3node/oltp_read_write"

building benchmark binaries for 830adcf: let's experiment drpc mux [bazel=true] 1/1 -
name                                           old time/op    new time/op    delta
Sysbench/SQL/3node/oltp_read_write-24            12.3ms ± 2%    12.8ms ± 2%  +3.98%  (p=0.000 n=19+19)
ParallelSysbench/SQL/3node/oltp_read_write-24     954µs ± 5%    1021µs ± 8%  +7.04%  (p=0.000 n=18+19)

name                                           old errs/op    new errs/op    delta
Sysbench/SQL/3node/oltp_read_write-24              0.00           0.00         ~     (all equal)
ParallelSysbench/SQL/3node/oltp_read_write-24      0.01 ±37%      0.01 ±67%    ~     (p=0.241 n=20+20)

name                                           old alloc/op   new alloc/op   delta
Sysbench/SQL/3node/oltp_read_write-24            1.56MB ± 5%    1.63MB ± 5%  +4.54%  (p=0.000 n=19+20)
ParallelSysbench/SQL/3node/oltp_read_write-24    1.58MB ± 8%    1.66MB ± 7%  +5.34%  (p=0.002 n=20+20)

name                                           old allocs/op  new allocs/op  delta
ParallelSysbench/SQL/3node/oltp_read_write-24     6.71k ± 5%     6.71k ± 5%    ~     (p=0.703 n=20+20)
Sysbench/SQL/3node/oltp_read_write-24             7.02k ± 0%     7.04k ± 1%  +0.29%  (p=0.010 n=17+17)

@shubhamdhama shubhamdhama marked this pull request as ready for review March 5, 2026 14:57
@shubhamdhama
Copy link
Author

shubhamdhama commented Mar 10, 2026

Below is an AI assisted review of impact on no-multiplexing:

"1" isn't a concern because we don't use it in cockroach for no-mux. I'll create an issue regarding this before I merge this PR to explore if we need framing for no-mux and mux and some benchmarks backing that decision.

"6" is something I can dig more deeper, likely in a follow-up.

"7" I can remove as a follow-up.

Details

#: 1
Change: SplitSize silently ignored
Impact: Large messages no longer split into multiple frames. The send/term
check between splits is gone.
Risk: Medium — could affect existing users relying on SplitSize for flow
control or memory management
────────────────────────────────────────
#: 2
Change: HandlePacket calls Put before term check for messages
Impact: In non-mux, syncPacketBuffer.Put returns immediately if closed, so
effectively safe
Risk: Low
────────────────────────────────────────
#: 3
Change: Reader allows cross-stream non-monotonic IDs
Impact: In non-mux (single stream), equivalent to before. New zero-ID
rejection
is stricter.
Risk: Low (positive)
────────────────────────────────────────
#: 4
Change: Writer.Reset() moved to manager
Impact: No impact if only used through Manager
Risk: Low
────────────────────────────────────────
#: 5
Change: Stream constructor takes interface
Impact: Public API change, source-compatible but breaks type assertions
Risk: Low
────────────────────────────────────────
#: 6
Change: packetBuffer → interface
Impact: One extra heap alloc per stream, interface dispatch overhead
Risk: Low (perf micro-regression)
────────────────────────────────────────
#: 7
Change: SplitSize is still a public option
Impact: Confusing since it's now dead code
Risk: Low (API hygiene)

@shubhamdhama
Copy link
Author

shubhamdhama commented Mar 10, 2026

Checklists

Explanations

  • Impact on non-multiplexing world. (Explained in this comment)
  • queuePacketBuffer
  • MuxManager
  • sharedWriteBuf
  • SplitSize

Immediate follow-ups,

  • add syscall benchmark numbers

Follow-ups

  • Move internal/drpcopts package to drpcstream package
  • Remove drpccache package. It seems practically unused.

Copy link
Author

@shubhamdhama shubhamdhama Mar 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[AcIm] Move the content of this file to mux_writer_test.go

"storj.io/drpc/drpcdebug"
)

// StreamWriter is the interface used by streams for writing packets. Each call
Copy link
Author

@shubhamdhama shubhamdhama Mar 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[AcIm] Keep this comment at interface level, don't leak the implementation detail here.

case fr.ID.Stream == 0 || fr.ID.Message == 0:
return Packet{}, drpc.ProtocolError.New("id monotonicity violation (fr:%v r:%v)", fr.ID, r.id)

case fr.ID.Stream == r.id.Stream && fr.ID.Less(r.id):
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is where I have relaxed the Frame ID constraint even for no-mux scenario.

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you elaborate on the reason to relax this?

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How would this work if frames of different streams arrive one after another and one of them is out of order?
Say, we got frames in this sequence (S: Stream, F: Frame)
S1 F1 --> allowed, set r.id = (1, 1)
S1 F2 --> allowed since F2 > F1, r.id = (1, 2)
S2 F1 --> allowed, since r.id.Stream != fr.ID.Stream, r.id = (2, 1)
S1 F2 --> allowed, since r.id.Stream != fr.ID.Stream, and the monotonicity check is skipped
We would need to track the last seen frame message id per stream and compare against that.

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah ok, I think this will not matter with the current implementation as we have now changed it to have exactly one frame per packet.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Exactly.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

But I plan to change this in upcoming PR. But that's sorta improvement without which this PR is still functional, though prone to application level HOL.


if s.sigs.term.IsSet() {
// Put must always be called for message packets to manage buffer
// ownership. Put returns the buffer to the pool if the stream is closed.
Copy link
Author

@shubhamdhama shubhamdhama Mar 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[AcIm] Add a comment why it's safe for no-mux.

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, clarify that this comment applies only for mux

@shubhamdhama shubhamdhama force-pushed the stream-mux-with-no-mux branch from eb90688 to 6232cfe Compare March 10, 2026 13:19
DRPC currently allows only one active stream per transport at a time.
This commit adds multiplexing as an opt-in mode (Options.Mux) that
enables concurrent streams over a single transport, while preserving
the original sequential behavior as the default.

Mux and non-mux paths have fundamentally different concurrency models,
so a single Manager with conditionals would add branching in hot paths.
Instead, two separate types are used: Manager (non-mux, unchanged) and
MuxManager (new, in manager_mux.go). MuxManager runs two goroutines:
manageReader routes packets to streams via a streamRegistry, and
manageWriter batches frames from a sharedWriteBuf into transport writes.

The Stream now accepts a StreamWriter interface instead of *Writer
directly, so both paths share the same Stream code. Non-mux uses *Writer
(direct transport writes), mux uses muxWriter (serializes each packet
atomically into the shared buffer to prevent frame interleaving).

Packet buffering is abstracted behind a packetStore interface with two
implementations: syncPacketBuffer (blocking single-slot, non-mux) and
queuePacketBuffer (non-blocking queue with sync.Pool recycling, mux).
RawRecv and MsgRecv branch on the mux flag for buffer lifecycle: non-mux
copies data then calls Done() to unblock the reader, mux takes ownership
and recycles after consumption.

HandlePacket now processes KindMessage before checking the term signal.
This is needed for mux mode where Put must always run to return pool
buffers. In non-mux mode, Put on a closed syncPacketBuffer returns
immediately, so there is no behavioral change. Though there is a brief
blocking window if terminate() has set the term signal but hasn't called
pbuf.Close() yet.

The reader's monotonicity check is relaxed from global to per-stream so
interleaved frames from different streams are accepted. Non-mux never
produces interleaved stream IDs, so this has no behavioral impact there.

Conn uses a streamManager interface satisfied by both Manager types,
branching once in the constructor. In mux mode, Invoke allocates a
per-call marshal buffer instead of reusing a shared one to support
concurrent calls. On the server side, ServeOne branches once to either
serveOneNonMux (sequential handleRPC) or serveOneMux (concurrent
handleRPC with WaitGroup).
@shubhamdhama shubhamdhama force-pushed the stream-mux-with-no-mux branch from 6232cfe to 783893f Compare March 10, 2026 13:33
// appropriate locks.
func (s *Stream) rawWriteLocked(kind drpcwire.Kind, data []byte) (err error) {
fr := s.newFrameLocked(kind)
n := s.opts.SplitSize
Copy link
Author

@shubhamdhama shubhamdhama Mar 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Optional AcIm] remove the SplitSize option.

Copy link

@cthumuluru-crdb cthumuluru-crdb left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@shubhamdhama - I'm still going through the entire PR but added some questions/comments for you to take a look.

case fr.ID.Stream == 0 || fr.ID.Message == 0:
return Packet{}, drpc.ProtocolError.New("id monotonicity violation (fr:%v r:%v)", fr.ID, r.id)

case fr.ID.Stream == r.id.Stream && fr.ID.Less(r.id):

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you elaborate on the reason to relax this?

// into grpc metadata in the context.
GRPCMetadataCompatMode bool

// Mux enables stream multiplexing on the transport, allowing multiple

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: Can you make it more descriptive? Since the RPC muxer is also names as Mux, I would like it if we rename this one. Also, this is a flag to indicate if stream multiplexing must be enabled or not.

if err := stream.HandlePacket(pkt); err != nil {
m.terminate(managerClosed.Wrap(err))
return
}

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Retain the check for ignoring old messages. I don't see a reason to drop it.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since the packets would come interleaved, there is no "old message" anymore.

}

// if any invoke sequence is being sent, forward it to be handled.
case pkt.Kind == drpcwire.KindInvoke || pkt.Kind == drpcwire.KindInvokeMetadata:

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You should still take care of the canceled message arriving first right? How do we protect against that?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Discussed below.


// silently drop packet for an unregistered stream
default:
m.log("DROP", pkt.String)

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You cannot really drop the packet silently in this case right? Reader goroutine is the one reading the packet and handing it off for the manager to create the stream. Lets assume, there is a race between "new stream" and "cancelation" packet. Depending on how the race you might end up creating a stream but endup silently dropping the cancel packet for that stream.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we have discussed this earlier, offline. If we receive a "cancellation" for a stream before its creation, it should be a protocol error. Otherwise, if we support the case of cancellation may appear before creation, then we have to keep some sort of memory of such stream IDs. Then we have to decide the duration for which we should persist that memory.

I am fine adding a support for this but this was never supported even in non-mux. I recall you observed this race, but IMO that was a bug rather than a norm.

} else if fr.Done {
return nil
}
if err := s.wr.WritePacket(pkt); err != nil {
Copy link

@suj-krishnan suj-krishnan Mar 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

About SplitSize, since we don't set it explicitly in cockroach, drpc would currently use a default frame size of 64 KB right? So this would be a change in the existing behaviour for the non-mux code path. Also, for the mux code path, there would be head-of-line blocking concerns if we don't split large messages.
Can we clarify the motivation for removing SplitSize in the PR?

Copy link
Author

@shubhamdhama shubhamdhama Mar 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ahh right, the default is n = 64 * 1024. Then I guess I should add it to the writer after-all. Thank you for catching this.

Reason for removing: for non-multiplexing I can't think why framing exists. I mean genuinely, for no-mux I don't understand it. For multiplexing I can understand the reasons of head of line blocking, which I think would further complicate the PR but still very much required, so it's the immediate follow-up of this PR.

For no-mux, if you have a packet of 1 MiB in memory, you chunk into 16 chunks and you write them one by one as 16 syscalls. On the other hand, if you write all of 1MiB at once, the TCP layer should be smart enough to chunk it. So without any benchmarking I don't know for sure why framing exists and that's why I have added an action item for it in #34 (comment).

That said, I also want to avoid deviating from existing behavior so I will add it back for no-mux.

Copy link
Author

@shubhamdhama shubhamdhama Mar 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[AcIm] add the SplitSize back for no-mux

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

On the other hand, if you write all of 1MiB at once, the TCP layer should be smart enough to chunk it.

I think it's naive way of looking at things. I'll do more digging here.

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't see a strong reason to keep it for non-mux (http 1.1 doesn't have any chunking for instance) but just wanted to point out the change in behaviour in the non-mux path.

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you are planning to benchmark in order to make the decision for the non-mux path, the existing chunking is something to keep in mind.

@shubhamdhama
Copy link
Author

HandlePacket: KindMessage before term check

HandlePacket now processes KindMessage (calls Put) before checking the term signal. In mux mode, Put must always run to return buffers to the pool via queuePacketBuffer. In non-mux mode, Put on a closed syncPacketBuffer checks pb.err != nil and returns immediately, so there's no behavioral change in steady state.

There is a race window worth noting: terminate() sets term before calling pbuf.Close(). If the reader goroutine enters HandlePacket with a KindMessage during that window, Put can enter syncPacketBuffer while it's not yet closed. Three cases:

  • Slot occupied (blocking at for pb.set && pb.err == nil): Close sets pb.err and broadcasts, waking Put, which exits immediately.
  • Slot empty, no consumer (Put places data, blocks at for pb.set || pb.held): Close sees pb.held == false, sets pb.set = false and pb.err, broadcasts. Put wakes and exits.
  • Slot empty, consumer active (Put places data, consumer calls Get then Done): Put completes normally, Close runs afterward.

No deadlock in any case. The third case means the message gets delivered during termination, which is fine. The goal is to close the stream, not to prevent an in-flight message from being consumed.


if s.sigs.term.IsSet() {
// Put must always be called for message packets to manage buffer
// ownership. Put returns the buffer to the pool if the stream is closed.

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, clarify that this comment applies only for mux

switch {
case fr.ID.Less(r.id):
case fr.ID.Stream == 0 || fr.ID.Message == 0:
return Packet{}, drpc.ProtocolError.New("id monotonicity violation (fr:%v r:%v)", fr.ID, r.id)
Copy link

@suj-krishnan suj-krishnan Mar 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

May want to rephrase the error since this is different from a monotonicity violation (as in, not really an out of order frame situation, more like an invalid id).

@suj-krishnan
Copy link

The non-mux code path looks good to me, except for the change w.r.t SplitSize 👍

Copy link
Author

@shubhamdhama shubhamdhama left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Answered some of the Chandra's comments.


// silently drop packet for an unregistered stream
default:
m.log("DROP", pkt.String)
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we have discussed this earlier, offline. If we receive a "cancellation" for a stream before its creation, it should be a protocol error. Otherwise, if we support the case of cancellation may appear before creation, then we have to keep some sort of memory of such stream IDs. Then we have to decide the duration for which we should persist that memory.

I am fine adding a support for this but this was never supported even in non-mux. I recall you observed this race, but IMO that was a bug rather than a norm.

}

// if any invoke sequence is being sent, forward it to be handled.
case pkt.Kind == drpcwire.KindInvoke || pkt.Kind == drpcwire.KindInvokeMetadata:
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Discussed below.

if err := stream.HandlePacket(pkt); err != nil {
m.terminate(managerClosed.Wrap(err))
return
}
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since the packets would come interleaved, there is no "old message" anymore.

rpc = string(pkt.Data)
streamCtx := ctx

if metadata := m.popMetadata(pkt.ID.Stream); metadata != nil {

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since this creates a new map for every stream, it may be good to keep some upper bound on the map size.

return
}

pb.data = append(pb.data, data)

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Need an upper bound on the size of this since we keep appending to it. If no MsgRecv() or the consumer lags for some reason, this will keep growing.

}
m.metaMu.Unlock()
// Cancel all active streams so they get a clear error.
m.reg.ForEach(func(_ uint64, s *drpcstream.Stream) {

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't need to cancel every stream explicitly here because the m.sigs.term.Set() at the beginning of terminate will unblock the select in all the manageStream goroutines, each of which will cancel its own stream. This is how it works in the non-mux as well.

m.sigs.tport.Set(m.tr.Close())
m.sw.Close()
m.metaMu.Lock()
for id := range m.meta {

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

May not need this either as this will be garbage collected.

defer r.mu.Unlock()

if r.closed {
return managerClosed.New("register")

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: rephrase to 'register called'

if m.sigs.term.Set(err) {
m.log("TERM", func() string { return fmt.Sprint(err) })
m.sigs.tport.Set(m.tr.Close())
m.sw.Close()

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Closing the shared write buffer and then closing the transport is cleaner. That way we avoid the unnecessary write to the transport (that will then fail) if the buffer has any leftover data.

// Drain swaps out accumulated bytes, giving the caller ownership of the
// returned slice. The internal buffer is replaced with spare (reset to zero
// length) so producers can continue appending without allocation.
func (sw *sharedWriteBuf) Drain(spare []byte) []byte {
Copy link

@suj-krishnan suj-krishnan Mar 20, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: Unused, only used in tests. Can consider moving to a test file.

m.meta[streamID] = metadata
}

func (m *MuxManager) popMetadata(streamID uint64) map[string]string {

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: Can we rename to getMetadata, since this is just a dict read? pop usually implies LIFO semantics (or atleast removing from end of a data structure), so it can be confusing.

case pkt.Kind == drpcwire.KindInvoke || pkt.Kind == drpcwire.KindInvokeMetadata:
select {
case m.pkts <- pkt:
m.pdone.Recv()

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The reader waits here for NewServerStream to pick up and process the packet each time an Invoke or InvokeMetadata packet arrives, which can block other packets from being read. I think this is because pkt.Data cannot be reused by the reader until the new stream is created. Since the number of invoke packets, in general, would be fewer than the number of message packets and invoke packets would tend to be small (just the rpc string), we can consider copying the packet here to remove this wait, incurring the cost of one small allocation per Invoke as a trade-off. This would also require that the size of metadata be fairly small but that isn't an unreasonable constraint, IMO.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants