go-redis/internal/semaphore.go

package internal

import (
	"context"
	"sync"
	"sync/atomic"
	"time"
)

var semTimers = sync.Pool{
	New: func() interface{} {
		t := time.NewTimer(time.Hour)
		t.Stop()
		return t
	},
}

// FastSemaphore is a counting semaphore implementation using atomic operations.
// It's optimized for the fast path (no blocking) while still supporting timeouts and context cancellation.
//
// Performance characteristics:
// - Fast path (no blocking): Single atomic CAS operation
// - Slow path (blocking): Falls back to channel-based waiting
// - Release: Single atomic decrement + optional channel notification
//
// This is significantly faster than a pure channel-based semaphore because:
// 1. The fast path avoids channel operations entirely (no scheduler involvement)
// 2. Atomic operations are much cheaper than channel send/receive
type FastSemaphore struct {
	// Current number of acquired tokens (atomic)
	count atomic.Int32

	// Maximum number of tokens (capacity)
	max int32

	// Channel for blocking waiters
	// Only used when fast path fails (semaphore is full)
	waitCh chan struct{}
}

// NewFastSemaphore creates a new fast semaphore with the given capacity.
func NewFastSemaphore(capacity int32) *FastSemaphore {
	return &FastSemaphore{
		max:    capacity,
		waitCh: make(chan struct{}, capacity),
	}
}

// TryAcquire attempts to acquire a token without blocking.
// Returns true if successful, false if the semaphore is full.
//
// This is the fast path - just a single CAS operation.
func (s *FastSemaphore) TryAcquire() bool {
	for {
		current := s.count.Load()
		if current >= s.max {
			return false // Semaphore is full
		}
		if s.count.CompareAndSwap(current, current+1) {
			return true // Successfully acquired
		}
		// CAS failed due to concurrent modification, retry
	}
}

// Acquire acquires a token, blocking if necessary until one is available or the context is cancelled.
// Returns an error if the context is cancelled or the timeout expires.
// Returns timeoutErr when the timeout expires.
//
// Performance optimization:
// 1. First try fast path (no blocking)
// 2. If that fails, fall back to channel-based waiting
func (s *FastSemaphore) Acquire(ctx context.Context, timeout time.Duration, timeoutErr error) error {
	// Fast path: try to acquire without blocking
	select {
	case <-ctx.Done():
		return ctx.Err()
	default:
	}

	// Try fast acquire first
	if s.TryAcquire() {
		return nil
	}

	// Fast path failed, need to wait
	// Use timer pool to avoid allocation
	timer := semTimers.Get().(*time.Timer)
	defer semTimers.Put(timer)
	timer.Reset(timeout)

	start := time.Now()

	for {
		select {
		case <-ctx.Done():
			if !timer.Stop() {
				<-timer.C
			}
			return ctx.Err()

		case <-s.waitCh:
			// Someone released a token, try to acquire it
			if s.TryAcquire() {
				if !timer.Stop() {
					<-timer.C
				}
				return nil
			}
			// Failed to acquire (race with another goroutine), continue waiting

		case <-timer.C:
			return timeoutErr
		}

		// Periodically check if we can acquire (handles race conditions)
		if time.Since(start) > timeout {
			return timeoutErr
		}
	}
}

// AcquireBlocking acquires a token, blocking indefinitely until one is available.
// This is useful for cases where you don't need timeout or context cancellation.
// Returns immediately if a token is available (fast path).
func (s *FastSemaphore) AcquireBlocking() {
	// Try fast path first
	if s.TryAcquire() {
		return
	}

	// Slow path: wait for a token
	for {
		<-s.waitCh
		if s.TryAcquire() {
			return
		}
		// Failed to acquire (race with another goroutine), continue waiting
	}
}

// Release releases a token back to the semaphore.
// This wakes up one waiting goroutine if any are blocked.
func (s *FastSemaphore) Release() {
	s.count.Add(-1)

	// Try to wake up a waiter (non-blocking)
	// If no one is waiting, this is a no-op
	select {
	case s.waitCh <- struct{}{}:
		// Successfully notified a waiter
	default:
		// No waiters, that's fine
	}
}

// Len returns the current number of acquired tokens.
// Used by tests to check semaphore state.
func (s *FastSemaphore) Len() int32 {
	return s.count.Load()
}