diff --git a/internal/grpcsync/refcounted.go b/internal/grpcsync/refcounted.go new file mode 100644 index 000000000000..2affcac7a8ae --- /dev/null +++ b/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 +} diff --git a/internal/grpcsync/refcounted_test.go b/internal/grpcsync/refcounted_test.go new file mode 100644 index 000000000000..0c9c92b73af5 --- /dev/null +++ b/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) + } +}