rewrite and fix reference counter

internal/utils/ref_count.go

@@ -1,42 +1,41 @@
 package utils
 
+import (
+	"sync"
+	"sync/atomic"
+)
+
 type RefCount struct {
 	_ NoCopy
 
-	refCh      chan bool
+	mu       sync.Mutex
-	notifyZero chan struct{}
+	cond     *sync.Cond
+	refCount uint32
+	zeroCh   chan struct{}
 }
 
 func NewRefCounter() *RefCount {
 	rc := &RefCount{
-		refCh:      make(chan bool, 1),
+		refCount: 1,
-		notifyZero: make(chan struct{}),
+		zeroCh:   make(chan struct{}),
 	}
-	go func() {
+	rc.cond = sync.NewCond(&rc.mu)
-		refCount := uint32(1)
-		for isAdd := range rc.refCh {
-			if isAdd {
-				refCount++
-			} else {
-				refCount--
-			}
-			if refCount <= 0 {
-				close(rc.notifyZero)
-				return
-			}
-		}
-	}()
 	return rc
 }
 
 func (rc *RefCount) Zero() <-chan struct{} {
-	return rc.notifyZero
+	return rc.zeroCh
 }
 
 func (rc *RefCount) Add() {
-	rc.refCh <- true
+	atomic.AddUint32(&rc.refCount, 1)
 }
 
 func (rc *RefCount) Sub() {
-	rc.refCh <- false
+	if atomic.AddUint32(&rc.refCount, ^uint32(0)) == 0 {
+		rc.mu.Lock()
+		close(rc.zeroCh)
+		rc.cond.Broadcast()
+		rc.mu.Unlock()
+	}
 }

internal/utils/ref_count_test.go

@@ -0,0 +1,77 @@
+package utils
+
+import (
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestRefCounter_AddSub(t *testing.T) {
+	rc := NewRefCounter()
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	go func() {
+		defer wg.Done()
+		rc.Add()
+	}()
+
+	go func() {
+		defer wg.Done()
+		rc.Sub()
+	}()
+
+	wg.Wait()
+
+	select {
+	case <-rc.Zero():
+		// Expected behavior
+	case <-time.After(1 * time.Second):
+		t.Fatal("Expected Zero channel to close, but it didn't")
+	}
+}
+
+func TestRefCounter_MultipleAddSub(t *testing.T) {
+	rc := NewRefCounter()
+
+	var wg sync.WaitGroup
+	numAdds := 5
+	numSubs := 5
+	wg.Add(numAdds + numSubs)
+
+	for range numAdds {
+		go func() {
+			defer wg.Done()
+			rc.Add()
+		}()
+	}
+
+	for range numSubs {
+		go func() {
+			defer wg.Done()
+			rc.Sub()
+		}()
+	}
+
+	wg.Wait()
+
+	select {
+	case <-rc.Zero():
+		// Expected behavior
+	case <-time.After(1 * time.Second):
+		t.Fatal("Expected Zero channel to close, but it didn't")
+	}
+}
+
+func TestRefCounter_ZeroInitially(t *testing.T) {
+	rc := NewRefCounter()
+	rc.Sub() // Bring count to zero
+
+	select {
+	case <-rc.Zero():
+		// Expected behavior
+	case <-time.After(1 * time.Second):
+		t.Fatal("Expected Zero channel to close, but it didn't")
+	}
+}