drpc: add compression by Nukitt · Pull Request #64 · cockroachdb/drpc

Nukitt · 2026-06-24T06:01:33Z

Introduce per-message compression for drpc streams. A Compression
enum type (CompressionNone, CompressionSnappy) in drpcwire provides
Compress/Decompress methods with built-in Snappy block-format support,
including buffer-reuse.

Compression is negotiated per-stream via the "drpc-compressor"
metadata key: the client injects the algorithm name into outgoing
metadata, and the server resolves it via CompressionFromName at
stream creation time, returning a ProtocolError immediately if the
algorithm is unknown.

Compression and decompression are performed in MsgSend/MsgRecv and
RawWrite/RawRecv via symmetric compressLocked/decompressLocked
helpers, keeping handlePacket and rawWriteLocked as pure transport
with zero compression awareness. Scratch buffers are cleanly owned
by their respective lock sides (cbuf under write, dbuf under read).
This also means the ring buffer enqueue copies smaller compressed
bytes rather than larger decompressed bytes, and decompression
errors are scoped to the affected stream's MsgRecv/RawRecv call
rather than tearing down all streams on the connection via the
manager's terminate path.

The compression metadata key is stripped from application-visible
metadata before populating the context. Only KindMessage frames
are compressed; control frames flow through unmodified.

Streams without compression are completely untouched so no new
wire framing, no per-message overhead, full backward compatibility.

shubhamdhama

Now that I see https://github.com/cockroachlabs/cockroach/pull/1150, it seems we are not re-using the compressors. If that's the case then I'm strongly not in favor of Compressor interface and this whole registration/init business. Instead of having WithCompressor, just have WithCompression where you pass the compression algorithm, Snappy, Zstd. That is, cockroach library should be able to use compression provided by DRPC library using just an enum from its perspective.

The reason gRPC has done it that way because it's a general purpose library. Here we own DRPC. We can add new compressors as we like in this repo.

Also, the code would be much less.

cthumuluru-crdb · 2026-06-27T09:39:20Z

Any reason the whole compression cannot be pushed to MsgSend/MsgReceive APIs? The sender compresses the buffer and replaces it as message data. Received receives compressed buffer, knows its compressed, uncompresses it and decodes it?

Nukitt · 2026-06-30T10:12:57Z

So some major changes I've incorporated after the initial round of reviews.

Starting with @shubhamdhama's suggestion of moving registration into drpc itself, now on cockroach a single dial option (WithCompression) is all that's needed and drpc automatically injects the compressor name into stream metadata, and the server extracts, validates, and applies it at stream creation time. This indeed keeps the boundary clean as drpc owns compression end-to-end and consumers opt in with a constant.

From @cthumuluru-crdb's suggestion, now compression and decompression happen entirely within the MsgSend/MsgRecv and RawWrite/RawRecv methods via symmetric compressLocked/decompressLocked helpers, keeping it out of the manager's read-loop hot path where handlePacket now just enqueues raw bytes without touching compression at all. This also means the enqueue copy is over smaller compressed bytes, rather than the earlier larger decompressed bytes. handlePacket and rawWriteLocked are back to pure transport with zero compression awareness.
The old compressAndWriteLocked helper was mixing compression with writing which is now been split so compressLocked and decompressLocked are both pure data transforms, with the scratch buffers cleanly owned by their respective lock sides (cbuf under write, dbuf under read). Elaborating further on moving decompression out of handlePacket fixes a blast-radius issue: a previous decompression error was returned from handlePacket, which caused the manager to call terminate() and tear down all streams on the connection. Now decompression errors surface as ProtocolError in the individual MsgRecv/RawRecv call, scoped to the affected stream. Also, only KindMessage frames are compressed and decompressed. Control frames (KindError, KindCancel, KindClose, KindCloseSend) flow uncompressed through handlePacket and are handled as protocol signals.

shubhamdhama

LGTM!

Please update the description and before i click approve I want to see it working.

Nukitt · 2026-07-01T10:03:24Z

TFTR! @shubhamdhama
Updated desc.

cthumuluru-crdb · 2026-07-02T03:06:35Z

