snapshot of https://github.com/grpc/grpc-go.git (master, 2026-09-15; blobs > 131072B omitted)

internal/grpcsync/callback_serializer.go

@@ -0,0 +1,124 @@
+/*
+ *
+ * Copyright 2022 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package grpcsync
+
+import (
+	"context"
+	"errors"
+
+	"google.golang.org/grpc/internal/buffer"
+)
+
+// ErrSerializerClosed is returned by ScheduleAndWait if the CallbackSerializer
+// was closed before the callback could be scheduled.
+var ErrSerializerClosed = errors.New("callback serializer is closed")
+
+// CallbackSerializer provides a mechanism to schedule callbacks in a
+// synchronized manner. It provides a FIFO guarantee on the order of execution
+// of scheduled callbacks. New callbacks can be scheduled by invoking the
+// Schedule() method.
+//
+// This type is safe for concurrent access.
+type CallbackSerializer struct {
+	// done is closed once the serializer is shut down completely, i.e all
+	// scheduled callbacks are executed and the serializer has deallocated all
+	// its resources.
+	done chan struct{}
+
+	callbacks *buffer.Unbounded[func(context.Context)]
+}
+
+// NewCallbackSerializer returns a new CallbackSerializer instance. The provided
+// context will be passed to the scheduled callbacks. Users should cancel the
+// provided context to shutdown the CallbackSerializer. It is guaranteed that no
+// callbacks will be added once this context is canceled, and any pending un-run
+// callbacks will be executed before the serializer is shut down.
+func NewCallbackSerializer(ctx context.Context) *CallbackSerializer {
+	cs := &CallbackSerializer{
+		done:      make(chan struct{}),
+		callbacks: buffer.NewUnbounded[func(context.Context)](),
+	}
+	go cs.run(ctx)
+	return cs
+}
+
+// TrySchedule tries to schedule the provided callback function f to be
+// executed in the order it was added. This is a best-effort operation. If the
+// context passed to NewCallbackSerializer was canceled before this method is
+// called, the callback will not be scheduled.
+//
+// Callbacks are expected to honor the context when performing any blocking
+// operations, and should return early when the context is canceled.
+func (cs *CallbackSerializer) TrySchedule(f func(ctx context.Context)) {
+	cs.callbacks.Put(f)
+}
+
+// ScheduleOr schedules the provided callback function f to be executed in the
+// order it was added. If the context passed to NewCallbackSerializer has been
+// canceled before this method is called, the onFailure callback will be
+// executed inline instead.
+//
+// Callbacks are expected to honor the context when performing any blocking
+// operations, and should return early when the context is canceled.
+func (cs *CallbackSerializer) ScheduleOr(f func(ctx context.Context), onFailure func()) {
+	if cs.callbacks.Put(f) != nil {
+		onFailure()
+	}
+}
+
+// ScheduleAndWait schedules the provided callback function f to be executed in
+// the order it was added and blocks until f has run. If the context passed to
+// NewCallbackSerializer was canceled before this method is called, f is not run
+// and ScheduleAndWait returns ErrSerializerClosed.
+//
+// Callbacks are expected to honor the context when performing any blocking
+// operations, and should return early when the context is canceled.
+func (cs *CallbackSerializer) ScheduleAndWait(f func(ctx context.Context)) error {
+	done := make(chan struct{})
+	var err error
+	cs.ScheduleOr(func(ctx context.Context) {
+		f(ctx)
+		close(done)
+	}, func() {
+		err = ErrSerializerClosed
+		close(done)
+	})
+	<-done
+	return err
+}
+
+func (cs *CallbackSerializer) run(ctx context.Context) {
+	defer close(cs.done)
+
+	// Close the buffer when the context is canceled
+	// to prevent new callbacks from being added.
+	context.AfterFunc(ctx, cs.callbacks.Close)
+
+	// Run all callbacks.
+	for cb := range cs.callbacks.Get() {
+		cs.callbacks.Load()
+		cb(ctx)
+	}
+}
+
+// Done returns a channel that is closed after the context passed to
+// NewCallbackSerializer is canceled and all callbacks have been executed.
+func (cs *CallbackSerializer) Done() <-chan struct{} {
+	return cs.done
+}

internal/grpcsync/callback_serializer_test.go

@@ -0,0 +1,247 @@
+/*
+ *
+ * Copyright 2022 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package grpcsync
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/google/go-cmp/cmp"
+)
+
+const (
+	defaultTestTimeout      = 5 * time.Second
+	defaultTestShortTimeout = 10 * time.Millisecond // For events expected to *not* happen.
+)
+
+// TestCallbackSerializer_Schedule_FIFO verifies that callbacks are executed in
+// the same order in which they were scheduled.
+func (s) TestCallbackSerializer_Schedule_FIFO(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	cs := NewCallbackSerializer(ctx)
+	defer cancel()
+
+	// We have two channels, one to record the order of scheduling, and the
+	// other to record the order of execution. We spawn a bunch of goroutines
+	// which record the order of scheduling and call the actual Schedule()
+	// method as well.  The callbacks record the order of execution.
+	//
+	// We need to grab a lock to record order of scheduling to guarantee that
+	// the act of recording and the act of calling Schedule() happen atomically.
+	const numCallbacks = 100
+	var mu sync.Mutex
+	scheduleOrderCh := make(chan int, numCallbacks)
+	executionOrderCh := make(chan int, numCallbacks)
+	for i := 0; i < numCallbacks; i++ {
+		go func(id int) {
+			mu.Lock()
+			defer mu.Unlock()
+			scheduleOrderCh <- id
+			cs.TrySchedule(func(ctx context.Context) {
+				select {
+				case <-ctx.Done():
+					return
+				case executionOrderCh <- id:
+				}
+			})
+		}(i)
+	}
+
+	// Spawn a couple of goroutines to capture the order or scheduling and the
+	// order of execution.
+	scheduleOrder := make([]int, numCallbacks)
+	executionOrder := make([]int, numCallbacks)
+	var wg sync.WaitGroup
+	wg.Add(2)
+	go func() {
+		defer wg.Done()
+		for i := 0; i < numCallbacks; i++ {
+			select {
+			case <-ctx.Done():
+				return
+			case id := <-scheduleOrderCh:
+				scheduleOrder[i] = id
+			}
+		}
+	}()
+	go func() {
+		defer wg.Done()
+		for i := 0; i < numCallbacks; i++ {
+			select {
+			case <-ctx.Done():
+				return
+			case id := <-executionOrderCh:
+				executionOrder[i] = id
+			}
+		}
+	}()
+	wg.Wait()
+
+	if diff := cmp.Diff(executionOrder, scheduleOrder); diff != "" {
+		t.Fatalf("Callbacks are not executed in scheduled order. diff(-want, +got):\n%s", diff)
+	}
+}
+
+// TestCallbackSerializer_Schedule_Concurrent verifies that all concurrently
+// scheduled callbacks get executed.
+func (s) TestCallbackSerializer_Schedule_Concurrent(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	cs := NewCallbackSerializer(ctx)
+	defer cancel()
+
+	// Schedule callbacks concurrently by calling Schedule() from goroutines.
+	// The execution of the callbacks call Done() on the waitgroup, which
+	// eventually unblocks the test and allows it to complete.
+	const numCallbacks = 100
+	var wg sync.WaitGroup
+	wg.Add(numCallbacks)
+	for i := 0; i < numCallbacks; i++ {
+		go func() {
+			cs.TrySchedule(func(context.Context) {
+				wg.Done()
+			})
+		}()
+	}
+
+	// We call Wait() on the waitgroup from a goroutine so that we can select on
+	// the Wait() being unblocked and the overall test deadline expiring.
+	done := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-ctx.Done():
+		t.Fatal("Timeout waiting for all scheduled callbacks to be executed")
+	case <-done:
+	}
+}
+
+// TestCallbackSerializer_Schedule_Close verifies that callbacks in the queue
+// are not executed once Close() returns.
+func (s) TestCallbackSerializer_Schedule_Close(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	serializerCtx, serializerCancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	cs := NewCallbackSerializer(serializerCtx)
+
+	// Schedule a callback which blocks until the context passed to it is
+	// canceled. It also closes a channel to signal that it has started.
+	firstCallbackStartedCh := make(chan struct{})
+	cs.TrySchedule(func(ctx context.Context) {
+		close(firstCallbackStartedCh)
+		<-ctx.Done()
+	})
+
+	// Schedule a bunch of callbacks. These should be executed since they are
+	// scheduled before the serializer is closed.
+	const numCallbacks = 10
+	callbackCh := make(chan int, numCallbacks)
+	for i := 0; i < numCallbacks; i++ {
+		num := i
+		callback := func(context.Context) { callbackCh <- num }
+		onFailure := func() { t.Fatal("Schedule failed to accept a callback when the serializer is yet to be closed") }
+		cs.ScheduleOr(callback, onFailure)
+	}
+
+	// Ensure that none of the newer callbacks are executed at this point.
+	select {
+	case <-time.After(defaultTestShortTimeout):
+	case <-callbackCh:
+		t.Fatal("Newer callback executed when older one is still executing")
+	}
+
+	// Wait for the first callback to start before closing the scheduler.
+	<-firstCallbackStartedCh
+
+	// Cancel the context which will unblock the first callback. All of the
+	// other callbacks (which have not started executing at this point) should
+	// be executed after this.
+	serializerCancel()
+
+	// Ensure that the newer callbacks are executed.
+	for i := 0; i < numCallbacks; i++ {
+		select {
+		case <-ctx.Done():
+			t.Fatal("Timeout when waiting for callback scheduled before close to be executed")
+		case num := <-callbackCh:
+			if num != i {
+				t.Fatalf("Executing callback %d, want %d", num, i)
+			}
+		}
+	}
+	<-cs.Done()
+
+	// Ensure that a callback cannot be scheduled after the serializer is
+	// closed.
+	done := make(chan struct{})
+	callback := func(context.Context) { t.Fatal("Scheduled a callback after closing the serializer") }
+	onFailure := func() { close(done) }
+	cs.ScheduleOr(callback, onFailure)
+	select {
+	case <-time.After(defaultTestTimeout):
+		t.Fatal("Successfully scheduled callback after serializer is closed")
+	case <-done:
+	}
+}
+
+// TestCallbackSerializer_ScheduleAndWait verifies that ScheduleAndWait runs the
+// provided callback and blocks until it has finished executing, returning a nil
+// error.
+func (s) TestCallbackSerializer_ScheduleAndWait(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	cs := NewCallbackSerializer(ctx)
+
+	executed := false
+	if err := cs.ScheduleAndWait(func(context.Context) {
+		executed = true
+	}); err != nil {
+		t.Fatalf("ScheduleAndWait() returned error: %v, want nil", err)
+	}
+
+	// The callback must have run by the time ScheduleAndWait returns. No
+	// synchronization is needed here because ScheduleAndWait is expected to
+	// block until the callback completes.
+	if !executed {
+		t.Fatal("ScheduleAndWait() returned before the callback was executed")
+	}
+}
+
+// TestCallbackSerializer_ScheduleAndWait_Closed verifies that ScheduleAndWait
+// does not run the callback and returns ErrSerializerClosed when the serializer
+// has already been closed.
+func (s) TestCallbackSerializer_ScheduleAndWait_Closed(t *testing.T) {
+	serializerCtx, serializerCancel := context.WithCancel(context.Background())
+	cs := NewCallbackSerializer(serializerCtx)
+
+	// Close the serializer and wait for it to finish shutting down.
+	serializerCancel()
+	<-cs.Done()
+
+	if err := cs.ScheduleAndWait(func(context.Context) {
+		t.Fatal("Callback executed on a closed serializer")
+	}); err != ErrSerializerClosed {
+		t.Fatalf("ScheduleAndWait() returned error: %v, want %v", err, ErrSerializerClosed)
+	}
+}

internal/grpcsync/event.go

@@ -0,0 +1,58 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+// Package grpcsync implements additional synchronization primitives built upon
+// the sync package.
+package grpcsync
+
+import (
+	"sync/atomic"
+)
+
+// Event represents a one-time event that may occur in the future.
+type Event struct {
+	fired atomic.Bool
+	c     chan struct{}
+}
+
+// Fire causes e to complete.  It is safe to call multiple times, and
+// concurrently.  It returns true iff this call to Fire caused the signaling
+// channel returned by Done to close. If Fire returns false, it is possible
+// the Done channel has not been closed yet.
+func (e *Event) Fire() bool {
+	if e.fired.CompareAndSwap(false, true) {
+		close(e.c)
+		return true
+	}
+	return false
+}
+
+// Done returns a channel that will be closed when Fire is called.
+func (e *Event) Done() <-chan struct{} {
+	return e.c
+}
+
+// HasFired returns true if Fire has been called.
+func (e *Event) HasFired() bool {
+	return e.fired.Load()
+}
+
+// NewEvent returns a new, ready-to-use Event.
+func NewEvent() *Event {
+	return &Event{c: make(chan struct{})}
+}

internal/grpcsync/event_test.go

@@ -0,0 +1,81 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package grpcsync
+
+import (
+	"testing"
+
+	"google.golang.org/grpc/internal/grpctest"
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+func (s) TestEventHasFired(t *testing.T) {
+	e := NewEvent()
+	if e.HasFired() {
+		t.Fatal("e.HasFired() = true; want false")
+	}
+	if !e.Fire() {
+		t.Fatal("e.Fire() = false; want true")
+	}
+	if !e.HasFired() {
+		t.Fatal("e.HasFired() = false; want true")
+	}
+}
+
+func (s) TestEventDoneChannel(t *testing.T) {
+	e := NewEvent()
+	select {
+	case <-e.Done():
+		t.Fatal("e.HasFired() = true; want false")
+	default:
+	}
+	if !e.Fire() {
+		t.Fatal("e.Fire() = false; want true")
+	}
+	select {
+	case <-e.Done():
+	default:
+		t.Fatal("e.HasFired() = false; want true")
+	}
+}
+
+func (s) TestEventMultipleFires(t *testing.T) {
+	e := NewEvent()
+	if e.HasFired() {
+		t.Fatal("e.HasFired() = true; want false")
+	}
+	if !e.Fire() {
+		t.Fatal("e.Fire() = false; want true")
+	}
+	for i := 0; i < 3; i++ {
+		if !e.HasFired() {
+			t.Fatal("e.HasFired() = false; want true")
+		}
+		if e.Fire() {
+			t.Fatal("e.Fire() = true; want false")
+		}
+	}
+}

internal/grpcsync/pubsub.go

@@ -0,0 +1,121 @@
+/*
+ *
+ * Copyright 2023 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package grpcsync
+
+import (
+	"context"
+	"sync"
+)
+
+// Subscriber represents an entity that is subscribed to messages published on
+// a PubSub. It wraps the callback to be invoked by the PubSub when a new
+// message is published.
+type Subscriber interface {
+	// OnMessage is invoked when a new message is published. Implementations
+	// must not block in this method.
+	OnMessage(msg any)
+}
+
+// PubSub is a simple one-to-many publish-subscribe system that supports
+// messages of arbitrary type. It guarantees that messages are delivered in
+// the same order in which they were published.
+//
+// Publisher invokes the Publish() method to publish new messages, while
+// subscribers interested in receiving these messages register a callback
+// via the Subscribe() method.
+//
+// Once a PubSub is stopped, no more messages can be published, but any pending
+// published messages will be delivered to the subscribers.  Done may be used
+// to determine when all published messages have been delivered.
+type PubSub struct {
+	cs *CallbackSerializer
+
+	// Access to the below fields are guarded by this mutex.
+	mu          sync.Mutex
+	msg         any
+	subscribers map[Subscriber]bool
+}
+
+// NewPubSub returns a new PubSub instance.  Users should cancel the
+// provided context to shutdown the PubSub.
+func NewPubSub(ctx context.Context) *PubSub {
+	return &PubSub{
+		cs:          NewCallbackSerializer(ctx),
+		subscribers: map[Subscriber]bool{},
+	}
+}
+
+// Subscribe registers the provided Subscriber to the PubSub.
+//
+// If the PubSub contains a previously published message, the Subscriber's
+// OnMessage() callback will be invoked asynchronously with the existing
+// message to begin with, and subsequently for every newly published message.
+//
+// The caller is responsible for invoking the returned cancel function to
+// unsubscribe itself from the PubSub.
+func (ps *PubSub) Subscribe(sub Subscriber) (cancel func()) {
+	ps.mu.Lock()
+	defer ps.mu.Unlock()
+
+	ps.subscribers[sub] = true
+
+	if ps.msg != nil {
+		msg := ps.msg
+		ps.cs.TrySchedule(func(context.Context) {
+			ps.mu.Lock()
+			defer ps.mu.Unlock()
+			if !ps.subscribers[sub] {
+				return
+			}
+			sub.OnMessage(msg)
+		})
+	}
+
+	return func() {
+		ps.mu.Lock()
+		defer ps.mu.Unlock()
+		delete(ps.subscribers, sub)
+	}
+}
+
+// Publish publishes the provided message to the PubSub, and invokes
+// callbacks registered by subscribers asynchronously.
+func (ps *PubSub) Publish(msg any) {
+	ps.mu.Lock()
+	defer ps.mu.Unlock()
+
+	ps.msg = msg
+	for sub := range ps.subscribers {
+		s := sub
+		ps.cs.TrySchedule(func(context.Context) {
+			ps.mu.Lock()
+			defer ps.mu.Unlock()
+			if !ps.subscribers[s] {
+				return
+			}
+			s.OnMessage(msg)
+		})
+	}
+}
+
+// Done returns a channel that is closed after the context passed to NewPubSub
+// is canceled and all updates have been sent to subscribers.
+func (ps *PubSub) Done() <-chan struct{} {
+	return ps.cs.Done()
+}

internal/grpcsync/pubsub_test.go

@@ -0,0 +1,196 @@
+/*
+ *
+ * Copyright 2023 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package grpcsync
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+)
+
+type testSubscriber struct {
+	onMsgCh chan int
+}
+
+func newTestSubscriber(chSize int) *testSubscriber {
+	return &testSubscriber{onMsgCh: make(chan int, chSize)}
+}
+
+func (ts *testSubscriber) OnMessage(msg any) {
+	select {
+	case ts.onMsgCh <- msg.(int):
+	default:
+	}
+}
+
+func (s) TestPubSub_PublishNoMsg(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	pubsub := NewPubSub(ctx)
+
+	ts := newTestSubscriber(1)
+	pubsub.Subscribe(ts)
+
+	select {
+	case <-ts.onMsgCh:
+		t.Fatal("Subscriber callback invoked when no message was published")
+	case <-time.After(defaultTestShortTimeout):
+	}
+}
+
+func (s) TestPubSub_PublishMsgs_RegisterSubs_And_Stop(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	pubsub := NewPubSub(ctx)
+
+	const numPublished = 10
+
+	ts1 := newTestSubscriber(numPublished)
+	pubsub.Subscribe(ts1)
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+	// Publish ten messages on the pubsub and ensure that they are received in order by the subscriber.
+	go func() {
+		for i := 0; i < numPublished; i++ {
+			pubsub.Publish(i)
+		}
+		wg.Done()
+	}()
+
+	go func() {
+		defer wg.Done()
+		for i := 0; i < numPublished; i++ {
+			select {
+			case m := <-ts1.onMsgCh:
+				if m != i {
+					t.Errorf("Received unexpected message: %q; want: %q", m, i)
+					return
+				}
+			case <-time.After(defaultTestTimeout):
+				t.Error("Timeout when expecting the onMessage() callback to be invoked")
+				return
+			}
+		}
+	}()
+	wg.Wait()
+	if t.Failed() {
+		t.FailNow()
+	}
+
+	// Register another subscriber and ensure that it receives the last published message.
+	ts2 := newTestSubscriber(numPublished)
+	pubsub.Subscribe(ts2)
+
+	select {
+	case m := <-ts2.onMsgCh:
+		if m != numPublished-1 {
+			t.Fatalf("Received unexpected message: %q; want: %q", m, numPublished-1)
+		}
+	case <-time.After(defaultTestShortTimeout):
+		t.Fatal("Timeout when expecting the onMessage() callback to be invoked")
+	}
+
+	wg.Add(3)
+	// Publish ten messages on the pubsub and ensure that they are received in order by the subscribers.
+	go func() {
+		for i := 0; i < numPublished; i++ {
+			pubsub.Publish(i)
+		}
+		wg.Done()
+	}()
+	go func() {
+		defer wg.Done()
+		for i := 0; i < numPublished; i++ {
+			select {
+			case m := <-ts1.onMsgCh:
+				if m != i {
+					t.Errorf("Received unexpected message: %q; want: %q", m, i)
+					return
+				}
+			case <-time.After(defaultTestTimeout):
+				t.Error("Timeout when expecting the onMessage() callback to be invoked")
+				return
+			}
+		}
+
+	}()
+	go func() {
+		defer wg.Done()
+		for i := 0; i < numPublished; i++ {
+			select {
+			case m := <-ts2.onMsgCh:
+				if m != i {
+					t.Errorf("Received unexpected message: %q; want: %q", m, i)
+					return
+				}
+			case <-time.After(defaultTestTimeout):
+				t.Error("Timeout when expecting the onMessage() callback to be invoked")
+				return
+			}
+		}
+	}()
+	wg.Wait()
+	if t.Failed() {
+		t.FailNow()
+	}
+
+	cancel()
+	<-pubsub.Done()
+
+	go func() {
+		pubsub.Publish(99)
+	}()
+	// Ensure that the subscriber callback is not invoked as instantiated
+	// pubsub has already closed.
+	select {
+	case <-ts1.onMsgCh:
+		t.Fatal("The callback was invoked after pubsub being stopped")
+	case <-ts2.onMsgCh:
+		t.Fatal("The callback was invoked after pubsub being stopped")
+	case <-time.After(defaultTestShortTimeout):
+	}
+}
+
+func (s) TestPubSub_PublishMsgs_BeforeRegisterSub(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	pubsub := NewPubSub(ctx)
+
+	const numPublished = 3
+	for i := 0; i < numPublished; i++ {
+		pubsub.Publish(i)
+	}
+
+	ts := newTestSubscriber(numPublished)
+	pubsub.Subscribe(ts)
+
+	// Ensure that the subscriber callback is invoked with a previously
+	// published message.
+	select {
+	case d := <-ts.onMsgCh:
+		if d != numPublished-1 {
+			t.Fatalf("Unexpected message received: %q; %q", d, numPublished-1)
+		}
+
+	case <-time.After(defaultTestShortTimeout):
+		t.Fatal("Timeout when expecting the onMessage() callback to be invoked")
+	}
+}

internal/grpcsync/refcounted.go

@@ -0,0 +1,109 @@
+/*
+ *
+ * Copyright 2026 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package grpcsync
+
+import (
+	"sync/atomic"
+
+	"google.golang.org/grpc/grpclog"
+)
+
+var logger = grpclog.Component("grpcsync")
+
+// RefCounted is a reference counted wrapper of type T. It tracks the number of
+// active references and runs a cleanup when the last reference is released.
+type RefCounted[T any] struct {
+	val      T
+	refCount atomic.Int32
+	onZero   func()
+}
+
+// NewRefCounted creates a new RefCounted instance wrapping the given value with
+// initial refcount of one.
+//
+// The value should typically be a pointer, interface, or handle rather than a
+// plain value type (such as a struct or primitive value).
+//
+// The provided onZero callback must not be nil, and is executed exactly once
+// when the reference count drops to zero.  Panics if onZero is nil.
+//
+// WARNING: onZero runs synchronously inside Decrement; it must not acquire
+// locks held by Decrement callers.
+func NewRefCounted[T any](val T, onZero func()) *RefCounted[T] {
+	if onZero == nil {
+		panic("grpcsync: onZero callback cannot be nil")
+	}
+	rc := &RefCounted[T]{
+		val:    val,
+		onZero: onZero,
+	}
+	rc.refCount.Store(1)
+	return rc
+}
+
+// Value returns the encapsulated resource.
+func (rc *RefCounted[T]) Value() T {
+	return rc.val
+}
+
+// TryIncrement attempts to increment the reference count, returning true if
+// successful. It returns false if the count has already reached 0, indicating
+// the resource has been cleaned up and cannot be resurrected.
+//
+// WARNING: Avoid calling TryIncrement on hot paths when an active reference is
+// already guaranteed by the caller; use Increment instead to bypass the
+// CompareAndSwap loop overhead. TryIncrement should be reserved for speculative
+// lookups (such as fetching from a cache or map) where the resource might be
+// dead.
+func (rc *RefCounted[T]) TryIncrement() bool {
+	// Utilize a CompareAndSwap loop to prevent race conditions where a concurrent
+	// decrement could drop the count to zero between the read and the increment
+	// operation, which would otherwise inadvertently resurrect a closed resource.
+	for {
+		count := rc.refCount.Load()
+		if count <= 0 {
+			return false // Already dead or dying
+		}
+		if rc.refCount.CompareAndSwap(count, count+1) {
+			return true
+		}
+	}
+}
+
+// Increment increments the reference count.
+//
+// WARNING: Call Increment only when there is a guarantee that an active
+// reference is already present, ensuring the resource is not dead. If there is
+// a possibility that the resource is dead or its reference count might have
+// reached zero, call TryIncrement instead.
+func (rc *RefCounted[T]) Increment() {
+	if rc.refCount.Add(1) <= 1 {
+		logger.Errorf("Resource already closed or dead")
+	}
+}
+
+// Decrement decrements the reference count. If it drops to zero, the onZero
+// callback is executed synchronously before this method returns.
+func (rc *RefCounted[T]) Decrement() {
+	if v := rc.refCount.Add(-1); v < 0 {
+		logger.Errorf("Refcount cannot be negative")
+	} else if v == 0 {
+		rc.onZero()
+	}
+}

internal/grpcsync/refcounted_test.go

@@ -0,0 +1,155 @@
+/*
+ *
+ * Copyright 2026 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package grpcsync
+
+import (
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"google.golang.org/grpc/internal/grpctest"
+)
+
+// Test verifies the scenario where a RefCounted instance is created and its
+// reference count is decremented. It verifies that the encapsulated value is
+// correctly retrieved and that the onZero callback is invoked when the
+// reference count drops to zero.
+func (s) TestRefCounted(t *testing.T) {
+	val := 666
+	var onZeroCalled atomic.Bool
+	rc := NewRefCounted(&val, func() { onZeroCalled.Store(true) })
+	if got := rc.Value(); *got != val {
+		t.Fatalf("Value() = %v, want %v", *got, val)
+	}
+
+	if got := onZeroCalled.Load(); got {
+		t.Fatalf("Before decrement, onZeroCalled = %v, want false", got)
+	}
+
+	rc.Decrement()
+
+	if got := onZeroCalled.Load(); !got {
+		t.Fatalf("After decrementing count to zero, onZeroCalled = %v, want true", got)
+	}
+}
+
+// Test verifies the scenario where the reference count of a resource is
+// explicitly incremented and decremented multiple times. It verifies that the
+// onZero callback is only executed when the count drops to zero, and not
+// beforehand.
+func (s) TestRefCounted_IncrementDecrement(t *testing.T) {
+	val := 666
+	var onZeroCount atomic.Int32
+	rc := NewRefCounted(&val, func() { onZeroCount.Add(1) })
+	rc.Increment()
+	rc.Decrement()
+
+	if got := onZeroCount.Load(); got != 0 {
+		t.Fatal("OnZero called unexpectedly after decrementing refcount from 2 to 1")
+	}
+
+	rc.Decrement()
+
+	if got := onZeroCount.Load(); got != 1 {
+		t.Fatalf("After decrementing count from 1 to 0, onZero called = %v times, want 1", got)
+	}
+}
+
+// Test verifies the scenario where TryIncrement is called on both an active and
+// a dead resource. It verifies that TryIncrement successfully increments active
+// resources but fails on resources whose reference count has already dropped to
+// zero.
+func (s) TestRefCounted_TryIncrement(t *testing.T) {
+	val := 666
+	var onZeroCount atomic.Int32
+	rc := NewRefCounted(&val, func() { onZeroCount.Add(1) })
+	if got := rc.TryIncrement(); !got {
+		t.Fatalf("TryIncrement() on active resource = %v, want true", got)
+	}
+
+	rc.Decrement()
+	rc.Decrement()
+
+	if got := onZeroCount.Load(); got != 1 {
+		t.Fatalf("After decrementing count to zero, onZero called = %v times, want 1", got)
+	}
+
+	if got := rc.TryIncrement(); got {
+		t.Fatalf("TryIncrement() on dead resource = %v, want false", got)
+	}
+}
+
+// Test verifies the scenario where multiple goroutines concurrently increment
+// and decrement the reference count. It verifies that the reference counting is
+// thread-safe and the onZero callback is invoked exactly once.
+func (s) TestRefCounted_Concurrent(t *testing.T) {
+	val := 666
+	var onZeroCount atomic.Int32
+	rc := NewRefCounted(&val, func() { onZeroCount.Add(1) })
+
+	var wg sync.WaitGroup
+	const numGoroutines = 100
+	for i := 0; i < numGoroutines; i++ {
+		wg.Go(func() {
+			if rc.TryIncrement() {
+				rc.Decrement()
+			}
+		})
+	}
+	rc.Decrement()
+	wg.Wait()
+	if got := onZeroCount.Load(); got != 1 {
+		t.Fatalf("After concurrent increments/decrements and final decrement, onZeroCount = %v, want 1", got)
+	}
+}
+
+// Test verifies that Decrementing a resource whose reference count has already
+// dropped to zero results in a negative refcount log error.
+func (s) TestRefCounted_DecrementNegative(_ *testing.T) {
+	val := 666
+	rc := NewRefCounted(&val, func() {})
+	rc.Decrement()
+
+	grpctest.ExpectError("Refcount cannot be negative")
+	rc.Decrement()
+}
+
+// Test verifies that Incrementing a dead resource results in a closed resource
+// log error.
+func (s) TestRefCounted_IncrementDead(_ *testing.T) {
+	val := 666
+	rc := NewRefCounted(&val, func() {})
+	rc.Decrement()
+
+	grpctest.ExpectError("Resource already closed or dead")
+	rc.Increment()
+}
+
+// Test verifies that NewRefCounted returns an error if the provided onZero
+// callback is nil.
+func (s) TestNilOnZero(t *testing.T) {
+	const wantErr = "grpcsync: onZero callback cannot be nil"
+	defer func() {
+		if r := recover(); r == nil || r != wantErr {
+			t.Fatalf("NewRefCounted(_, nil) panic = %v, want: %q", r, wantErr)
+		}
+	}()
+	val := 666
+	NewRefCounted(&val, nil)
+}