drpcmanager: replace streamBuffer with a streams registry by shubhamdhama · Pull Request #47 · cockroachdb/drpc

shubhamdhama · 2026-03-30T04:33:16Z

Replace the single-stream streamBuffer with a stream registry that maps
stream IDs to stream objects. The registry currently holds at most one
active stream (two briefly during handoff), but provides the foundation
for stream multiplexing where callers will look up streams by ID directly.

drpcmanager/registry.go

cthumuluru-crdb · 2026-04-02T09:51:43Z

drpcmanager/registry.go

+	r.mu.RLock()
+	defer r.mu.RUnlock()
+
+	s, ok := r.streams[id]


Don't you have to check if registry is closed?

We can do that now since Close means remove all the entries. I can do it in next PR.

cthumuluru-crdb · 2026-04-02T09:52:06Z

drpcmanager/registry.go

+func (r *streamRegistry) Unregister(id uint64) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+


Don't you have to check if registry is closed?

I don't think so. If registry is closed the entry would have most likely removed and doing this is noop.

drpcmanager/registry.go

drpcmanager/manager.go

drpcmanager/registry.go

cthumuluru-crdb · 2026-04-02T10:02:32Z

drpcmanager/manager.go

 		m.log("TERM", func() string { return fmt.Sprint(err) })
 		m.sigs.tport.Set(m.tr.Close())
-		m.sbuf.Close()
+		m.reg.ForEach(func(s *drpcstream.Stream) { s.Close() })


Should this be s.Cancel() instead of s.Close()? s.Close() writes to the transport, and the current code already holds the lock and writes to the transport for every stream. I'm assuming this path isn't intended for mux streams?

This is just closing the registry. Anyways, Close will Cancel the streams now.

shubhamdhama · 2026-04-06T11:36:18Z

drpcmanager/manager.go

 	case <-m.sigs.term.Signal():
 		err := m.sigs.term.Err()
 		if errors.Is(err, io.EOF) {
 			err = context.Canceled
 			if stream.Kind() == drpc.StreamKindClient {
 				err = drpc.ClosedError.New("connection closed")
 			}
 		}
 		stream.Cancel(err)
 		<-stream.Finished()
 		m.sem.Recv()


[pp] we can remove it and cancel all the streams in reg.Close().

shubhamdhama · 2026-04-06T11:44:36Z

drpcmanager/manager.go

@@ -217,7 +220,7 @@
 	if m.sigs.term.Set(err) {


[pp] Let's depend on .term in the registry.Register. And close the streams in reg.Close and remove the .term dependency from manageStreams

I'll do in follow-up.

shubhamdhama

I'm going to address some of the feedback in a separate PR.

drpcmanager/registry.go

shubhamdhama · 2026-04-07T04:46:03Z

drpcmanager/registry.go

+	r.mu.RLock()
+	defer r.mu.RUnlock()
+
+	s, ok := r.streams[id]


We can do that now since Close means remove all the entries. I can do it in next PR.

drpcmanager/manager.go

shubhamdhama · 2026-04-07T04:48:14Z

drpcmanager/manager.go

@@ -217,7 +220,7 @@ func (m *Manager) terminate(err error) {
 	if m.sigs.term.Set(err) {


I'll do in follow-up.

shubhamdhama · 2026-04-07T04:49:10Z

drpcmanager/manager.go

 		m.log("TERM", func() string { return fmt.Sprint(err) })
 		m.sigs.tport.Set(m.tr.Close())
-		m.sbuf.Close()
+		m.reg.ForEach(func(s *drpcstream.Stream) { s.Close() })


This is just closing the registry. Anyways, Close will Cancel the streams now.

drpcmanager/registry.go

Replace the single-stream streamBuffer with a stream registry that maps stream IDs to stream objects. The registry currently holds at most one active stream (two briefly during handoff), but provides the foundation for stream multiplexing where callers will look up streams by ID directly.

cthumuluru-crdb

Just added one comment about a test that's removed. Otherwise changes look good to me.

cthumuluru-crdb · 2026-04-09T00:28:05Z

drpcmanager/active_streams.go

+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	if r.closed {


(for future commits), if you make this an atomic boolean, you can do this check outside. You can also move this check outside once you rely on manager closed signal.

drpcmanager/active_streams.go

drpcmanager/manager_test.go

shubhamdhama force-pushed the shubham/use-atomic-counter branch from f5869a5 to f204e04 Compare March 30, 2026 14:07

shubhamdhama force-pushed the shubham/add-stream-registry branch from 86bb186 to e16d021 Compare March 30, 2026 14:07

shubhamdhama force-pushed the shubham/use-atomic-counter branch from f204e04 to 6abe605 Compare March 30, 2026 14:25

shubhamdhama force-pushed the shubham/add-stream-registry branch from e16d021 to 0129d8f Compare March 30, 2026 14:25

shubhamdhama mentioned this pull request Mar 31, 2026

[wip] drpc refactor benchmarking playground cockroachdb/cockroach#167104

Open

cthumuluru-crdb reviewed Apr 2, 2026

View reviewed changes

shubhamdhama commented Apr 6, 2026

View reviewed changes

shubhamdhama force-pushed the shubham/use-atomic-counter branch from 6abe605 to e3212b7 Compare April 7, 2026 03:55

Base automatically changed from shubham/use-atomic-counter to stream-multiplexing April 7, 2026 03:56

shubhamdhama force-pushed the shubham/add-stream-registry branch from a4092e3 to 9155272 Compare April 7, 2026 05:09

shubhamdhama commented Apr 7, 2026

View reviewed changes

shubhamdhama changed the title ~~drpcmanager: replace streamBuffer with streamRegistry~~ drpcmanager: replace streamBuffer with a streams registry Apr 7, 2026

shubhamdhama force-pushed the shubham/add-stream-registry branch from 9155272 to e14c876 Compare April 8, 2026 07:27

cthumuluru-crdb approved these changes Apr 9, 2026

View reviewed changes

shubhamdhama merged commit 54db735 into stream-multiplexing Apr 13, 2026

shubhamdhama deleted the shubham/add-stream-registry branch April 13, 2026 11:53

		@@ -217,7 +220,7 @@ func (m *Manager) terminate(err error) {
		if m.sigs.term.Set(err) {

Conversation

shubhamdhama commented Mar 30, 2026

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shubhamdhama Apr 7, 2026 • 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shubhamdhama left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

shubhamdhama Apr 7, 2026 • 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

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!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

shubhamdhama Apr 7, 2026 •

edited

Loading

shubhamdhama Apr 7, 2026 •

edited

Loading