internal: adds utility for refcounting and cleanup

internal/grpcsync/refcounted.go

@@ -0,0 +1,106 @@
+/*
+ *
+ * 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 (
+	"fmt"
+	"sync/atomic"
+)
+
+// RefCounted tracks how many consumers hold a value of type V and runs a
+// cleanup when the last reference is released
+type RefCounted[V any] struct {
+	val      V
+	refCount atomic.Int32
+	onZero   func()
+}
+
+// NewRefCounted creates a new RefCounted instance wrapping the given value with
+// initial refcount of one. The provided onZero callback must not be nil, and is
+// executed exactly once when the reference count drops to zero.
+//
+// The value should typically be a pointer, interface, or handle rather than a
+// plain value type (such as a struct or primitive value).
+//
+// WARNING: onZero runs synchronously inside Decrement; it must not acquire
+// locks held by Decrement callers or invert lock ordering.
+func NewRefCounted[V any](val V, onZero func()) (*RefCounted[V], error) {
+	if onZero == nil {
+		return nil, fmt.Errorf("grpcsync: onZero callback cannot be nil")
+	}
+	rc := &RefCounted[V]{
+		val:    val,
+		onZero: onZero,
+	}
+	rc.refCount.Store(1)
+	return rc, nil
+}
+
+// Value returns the encapsulated resource.
+func (rc *RefCounted[V]) Value() V {
+	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[V]) 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[V]) Increment() error {
+	if rc.refCount.Add(1) <= 1 {
+		return fmt.Errorf("grpcsync: resource already closed or dead")
+	}
+	return nil
+}
+
+// Decrement decrements the reference count. If it drops to zero, the onZero
+// callback is executed synchronously before this method returns.
+func (rc *RefCounted[V]) Decrement() error {
+	if v := rc.refCount.Add(-1); v < 0 {
+		return fmt.Errorf("grpcsync: refcount cannot be negative")
+	} else if v == 0 {
+		rc.onZero()
+	}
+	return nil
+}

internal/grpcsync/refcounted_test.go

@@ -0,0 +1,167 @@
+/*
+ *
+ * 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"
+)
+
+// 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 := "test-value"
+	var onZeroCalled atomic.Bool
+
+	rc, err := NewRefCounted(val, func() {
+		onZeroCalled.Store(true)
+	})
+	if err != nil {
+		t.Fatalf("NewRefCounted() failed: %v", err)
+	}
+
+	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 := 42
+	var onZeroCount atomic.Int32
+
+	rc, err := NewRefCounted(val, func() {
+		onZeroCount.Add(1)
+	})
+	if err != nil {
+		t.Fatalf("NewRefCounted() failed: %v", err)
+	}
+
+	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) {
+	var onZeroCount atomic.Int32
+	rc, err := NewRefCounted("val", func() {
+		onZeroCount.Add(1)
+	})
+	if err != nil {
+		t.Fatalf("NewRefCounted() failed: %v", err)
+	}
+
+	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. It also verifies
+// that once onZero has been called, subsequent increament/decrement return
+// error.
+func (s) TestRefCounted_Concurrent(t *testing.T) {
+	const numGoroutines = 10
+	var onZeroCount atomic.Int32
+
+	rc, err := NewRefCounted("concurrent-val", func() {
+		onZeroCount.Add(1)
+	})
+	if err != nil {
+		t.Fatalf("NewRefCounted() failed: %v", err)
+	}
+
+	var wg sync.WaitGroup
+	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)
+	}
+
+	wantErr := "grpcsync: refcount cannot be negative"
+	if err := rc.Decrement(); err == nil || err.Error() != wantErr {
+		t.Fatalf("Decrement() on dead resource = %v, want %q", err, wantErr)
+	}
+
+	wantErr = "grpcsync: resource already closed or dead"
+	if err := rc.Increment(); err == nil || err.Error() != wantErr {
+		t.Fatalf("Increment() on dead resource = %v, want %q", err, wantErr)
+	}
+
+	if got := rc.TryIncrement(); got {
+		t.Fatalf("TryIncrement() on dead resource = %v, want false", got)
+	}
+}
+
+// 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"
+	if _, err := NewRefCounted("val", nil); err == nil || err.Error() != wantErr {
+		t.Fatalf("NewRefCounted(_, nil) returned error: %v, want: %q", err, wantErr)
+	}
+}