-func (m *Manager) newStream(ctx context.Context, sid uint64, kind drpc.StreamKind, rpc string) (*drpcstream.Stream, error) {
+func (m *Manager) newStream(ctx context.Context, sid uint64, kind drpc.StreamKind, rpc string, compressor drpcwire.Compression) (*drpcstream.Stream, error) {
 	opts := m.opts.Stream
+	if compressor != drpcwire.CompressionNone {


How about implementing no-op compressor and make rest of the code free from checking the compression type?

cthumuluru-crdb · 2026-07-02T03:07:52Z

+		var compressor drpcwire.Compression
 		if pkt.metadata != nil {
+			if name, ok := pkt.metadata[drpcwire.CompressorMetadataKey]; ok {
+				comp, found := drpcwire.CompressionFromName(name)


CompressionFromName -> CompressorFromName ?

cthumuluru-crdb · 2026-07-02T03:11:44Z

+			if name, ok := pkt.metadata[drpcwire.CompressorMetadataKey]; ok {
+				comp, found := drpcwire.CompressionFromName(name)
+				if !found {
+					m.pdone.Send()


nit: separately, we should consider removing this hand-off. I don't see value in it except for the overhead and code complexity.

I'll keep it in for now and track it separately since this handoff ensures manageReader waits until the stream is registered before reading the next frame.

cthumuluru-crdb · 2026-07-02T03:17:01Z

 	wr        *drpcwire.MuxWriter
 	recvQueue ringBuffer
 	wbuf      []byte
+	cbuf      []byte // compression scratch buffer


I think we can avoid this shared state altogether. Can you check?

cbuf/dbuf provide buffer reuse across messages hence avoiding per-message allocations for snappy so i would keep it in.

cthumuluru-crdb · 2026-07-02T03:18:11Z

 	defer s.write.Unlock()

+	if kind == drpcwire.KindMessage {
+		data = s.compressLocked(data)


How about making this call inline and use the buffer to compress, write and let go of the buffer? That way you can avoid the buffer escaping to heap?

cthumuluru-crdb · 2026-07-02T03:18:55Z

 	}
+	defer s.recvQueue.Done()
+
+	b, err = s.decompressLocked(b)


I would suggest the same here too. Inline the code and use a local buffer?

cthumuluru-crdb · 2026-07-02T03:26:14Z

+	ctx := drpctest.NewTracker(t)
+	defer ctx.Close()
+
+	cli, close := createCompressedConnection(t, standardImpl)


nit: createCompressedConnection -> createConnectionWithCompression or let the createConnection take an optional compression type?

cthumuluru-crdb · 2026-07-02T03:30:47Z

So some major changes I've incorporated after the initial round of reviews.

Starting with @shubhamdhama's suggestion of moving registration into drpc itself, now on cockroach a single dial option (WithCompression) is all that's needed and drpc automatically injects the compressor name into stream metadata, and the server extracts, validates, and applies it at stream creation time. This indeed keeps the boundary clean as drpc owns compression end-to-end and consumers opt in with a constant.

From @cthumuluru-crdb's suggestion, now compression and decompression happen entirely within the MsgSend/MsgRecv and RawWrite/RawRecv methods via symmetric compressLocked/decompressLocked helpers, keeping it out of the manager's read-loop hot path where handlePacket now just enqueues raw bytes without touching compression at all. This also means the enqueue copy is over smaller compressed bytes, rather than the earlier larger decompressed bytes. handlePacket and rawWriteLocked are back to pure transport with zero compression awareness. The old compressAndWriteLocked helper was mixing compression with writing which is now been split so compressLocked and decompressLocked are both pure data transforms, with the scratch buffers cleanly owned by their respective lock sides (cbuf under write, dbuf under read). Elaborating further on moving decompression out of handlePacket fixes a blast-radius issue: a previous decompression error was returned from handlePacket, which caused the manager to call terminate() and tear down all streams on the connection. Now decompression errors surface as ProtocolError in the individual MsgRecv/RawRecv call, scoped to the affected stream. Also, only KindMessage frames are compressed and decompressed. Control frames (KindError, KindCancel, KindClose, KindCloseSend) flow uncompressed through handlePacket and are handled as protocol signals.

@Nukitt, can you also details on how compression can be turned off or on. Do we have a cluster setting using which can turn compression on/off at runtime?

Introduce per-message compression for drpc streams. A Compression enum type (CompressionNone, CompressionSnappy) in drpcwire provides Compress/Decompress methods with built-in Snappy block-format support, including buffer-reuse via scratch buffers owned by each lock side (cbuf under write, dbuf under read). Compression is negotiated per-stream via the "drpc-compressor" metadata key: the client injects the algorithm name into outgoing metadata, and the server resolves it via CompressorFromName at stream creation time, returning a ProtocolError immediately if the algorithm is unknown. A CompressorFunc callback on the manager options allows dynamic per-stream compression selection (e.g. after a version gate activates) without recycling the connection. Compression and decompression are performed in `MsgSend`/`MsgRecv` and `RawWrite`/`RawRecv` via symmetric `compressLocked`/`decompressLocked` helpers, keeping `handlePacket` and `rawWriteLocked` as pure transport with zero compression awareness. Scratch buffers are cleanly owned by their respective lock sides (`cbuf` under write, `dbuf` under read). This also means the ring buffer enqueue copies smaller compressed bytes rather than larger decompressed bytes, and decompression errors are scoped to the affected stream's `MsgRecv`/`RawRecv` call rather than tearing down all streams on the connection via the manager's terminate path. The compression metadata key is stripped from application-visible metadata before populating the context. Only `KindMessage` frames are compressed; control frames flow through unmodified. Streams without compression are completely untouched so no new wire framing, no per-message overhead, full backward compatibility.

shubhamdhama · 2026-07-06T12:32:06Z

+
+// CompressorMetadataKey is the metadata key used to signal the compression
+// algorithm from client to server during stream invocation.
+const CompressorMetadataKey = "drpc-compressor"


We should make this "drpc-compression".

Let's avoid using the word "compressor". In most of the context, we can just use compression.

shubhamdhama · 2026-07-06T12:33:42Z

+// per-stream to determine the compression algorithm. This allows
+// compression to be enabled dynamically (e.g. after a version gate
+// activates) without recycling the connection.
+func WithCompressorFunc(f func() drpcwire.Compression) DialOption {


WithCompression is enough unless there is an overloading.

shubhamdhama · 2026-07-06T12:34:38Z

 func (c *Conn) NewStream(ctx context.Context, rpc string, enc drpc.Encoding) (_ drpc.Stream, err error) {
 	defer func() { err = drpc.ToRPCErr(err) }()

+	comp := c.resolveCompressor()


Resolve this once for the whole connection.

shubhamdhama · 2026-07-06T12:52:37Z

+		maxLen := snappy.MaxEncodedLen(len(src))
+		if maxLen < 0 {
+			panic("drpcwire: snappy source too large for encoding")
+		}
+		if cap(dst) >= maxLen {
+			return snappy.Encode(dst[:cap(dst)], src)
+		}
+		return snappy.Encode(nil, src)


This whole block is redundant as it's present in snappy.Encode. The only line we need to keep is return snappy.Encode(dst[:cap(dst)], src) explaining why we are resetting the length to the capacity for destination.

Same for Decompress

More updates needed.

shubhamdhama · 2026-07-06T12:57:02Z

I think you are yet to update the Stream.

Nukitt requested review from cthumuluru-crdb, shubhamdhama and suj-krishnan June 24, 2026 06:04

Nukitt force-pushed the add-compression branch 3 times, most recently from bab59b3 to bd6926a Compare June 24, 2026 06:52

shubhamdhama reviewed Jun 24, 2026

View reviewed changes

Comment thread drpcstream/stream.go Outdated

Comment thread drpcconn/conn.go Outdated

Comment thread drpcmanager/manager.go Outdated

shubhamdhama requested changes Jun 24, 2026

View reviewed changes

Nukitt force-pushed the add-compression branch from bd6926a to 89f6464 Compare June 30, 2026 06:29

Nukitt requested a review from shubhamdhama June 30, 2026 10:13

shubhamdhama reviewed Jul 1, 2026

View reviewed changes

Comment thread drpcstream/stream.go Outdated

Nukitt force-pushed the add-compression branch from 89f6464 to f4be0fb Compare July 1, 2026 10:01

shubhamdhama previously approved these changes Jul 1, 2026

View reviewed changes

cthumuluru-crdb reviewed Jul 2, 2026

View reviewed changes

Nukitt force-pushed the add-compression branch from f4be0fb to ebb6054 Compare July 3, 2026 08:07

shubhamdhama self-requested a review July 6, 2026 12:17

shubhamdhama reviewed Jul 6, 2026

View reviewed changes

Uh oh!

Conversation

Nukitt commented Jun 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

shubhamdhama left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cthumuluru-crdb commented Jun 27, 2026

Uh oh!

Nukitt commented Jun 30, 2026

Uh oh!

shubhamdhama left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Nukitt commented Jul 1, 2026

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!

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!

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cthumuluru-crdb commented Jul 2, 2026

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!

Choose a reason for hiding this comment

Uh oh!

shubhamdhama commented Jul 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Nukitt commented Jun 24, 2026 •

edited

Loading

shubhamdhama left a comment •

edited

Loading