drpcstream: introduce shared BufferPool for ring buffer by shubhamdhama · Pull Request #55 · cockroachdb/drpc

shubhamdhama · 2026-04-17T11:06:21Z

Add a BufferPool backed by sync.Pool that is shared across all streams
within a Manager. The ring buffer now obtains buffers from the pool on
Enqueue and transfers ownership to the caller on Dequeue, which advances
the tail immediately. This removes the two-step Dequeue/Done protocol
and simplifies Close (no longer needs to wait for held buffers).

The pool is a required parameter in the Stream constructor, created once
per Manager and passed to all streams it creates.

shubhamdhama · 2026-04-17T11:18:31Z

I had an idea which I ran through claude and below is the summary of it. I'm not planning to do it but in future we can re-consider if profiling shows any gain.

Buffer pool: further optimization ideas

Right now the data copy chain for an incoming message looks like this:

PacketAssembler.AppendFrame: copies frame data into pa.pk.Data via append
ringBuffer.Enqueue: copies pkt.Data into a pooled buffer via append
RawRecv copies the pooled buffer out, or MsgRecv unmarshals from it

We could eliminate copy #2 by having the packet assembler get its buffer from the pool directly. The assembler already has a TODO for buffer reuse. Instead of reusing its own backing array across packets (lines 84-87), it would pool.Get a buffer, assemble into it, hand it off through the ring buffer, and pool.Get a fresh one for the next packet.

Another idea: size-bucketed pools (e.g. 1KiB, 16KiB, 32KiB) so that Enqueue's append doesn't have to reallocate when messages are larger than the default capacity. You could even have a pool.Append(buf, data...) method that detects when the buffer needs to grow and fetches from the right bucket.

I think we should keep the pool simple for now. sync.Pool already gives you natural high-water-mark behavior: a buffer that grew to 32KiB stays at 32KiB when returned, so after warm-up the pool self-tunes to the workload's size distribution. Size buckets would add real maintenance cost (choosing boundaries, handling cross-bucket transitions) for a gain that append + sync.Pool already provides. Latency here is dominated by network IO anyway, not memcpy.

The assembler integration is the more interesting optimization since it removes a full copy per message. Worth revisiting once we have benchmarks to measure the actual impact.

cthumuluru-crdb · 2026-06-10T10:08:11Z

 	}

-	rb.buf[rb.head] = append(rb.buf[rb.head][:0], data...)
+	b := rb.pool.Get()


These can me move outside the lock now.

cthumuluru-crdb · 2026-06-10T10:11:58Z

-	data = append([]byte(nil), data...)
-	s.recvQueue.Done()
+	data = append([]byte(nil), *b...)
+	s.pool.Put(b)


We should continue to use s.recvQueue.Done(). Refer to other comments for context.

cthumuluru-crdb

Added a couple of comments. Besides that the changes look good.

cthumuluru-crdb · 2026-06-12T05:55:47Z

-	head  int      // next write position (producer)
-	tail  int      // next read position (consumer)
-	count int      // number of occupied slots
+	pool  *BufferPool // shared pool buffers are obtained from


nit: in case of Close() we are no longer waiting for the held buffer to be released. I don't see any risk in that but that would also mean, the consumer may return the buffer back to the pool after the packet buffer is closed (maybe due to context cancellation). The buffer should still be returned to the pool since the pool is shared across streams and is at the connection level.

Can you add a comment that buffer pool is shared across streams and is owned by the manager and not by ring buffer (just to be more specific about the ownership)?

I will add the comment but as for the rest of the comment, i couldn't understand the problem outlined.

I see you were talking about adding the comment only. Updated.

cthumuluru-crdb · 2026-06-12T05:59:34Z

 // increasing stream ids within a single transport.
-func New(ctx context.Context, sid uint64, wr *drpcwire.MuxWriter) *Stream {
-	return NewWithOptions(ctx, sid, wr, Options{})
+func New(ctx context.Context, sid uint64, wr *drpcwire.MuxWriter, pool *BufferPool) *Stream {


Since buffer pool is an internal option to a stream, I prefer Internal options to pass it around.

I don't understand this one. It's not an optional parameter so I don't see the value in wrapping it behind Internal struct.

cthumuluru-crdb · 2026-06-12T08:35:48Z

 	m.pendingStreams = make(map[uint64]*pendingStream)

 	m.streams = newActiveStreams()
+	m.recvPool = drpcstream.NewBufferPool()


Since manager uses it, how about we move it to the package where manager belongs? Also, keeping the name as buffer pool would enable us to use it for send buffer if needed in the future.

I want to move it to drpcwire in an upcoming PR as it's a lower in dependency hierarchy and contains similar fundamental pieces. I don't want to move it to drpcmanager or drpcstream because we have to create interfaces all over the place.

Add a BufferPool backed by sync.Pool that is shared across all streams within a Manager. The ring buffer now obtains a pooled buffer on Enqueue and copies the message into it, instead of growing a fixed slice per slot. Dequeue hands the buffer's data to the consumer and advances the tail immediately, and Done releases the buffer back to the pool once the consumer is finished with it. Keeping the Dequeue/Done contract means the consumer works with a plain []byte and never has to know whether the queue is backed by a pool or by fixed buffers. Because Dequeue advances the tail right away rather than waiting for Done, Close no longer has to block until in-flight buffers are released. The pool is a required parameter in the Stream constructor, created once per Manager and passed to all streams it creates. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

shubhamdhama force-pushed the shubham/enable-stream-multiplexing branch from a17330d to b91bf1b Compare April 17, 2026 14:57

shubhamdhama force-pushed the shubham/buffer-pool-for-ringbuffer branch from cafa1dc to b3d2355 Compare April 17, 2026 14:57

shubhamdhama force-pushed the shubham/enable-stream-multiplexing branch from b91bf1b to a58986c Compare April 17, 2026 15:00

shubhamdhama force-pushed the shubham/buffer-pool-for-ringbuffer branch from b3d2355 to 38c84dc Compare April 17, 2026 15:00

Base automatically changed from shubham/enable-stream-multiplexing to stream-multiplexing April 17, 2026 16:30

shubhamdhama force-pushed the shubham/buffer-pool-for-ringbuffer branch from 38c84dc to f2f767f Compare May 11, 2026 11:03

shubhamdhama force-pushed the shubham/buffer-pool-for-ringbuffer branch from f2f767f to a305254 Compare June 9, 2026 07:36

shubhamdhama changed the base branch from stream-multiplexing to main June 9, 2026 07:37

cthumuluru-crdb reviewed Jun 10, 2026

View reviewed changes

shubhamdhama force-pushed the shubham/buffer-pool-for-ringbuffer branch from a305254 to f5edd92 Compare June 11, 2026 14:50

shubhamdhama requested a review from cthumuluru-crdb June 11, 2026 14:57

cthumuluru-crdb approved these changes Jun 12, 2026

View reviewed changes

shubhamdhama force-pushed the shubham/buffer-pool-for-ringbuffer branch from f5edd92 to 03f1339 Compare June 12, 2026 09:33

shubhamdhama merged commit 9029d2e into main Jun 12, 2026
3 checks passed

shubhamdhama deleted the shubham/buffer-pool-for-ringbuffer branch June 12, 2026 09:33

Conversation

shubhamdhama commented Apr 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

shubhamdhama commented Apr 17, 2026

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cthumuluru-crdb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

shubhamdhama commented Apr 17, 2026 •

edited

Loading