fix(go.sum): update ResolveSpec dependency to v1.0.87
This commit is contained in:
+990
@@ -0,0 +1,990 @@
|
||||
// Package pool implements the pool management
|
||||
package pool
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9/internal"
|
||||
"github.com/redis/go-redis/v9/internal/maintnotifications/logs"
|
||||
"github.com/redis/go-redis/v9/internal/proto"
|
||||
uberatomic "go.uber.org/atomic"
|
||||
)
|
||||
|
||||
var noDeadline = time.Time{}
|
||||
|
||||
// Preallocated errors for hot paths to avoid allocations
|
||||
var (
|
||||
errAlreadyMarkedForHandoff = errors.New("connection is already marked for handoff")
|
||||
errNotMarkedForHandoff = errors.New("connection was not marked for handoff")
|
||||
errHandoffStateChanged = errors.New("handoff state changed during marking")
|
||||
errConnectionNotAvailable = errors.New("redis: connection not available")
|
||||
errConnNotAvailableForWrite = errors.New("redis: connection not available for write operation")
|
||||
)
|
||||
|
||||
// getCachedTimeNs returns the current time in nanoseconds.
|
||||
// This function previously used a global cache updated by a background goroutine,
|
||||
// but that caused unnecessary CPU usage when the client was idle (ticker waking up
|
||||
// the scheduler every 50ms). We now use time.Now() directly, which is fast enough
|
||||
// on modern systems (vDSO on Linux) and only adds ~1-2% overhead in extreme
|
||||
// high-concurrency benchmarks while eliminating idle CPU usage.
|
||||
func getCachedTimeNs() int64 {
|
||||
return time.Now().UnixNano()
|
||||
}
|
||||
|
||||
// GetCachedTimeNs returns the current time in nanoseconds.
|
||||
// Exported for use by other packages that need fast time access.
|
||||
func GetCachedTimeNs() int64 {
|
||||
return getCachedTimeNs()
|
||||
}
|
||||
|
||||
// Global atomic counter for connection IDs
|
||||
var connIDCounter uint64
|
||||
|
||||
// HandoffState represents the atomic state for connection handoffs
|
||||
// This struct is stored atomically to prevent race conditions between
|
||||
// checking handoff status and reading handoff parameters
|
||||
type HandoffState struct {
|
||||
ShouldHandoff bool // Whether connection should be handed off
|
||||
Endpoint string // New endpoint for handoff
|
||||
SeqID int64 // Sequence ID from MOVING notification
|
||||
}
|
||||
|
||||
// atomicNetConn is a wrapper to ensure consistent typing in atomic.Value
|
||||
type atomicNetConn struct {
|
||||
conn net.Conn
|
||||
}
|
||||
|
||||
// generateConnID generates a fast unique identifier for a connection with zero allocations
|
||||
func generateConnID() uint64 {
|
||||
return atomic.AddUint64(&connIDCounter, 1)
|
||||
}
|
||||
|
||||
type Conn struct {
|
||||
// Connection identifier for unique tracking
|
||||
id uint64
|
||||
|
||||
usedAt atomic.Int64
|
||||
lastPutAt atomic.Int64
|
||||
dialStartNs atomic.Int64 // Time when dial started (for connection create time metric)
|
||||
|
||||
// Lock-free netConn access using atomic.Value
|
||||
// Contains *atomicNetConn wrapper, accessed atomically for better performance
|
||||
netConnAtomic atomic.Value // stores *atomicNetConn
|
||||
|
||||
rd *proto.Reader
|
||||
bw *bufio.Writer
|
||||
wr *proto.Writer
|
||||
|
||||
// Lightweight mutex to protect reader operations during handoff
|
||||
// Only used for the brief period during SetNetConn and HasBufferedData/PeekReplyTypeSafe
|
||||
readerMu sync.RWMutex
|
||||
|
||||
// State machine for connection state management
|
||||
// Replaces: usable, Inited, used
|
||||
// Provides thread-safe state transitions with FIFO waiting queue
|
||||
// States: CREATED → INITIALIZING → IDLE ⇄ IN_USE
|
||||
// ↓
|
||||
// UNUSABLE (handoff/reauth)
|
||||
// ↓
|
||||
// IDLE/CLOSED
|
||||
stateMachine *ConnStateMachine
|
||||
|
||||
// Handoff metadata - managed separately from state machine
|
||||
// These are atomic for lock-free access during handoff operations
|
||||
handoffStateAtomic atomic.Value // stores *HandoffState
|
||||
handoffRetriesAtomic atomic.Uint32 // retry counter
|
||||
|
||||
pooled bool
|
||||
pubsub bool
|
||||
createdAt time.Time
|
||||
expiresAt time.Time
|
||||
poolName string // Name of the pool this connection belongs to (for metrics)
|
||||
|
||||
// When a goroutine closes a connection, it usually knows the reason, so closeReason is not needed.
|
||||
// closeReason is only used when an in-use connection is closed by another goroutine,
|
||||
// to inform the goroutine using the connection why the connection was closed.
|
||||
closeReason uberatomic.String
|
||||
|
||||
// maintenanceNotifications upgrade support: relaxed timeouts during migrations/failovers
|
||||
|
||||
// Using atomic operations for lock-free access to avoid mutex contention
|
||||
relaxedReadTimeoutNs atomic.Int64 // time.Duration as nanoseconds
|
||||
relaxedWriteTimeoutNs atomic.Int64 // time.Duration as nanoseconds
|
||||
relaxedDeadlineNs atomic.Int64 // time.Time as nanoseconds since epoch
|
||||
|
||||
// Counter to track multiple relaxed timeout setters if we have nested calls
|
||||
// will be decremented when ClearRelaxedTimeout is called or deadline is reached
|
||||
// if counter reaches 0, we clear the relaxed timeouts
|
||||
relaxedCounter atomic.Int32
|
||||
|
||||
// Connection initialization function for reconnections
|
||||
initConnFunc func(context.Context, *Conn) error
|
||||
|
||||
onClose func() error
|
||||
}
|
||||
|
||||
func NewConn(netConn net.Conn) *Conn {
|
||||
return NewConnWithBufferSize(netConn, proto.DefaultBufferSize, proto.DefaultBufferSize)
|
||||
}
|
||||
|
||||
func NewConnWithBufferSize(netConn net.Conn, readBufSize, writeBufSize int) *Conn {
|
||||
now := time.Now()
|
||||
cn := &Conn{
|
||||
createdAt: now,
|
||||
id: generateConnID(), // Generate unique ID for this connection
|
||||
stateMachine: NewConnStateMachine(),
|
||||
}
|
||||
|
||||
// Use specified buffer sizes, or fall back to 32KiB defaults if 0
|
||||
if readBufSize > 0 {
|
||||
cn.rd = proto.NewReaderSize(netConn, readBufSize)
|
||||
} else {
|
||||
cn.rd = proto.NewReader(netConn) // Uses 32KiB default
|
||||
}
|
||||
|
||||
if writeBufSize > 0 {
|
||||
cn.bw = bufio.NewWriterSize(netConn, writeBufSize)
|
||||
} else {
|
||||
cn.bw = bufio.NewWriterSize(netConn, proto.DefaultBufferSize)
|
||||
}
|
||||
|
||||
// Store netConn atomically for lock-free access using wrapper
|
||||
cn.netConnAtomic.Store(&atomicNetConn{conn: netConn})
|
||||
|
||||
cn.wr = proto.NewWriter(cn.bw)
|
||||
cn.SetUsedAt(now)
|
||||
// Initialize handoff state atomically
|
||||
initialHandoffState := &HandoffState{
|
||||
ShouldHandoff: false,
|
||||
Endpoint: "",
|
||||
SeqID: 0,
|
||||
}
|
||||
cn.handoffStateAtomic.Store(initialHandoffState)
|
||||
return cn
|
||||
}
|
||||
|
||||
func (cn *Conn) UsedAt() time.Time {
|
||||
return time.Unix(0, cn.usedAt.Load())
|
||||
}
|
||||
func (cn *Conn) SetUsedAt(tm time.Time) {
|
||||
cn.usedAt.Store(tm.UnixNano())
|
||||
}
|
||||
|
||||
func (cn *Conn) UsedAtNs() int64 {
|
||||
return cn.usedAt.Load()
|
||||
}
|
||||
func (cn *Conn) SetUsedAtNs(ns int64) {
|
||||
cn.usedAt.Store(ns)
|
||||
}
|
||||
|
||||
func (cn *Conn) LastPutAtNs() int64 {
|
||||
return cn.lastPutAt.Load()
|
||||
}
|
||||
func (cn *Conn) SetLastPutAtNs(ns int64) {
|
||||
cn.lastPutAt.Store(ns)
|
||||
}
|
||||
|
||||
// GetDialStartNs returns the time when the dial started (in nanoseconds since epoch).
|
||||
// This is used to calculate the full connection creation time (TCP + handshake).
|
||||
func (cn *Conn) GetDialStartNs() int64 {
|
||||
return cn.dialStartNs.Load()
|
||||
}
|
||||
|
||||
// PoolName returns the name of the pool this connection belongs to.
|
||||
// This is used for metrics to identify which pool a connection is from.
|
||||
func (cn *Conn) PoolName() string {
|
||||
return cn.poolName
|
||||
}
|
||||
|
||||
// SetPoolName sets the name of the pool this connection belongs to.
|
||||
// This should be called when the connection is added to a pool.
|
||||
func (cn *Conn) SetPoolName(name string) {
|
||||
cn.poolName = name
|
||||
}
|
||||
|
||||
// Backward-compatible wrapper methods for state machine
|
||||
// These maintain the existing API while using the new state machine internally
|
||||
|
||||
// CompareAndSwapUsable atomically compares and swaps the usable flag (lock-free).
|
||||
//
|
||||
// This is used by background operations (handoff, re-auth) to acquire exclusive
|
||||
// access to a connection. The operation sets usable to false, preventing the pool
|
||||
// from returning the connection to clients.
|
||||
//
|
||||
// Returns true if the swap was successful (old value matched), false otherwise.
|
||||
//
|
||||
// Implementation note: This is a compatibility wrapper around the state machine.
|
||||
// It checks if the current state is "usable" (IDLE or IN_USE) and transitions accordingly.
|
||||
// Deprecated: Use GetStateMachine().TryTransition() directly for better state management.
|
||||
func (cn *Conn) CompareAndSwapUsable(old, new bool) bool {
|
||||
currentState := cn.stateMachine.GetState()
|
||||
|
||||
// Check if current state matches the "old" usable value
|
||||
currentUsable := (currentState == StateIdle || currentState == StateInUse)
|
||||
if currentUsable != old {
|
||||
return false
|
||||
}
|
||||
|
||||
// If we're trying to set to the same value, succeed immediately
|
||||
if old == new {
|
||||
return true
|
||||
}
|
||||
|
||||
// Transition based on new value
|
||||
if new {
|
||||
// Trying to make usable - transition from UNUSABLE to IDLE
|
||||
// This should only work from UNUSABLE or INITIALIZING states
|
||||
// Use predefined slice to avoid allocation
|
||||
_, err := cn.stateMachine.TryTransition(
|
||||
validFromInitializingOrUnusable,
|
||||
StateIdle,
|
||||
)
|
||||
return err == nil
|
||||
}
|
||||
// Trying to make unusable - transition from IDLE to UNUSABLE
|
||||
// This is typically for acquiring the connection for background operations
|
||||
// Use predefined slice to avoid allocation
|
||||
_, err := cn.stateMachine.TryTransition(
|
||||
validFromIdle,
|
||||
StateUnusable,
|
||||
)
|
||||
return err == nil
|
||||
}
|
||||
|
||||
// IsUsable returns true if the connection is safe to use for new commands (lock-free).
|
||||
//
|
||||
// A connection is "usable" when it's in a stable state and can be returned to clients.
|
||||
// It becomes unusable during:
|
||||
// - Handoff operations (network connection replacement)
|
||||
// - Re-authentication (credential updates)
|
||||
// - Other background operations that need exclusive access
|
||||
//
|
||||
// Note: CREATED state is considered usable because new connections need to pass OnGet() hook
|
||||
// before initialization. The initialization happens after OnGet() in the client code.
|
||||
func (cn *Conn) IsUsable() bool {
|
||||
state := cn.stateMachine.GetState()
|
||||
// CREATED, IDLE, and IN_USE states are considered usable
|
||||
// CREATED: new connection, not yet initialized (will be initialized by client)
|
||||
// IDLE: initialized and ready to be acquired
|
||||
// IN_USE: usable but currently acquired by someone
|
||||
return state == StateCreated || state == StateIdle || state == StateInUse
|
||||
}
|
||||
|
||||
// SetUsable sets the usable flag for the connection (lock-free).
|
||||
//
|
||||
// Deprecated: Use GetStateMachine().Transition() directly for better state management.
|
||||
// This method is kept for backwards compatibility.
|
||||
//
|
||||
// This should be called to mark a connection as usable after initialization or
|
||||
// to release it after a background operation completes.
|
||||
//
|
||||
// Prefer CompareAndSwapUsable() when acquiring exclusive access to avoid race conditions.
|
||||
// Deprecated: Use GetStateMachine().Transition() directly for better state management.
|
||||
func (cn *Conn) SetUsable(usable bool) {
|
||||
if usable {
|
||||
// Transition to IDLE state (ready to be acquired)
|
||||
cn.stateMachine.Transition(StateIdle)
|
||||
} else {
|
||||
// Transition to UNUSABLE state (for background operations)
|
||||
cn.stateMachine.Transition(StateUnusable)
|
||||
}
|
||||
}
|
||||
|
||||
// IsInited returns true if the connection has been initialized.
|
||||
// This is a backward-compatible wrapper around the state machine.
|
||||
func (cn *Conn) IsInited() bool {
|
||||
state := cn.stateMachine.GetState()
|
||||
// Connection is initialized if it's in IDLE or any post-initialization state
|
||||
return state != StateCreated && state != StateInitializing && state != StateClosed
|
||||
}
|
||||
|
||||
// Used - State machine based implementation
|
||||
|
||||
// CompareAndSwapUsed atomically compares and swaps the used flag (lock-free).
|
||||
// This method is kept for backwards compatibility.
|
||||
//
|
||||
// This is the preferred method for acquiring a connection from the pool, as it
|
||||
// ensures that only one goroutine marks the connection as used.
|
||||
//
|
||||
// Implementation: Uses state machine transitions IDLE ⇄ IN_USE
|
||||
//
|
||||
// Returns true if the swap was successful (old value matched), false otherwise.
|
||||
// Deprecated: Use GetStateMachine().TryTransition() directly for better state management.
|
||||
func (cn *Conn) CompareAndSwapUsed(old, new bool) bool {
|
||||
if old == new {
|
||||
// No change needed
|
||||
currentState := cn.stateMachine.GetState()
|
||||
currentUsed := (currentState == StateInUse)
|
||||
return currentUsed == old
|
||||
}
|
||||
|
||||
if !old && new {
|
||||
// Acquiring: IDLE → IN_USE
|
||||
// Use predefined slice to avoid allocation
|
||||
_, err := cn.stateMachine.TryTransition(validFromCreatedOrIdle, StateInUse)
|
||||
return err == nil
|
||||
} else {
|
||||
// Releasing: IN_USE → IDLE
|
||||
// Use predefined slice to avoid allocation
|
||||
_, err := cn.stateMachine.TryTransition(validFromInUse, StateIdle)
|
||||
return err == nil
|
||||
}
|
||||
}
|
||||
|
||||
// IsUsed returns true if the connection is currently in use (lock-free).
|
||||
//
|
||||
// Deprecated: Use GetStateMachine().GetState() == StateInUse directly for better clarity.
|
||||
// This method is kept for backwards compatibility.
|
||||
//
|
||||
// A connection is "used" when it has been retrieved from the pool and is
|
||||
// actively processing a command. Background operations (like re-auth) should
|
||||
// wait until the connection is not used before executing commands.
|
||||
func (cn *Conn) IsUsed() bool {
|
||||
return cn.stateMachine.GetState() == StateInUse
|
||||
}
|
||||
|
||||
// SetUsed sets the used flag for the connection (lock-free).
|
||||
//
|
||||
// This should be called when returning a connection to the pool (set to false)
|
||||
// or when a single-connection pool retrieves its connection (set to true).
|
||||
//
|
||||
// Prefer CompareAndSwapUsed() when acquiring from a multi-connection pool to
|
||||
// avoid race conditions.
|
||||
// Deprecated: Use GetStateMachine().Transition() directly for better state management.
|
||||
func (cn *Conn) SetUsed(val bool) {
|
||||
if val {
|
||||
cn.stateMachine.Transition(StateInUse)
|
||||
} else {
|
||||
cn.stateMachine.Transition(StateIdle)
|
||||
}
|
||||
}
|
||||
|
||||
// getNetConn returns the current network connection using atomic load (lock-free).
|
||||
// This is the fast path for accessing netConn without mutex overhead.
|
||||
func (cn *Conn) getNetConn() net.Conn {
|
||||
if v := cn.netConnAtomic.Load(); v != nil {
|
||||
if wrapper, ok := v.(*atomicNetConn); ok {
|
||||
return wrapper.conn
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// setNetConn stores the network connection atomically (lock-free).
|
||||
// This is used for the fast path of connection replacement.
|
||||
func (cn *Conn) setNetConn(netConn net.Conn) {
|
||||
cn.netConnAtomic.Store(&atomicNetConn{conn: netConn})
|
||||
}
|
||||
|
||||
// Handoff state management - atomic access to handoff metadata
|
||||
|
||||
// ShouldHandoff returns true if connection needs handoff (lock-free).
|
||||
func (cn *Conn) ShouldHandoff() bool {
|
||||
if v := cn.handoffStateAtomic.Load(); v != nil {
|
||||
return v.(*HandoffState).ShouldHandoff
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// GetHandoffEndpoint returns the new endpoint for handoff (lock-free).
|
||||
func (cn *Conn) GetHandoffEndpoint() string {
|
||||
if v := cn.handoffStateAtomic.Load(); v != nil {
|
||||
return v.(*HandoffState).Endpoint
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// GetMovingSeqID returns the sequence ID from the MOVING notification (lock-free).
|
||||
func (cn *Conn) GetMovingSeqID() int64 {
|
||||
if v := cn.handoffStateAtomic.Load(); v != nil {
|
||||
return v.(*HandoffState).SeqID
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// GetHandoffInfo returns all handoff information atomically (lock-free).
|
||||
// This method prevents race conditions by returning all handoff state in a single atomic operation.
|
||||
// Returns (shouldHandoff, endpoint, seqID).
|
||||
func (cn *Conn) GetHandoffInfo() (bool, string, int64) {
|
||||
if v := cn.handoffStateAtomic.Load(); v != nil {
|
||||
state := v.(*HandoffState)
|
||||
return state.ShouldHandoff, state.Endpoint, state.SeqID
|
||||
}
|
||||
return false, "", 0
|
||||
}
|
||||
|
||||
// HandoffRetries returns the current handoff retry count (lock-free).
|
||||
func (cn *Conn) HandoffRetries() int {
|
||||
return int(cn.handoffRetriesAtomic.Load())
|
||||
}
|
||||
|
||||
// IncrementAndGetHandoffRetries atomically increments and returns handoff retries (lock-free).
|
||||
func (cn *Conn) IncrementAndGetHandoffRetries(n int) int {
|
||||
return int(cn.handoffRetriesAtomic.Add(uint32(n)))
|
||||
}
|
||||
|
||||
// IsPooled returns true if the connection is managed by a pool and will be pooled on Put.
|
||||
func (cn *Conn) IsPooled() bool {
|
||||
return cn.pooled
|
||||
}
|
||||
|
||||
// IsPubSub returns true if the connection is used for PubSub.
|
||||
func (cn *Conn) IsPubSub() bool {
|
||||
return cn.pubsub
|
||||
}
|
||||
|
||||
// SetRelaxedTimeout sets relaxed timeouts for this connection during maintenanceNotifications upgrades.
|
||||
// These timeouts will be used for all subsequent commands until the deadline expires.
|
||||
// Uses atomic operations for lock-free access.
|
||||
// Note: Metrics should be recorded by the caller (notification handler) which has context about
|
||||
// the notification type and pool name.
|
||||
func (cn *Conn) SetRelaxedTimeout(readTimeout, writeTimeout time.Duration) {
|
||||
cn.relaxedCounter.Add(1)
|
||||
cn.relaxedReadTimeoutNs.Store(int64(readTimeout))
|
||||
cn.relaxedWriteTimeoutNs.Store(int64(writeTimeout))
|
||||
}
|
||||
|
||||
// SetRelaxedTimeoutWithDeadline sets relaxed timeouts with an expiration deadline.
|
||||
// After the deadline, timeouts automatically revert to normal values.
|
||||
// Uses atomic operations for lock-free access.
|
||||
func (cn *Conn) SetRelaxedTimeoutWithDeadline(readTimeout, writeTimeout time.Duration, deadline time.Time) {
|
||||
cn.SetRelaxedTimeout(readTimeout, writeTimeout)
|
||||
cn.relaxedDeadlineNs.Store(deadline.UnixNano())
|
||||
}
|
||||
|
||||
// ClearRelaxedTimeout removes relaxed timeouts, returning to normal timeout behavior.
|
||||
// Uses atomic operations for lock-free access.
|
||||
func (cn *Conn) ClearRelaxedTimeout() {
|
||||
// Atomically decrement counter and check if we should clear
|
||||
newCount := cn.relaxedCounter.Add(-1)
|
||||
deadlineNs := cn.relaxedDeadlineNs.Load()
|
||||
if newCount <= 0 && (deadlineNs == 0 || time.Now().UnixNano() >= deadlineNs) {
|
||||
// Use atomic load to get current value for CAS to avoid stale value race
|
||||
current := cn.relaxedCounter.Load()
|
||||
if current <= 0 && cn.relaxedCounter.CompareAndSwap(current, 0) {
|
||||
cn.clearRelaxedTimeout()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (cn *Conn) clearRelaxedTimeout() {
|
||||
cn.relaxedReadTimeoutNs.Store(0)
|
||||
cn.relaxedWriteTimeoutNs.Store(0)
|
||||
cn.relaxedDeadlineNs.Store(0)
|
||||
cn.relaxedCounter.Store(0)
|
||||
|
||||
// Note: Metrics for timeout unrelaxing are not recorded here because we don't have
|
||||
// context about which notification type or pool triggered the relaxation.
|
||||
// In practice, relaxed timeouts expire automatically via deadline, so explicit
|
||||
// unrelaxing metrics are less critical than the initial relaxation metrics.
|
||||
}
|
||||
|
||||
// HasRelaxedTimeout returns true if relaxed timeouts are currently active on this connection.
|
||||
// This checks both the timeout values and the deadline (if set).
|
||||
// Uses atomic operations for lock-free access.
|
||||
func (cn *Conn) HasRelaxedTimeout() bool {
|
||||
// Fast path: no relaxed timeouts are set
|
||||
if cn.relaxedCounter.Load() <= 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
readTimeoutNs := cn.relaxedReadTimeoutNs.Load()
|
||||
writeTimeoutNs := cn.relaxedWriteTimeoutNs.Load()
|
||||
|
||||
// If no relaxed timeouts are set, return false
|
||||
if readTimeoutNs <= 0 && writeTimeoutNs <= 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
deadlineNs := cn.relaxedDeadlineNs.Load()
|
||||
// If no deadline is set, relaxed timeouts are active
|
||||
if deadlineNs == 0 {
|
||||
return true
|
||||
}
|
||||
|
||||
// If deadline is set, check if it's still in the future
|
||||
return time.Now().UnixNano() < deadlineNs
|
||||
}
|
||||
|
||||
// getEffectiveReadTimeout returns the timeout to use for read operations.
|
||||
// If relaxed timeout is set and not expired, it takes precedence over the provided timeout.
|
||||
// This method automatically clears expired relaxed timeouts using atomic operations.
|
||||
func (cn *Conn) getEffectiveReadTimeout(normalTimeout time.Duration) time.Duration {
|
||||
readTimeoutNs := cn.relaxedReadTimeoutNs.Load()
|
||||
|
||||
// Fast path: no relaxed timeout set
|
||||
if readTimeoutNs <= 0 {
|
||||
return normalTimeout
|
||||
}
|
||||
|
||||
deadlineNs := cn.relaxedDeadlineNs.Load()
|
||||
// If no deadline is set, use relaxed timeout
|
||||
if deadlineNs == 0 {
|
||||
return time.Duration(readTimeoutNs)
|
||||
}
|
||||
|
||||
// Use cached time to avoid expensive syscall (max 50ms staleness is acceptable for timeout checks)
|
||||
nowNs := getCachedTimeNs()
|
||||
// Check if deadline has passed
|
||||
if nowNs < deadlineNs {
|
||||
// Deadline is in the future, use relaxed timeout
|
||||
return time.Duration(readTimeoutNs)
|
||||
} else {
|
||||
// Deadline has passed, clear relaxed timeouts atomically and use normal timeout
|
||||
newCount := cn.relaxedCounter.Add(-1)
|
||||
if newCount <= 0 {
|
||||
internal.Logger.Printf(context.Background(), logs.UnrelaxedTimeoutAfterDeadline(cn.GetID()))
|
||||
cn.clearRelaxedTimeout()
|
||||
}
|
||||
return normalTimeout
|
||||
}
|
||||
}
|
||||
|
||||
// getEffectiveWriteTimeout returns the timeout to use for write operations.
|
||||
// If relaxed timeout is set and not expired, it takes precedence over the provided timeout.
|
||||
// This method automatically clears expired relaxed timeouts using atomic operations.
|
||||
func (cn *Conn) getEffectiveWriteTimeout(normalTimeout time.Duration) time.Duration {
|
||||
writeTimeoutNs := cn.relaxedWriteTimeoutNs.Load()
|
||||
|
||||
// Fast path: no relaxed timeout set
|
||||
if writeTimeoutNs <= 0 {
|
||||
return normalTimeout
|
||||
}
|
||||
|
||||
deadlineNs := cn.relaxedDeadlineNs.Load()
|
||||
// If no deadline is set, use relaxed timeout
|
||||
if deadlineNs == 0 {
|
||||
return time.Duration(writeTimeoutNs)
|
||||
}
|
||||
|
||||
// Use cached time to avoid expensive syscall (max 50ms staleness is acceptable for timeout checks)
|
||||
nowNs := getCachedTimeNs()
|
||||
// Check if deadline has passed
|
||||
if nowNs < deadlineNs {
|
||||
// Deadline is in the future, use relaxed timeout
|
||||
return time.Duration(writeTimeoutNs)
|
||||
} else {
|
||||
// Deadline has passed, clear relaxed timeouts atomically and use normal timeout
|
||||
newCount := cn.relaxedCounter.Add(-1)
|
||||
if newCount <= 0 {
|
||||
internal.Logger.Printf(context.Background(), logs.UnrelaxedTimeoutAfterDeadline(cn.GetID()))
|
||||
cn.clearRelaxedTimeout()
|
||||
}
|
||||
return normalTimeout
|
||||
}
|
||||
}
|
||||
|
||||
// SetOnClose installs fn as the callback invoked exactly once when this
|
||||
// connection is closed (via Conn.Close).
|
||||
//
|
||||
// IMPORTANT: SetOnClose OVERWRITES any previously installed callback — it
|
||||
// does not compose, chain, or deduplicate. A Conn has room for a single
|
||||
// onClose hook by design, because its lifecycle is bounded (a Conn is
|
||||
// created, optionally re-initialized on its own net.Conn, and then closed
|
||||
// once) and the pool's OnRemove hooks handle any registry-level cleanup
|
||||
// that must survive the net.Conn being swapped.
|
||||
//
|
||||
// This has a subtle implication for per-connection subscriptions such as
|
||||
// the unsubscribe function returned by StreamingCredentialsProvider
|
||||
// (e.g. EntraID token rotation): if SetOnClose is called twice on the
|
||||
// same Conn with DIFFERENT unsubscribe closures — for example because
|
||||
// initConn ran a second time and obtained a fresh Subscribe() —
|
||||
// the previous unsubscribe is dropped and will NEVER run, leaking a
|
||||
// subscription on the provider. Callers must therefore ensure either:
|
||||
//
|
||||
// - the provider's Subscribe is idempotent for the same listener (the
|
||||
// streaming credentials Manager deduplicates listeners by connection
|
||||
// id, so re-Subscribe returns an equivalent unsubscribe), OR
|
||||
// - the previous callback has already been invoked before SetOnClose is
|
||||
// called again.
|
||||
//
|
||||
// Design note: unlike the client-level onCloseHooks registry (see
|
||||
// redis.baseClient), there is intentionally NO named-hook dedup or
|
||||
// multi-callback support on Conn. This is a deliberate trade-off to keep
|
||||
// the Conn object slim — a pool can hold thousands of Conn values and
|
||||
// each one is a hot allocation, so paying for a sync.Mutex plus a
|
||||
// map[string]func() error per connection to support a feature that would
|
||||
// only be used by at most one subsystem today (streaming credentials) is
|
||||
// not worth the per-connection memory and allocation cost. For a single
|
||||
// Conn there is at most one meaningful close callback at any point in
|
||||
// time, and a richer registry here would not even solve the "stale
|
||||
// closure" hazard described above.
|
||||
func (cn *Conn) SetOnClose(fn func() error) {
|
||||
cn.onClose = fn
|
||||
}
|
||||
|
||||
// SetInitConnFunc sets the connection initialization function to be called on reconnections.
|
||||
func (cn *Conn) SetInitConnFunc(fn func(context.Context, *Conn) error) {
|
||||
cn.initConnFunc = fn
|
||||
}
|
||||
|
||||
// ExecuteInitConn runs the stored connection initialization function if available.
|
||||
func (cn *Conn) ExecuteInitConn(ctx context.Context) error {
|
||||
if cn.initConnFunc != nil {
|
||||
return cn.initConnFunc(ctx, cn)
|
||||
}
|
||||
return fmt.Errorf("redis: no initConnFunc set for conn[%d]", cn.GetID())
|
||||
}
|
||||
|
||||
func (cn *Conn) SetNetConn(netConn net.Conn) {
|
||||
// Store the new connection atomically first (lock-free)
|
||||
cn.setNetConn(netConn)
|
||||
// Protect reader reset operations to avoid data races
|
||||
// Use write lock since we're modifying the reader state
|
||||
cn.readerMu.Lock()
|
||||
cn.rd.Reset(netConn)
|
||||
cn.readerMu.Unlock()
|
||||
|
||||
cn.bw.Reset(netConn)
|
||||
}
|
||||
|
||||
// GetNetConn safely returns the current network connection using atomic load (lock-free).
|
||||
// This method is used by the pool for health checks and provides better performance.
|
||||
func (cn *Conn) GetNetConn() net.Conn {
|
||||
return cn.getNetConn()
|
||||
}
|
||||
|
||||
// SetNetConnAndInitConn replaces the underlying connection and executes the initialization.
|
||||
// This method ensures only one initialization can happen at a time by using atomic state transitions.
|
||||
// If another goroutine is currently initializing, this will wait for it to complete.
|
||||
func (cn *Conn) SetNetConnAndInitConn(ctx context.Context, netConn net.Conn) error {
|
||||
// Wait for and transition to INITIALIZING state - this prevents concurrent initializations
|
||||
// Valid from states: CREATED (first init), IDLE (reconnect), UNUSABLE (handoff/reauth)
|
||||
// If another goroutine is initializing, we'll wait for it to finish
|
||||
// if the context has a deadline, use that, otherwise use the connection read (relaxed) timeout
|
||||
// which should be set during handoff. If it is not set, use a 5 second default
|
||||
deadline, ok := ctx.Deadline()
|
||||
if !ok {
|
||||
deadline = time.Now().Add(cn.getEffectiveReadTimeout(5 * time.Second))
|
||||
}
|
||||
waitCtx, cancel := context.WithDeadline(ctx, deadline)
|
||||
defer cancel()
|
||||
// Use predefined slice to avoid allocation
|
||||
finalState, err := cn.stateMachine.AwaitAndTransition(
|
||||
waitCtx,
|
||||
validFromCreatedIdleOrUnusable,
|
||||
StateInitializing,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("cannot initialize connection from state %s: %w", finalState, err)
|
||||
}
|
||||
|
||||
// Replace the underlying connection
|
||||
cn.SetNetConn(netConn)
|
||||
|
||||
// Execute initialization
|
||||
// NOTE: ExecuteInitConn (via baseClient.initConn) will transition to IDLE on success
|
||||
// or CLOSED on failure. We don't need to do it here.
|
||||
// NOTE: Initconn returns conn in IDLE state
|
||||
initErr := cn.ExecuteInitConn(ctx)
|
||||
if initErr != nil {
|
||||
// ExecuteInitConn already transitioned to CLOSED, just return the error
|
||||
return initErr
|
||||
}
|
||||
|
||||
// ExecuteInitConn already transitioned to IDLE
|
||||
return nil
|
||||
}
|
||||
|
||||
// MarkForHandoff marks the connection for handoff due to MOVING notification.
|
||||
// Returns an error if the connection is already marked for handoff.
|
||||
// Note: This only sets metadata - the connection state is not changed until OnPut.
|
||||
// This allows the current user to finish using the connection before handoff.
|
||||
func (cn *Conn) MarkForHandoff(newEndpoint string, seqID int64) error {
|
||||
// Check if already marked for handoff
|
||||
if cn.ShouldHandoff() {
|
||||
return errAlreadyMarkedForHandoff
|
||||
}
|
||||
|
||||
// Set handoff metadata atomically
|
||||
cn.handoffStateAtomic.Store(&HandoffState{
|
||||
ShouldHandoff: true,
|
||||
Endpoint: newEndpoint,
|
||||
SeqID: seqID,
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// MarkQueuedForHandoff marks the connection as queued for handoff processing.
|
||||
// This makes the connection unusable until handoff completes.
|
||||
// This is called from OnPut hook, where the connection is typically in IN_USE state.
|
||||
// The pool will preserve the UNUSABLE state and not overwrite it with IDLE.
|
||||
func (cn *Conn) MarkQueuedForHandoff() error {
|
||||
// Get current handoff state
|
||||
currentState := cn.handoffStateAtomic.Load()
|
||||
if currentState == nil {
|
||||
return errNotMarkedForHandoff
|
||||
}
|
||||
|
||||
state := currentState.(*HandoffState)
|
||||
if !state.ShouldHandoff {
|
||||
return errNotMarkedForHandoff
|
||||
}
|
||||
|
||||
// Create new state with ShouldHandoff=false but preserve endpoint and seqID
|
||||
// This prevents the connection from being queued multiple times while still
|
||||
// allowing the worker to access the handoff metadata
|
||||
newState := &HandoffState{
|
||||
ShouldHandoff: false,
|
||||
Endpoint: state.Endpoint, // Preserve endpoint for handoff processing
|
||||
SeqID: state.SeqID, // Preserve seqID for handoff processing
|
||||
}
|
||||
|
||||
// Atomic compare-and-swap to update state
|
||||
if !cn.handoffStateAtomic.CompareAndSwap(currentState, newState) {
|
||||
// State changed between load and CAS - retry or return error
|
||||
return errHandoffStateChanged
|
||||
}
|
||||
|
||||
// Transition to UNUSABLE from IN_USE (normal flow), IDLE (edge cases), or CREATED (tests/uninitialized)
|
||||
// The connection is typically in IN_USE state when OnPut is called (normal Put flow)
|
||||
// But in some edge cases or tests, it might be in IDLE or CREATED state
|
||||
// The pool will detect this state change and preserve it (not overwrite with IDLE)
|
||||
// Use predefined slice to avoid allocation
|
||||
finalState, err := cn.stateMachine.TryTransition(validFromCreatedInUseOrIdle, StateUnusable)
|
||||
if err != nil {
|
||||
// Check if already in UNUSABLE state (race condition or retry)
|
||||
// ShouldHandoff should be false now, but check just in case
|
||||
if finalState == StateUnusable && !cn.ShouldHandoff() {
|
||||
// Already unusable - this is fine, keep the new handoff state
|
||||
return nil
|
||||
}
|
||||
// Restore the original state if transition fails for other reasons
|
||||
cn.handoffStateAtomic.Store(currentState)
|
||||
return fmt.Errorf("failed to mark connection as unusable: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetID returns the unique identifier for this connection.
|
||||
func (cn *Conn) GetID() uint64 {
|
||||
return cn.id
|
||||
}
|
||||
|
||||
// GetStateMachine returns the connection's state machine for advanced state management.
|
||||
// This is primarily used by internal packages like maintnotifications for handoff processing.
|
||||
func (cn *Conn) GetStateMachine() *ConnStateMachine {
|
||||
return cn.stateMachine
|
||||
}
|
||||
|
||||
// TryAcquire attempts to acquire the connection for use.
|
||||
// This is an optimized inline method for the hot path (Get operation).
|
||||
//
|
||||
// It tries to transition from IDLE -> IN_USE or CREATED -> CREATED.
|
||||
// Returns true if the connection was successfully acquired, false otherwise.
|
||||
// The CREATED->CREATED is done so we can keep the state correct for later
|
||||
// initialization of the connection in initConn.
|
||||
//
|
||||
// Performance: This is faster than calling GetStateMachine() + TryTransitionFast()
|
||||
//
|
||||
// NOTE: We directly access cn.stateMachine.state here instead of using the state machine's
|
||||
// methods. This breaks encapsulation but is necessary for performance.
|
||||
// The IDLE->IN_USE and CREATED->CREATED transitions don't need
|
||||
// waiter notification, and benchmarks show 1-3% improvement. If the state machine ever
|
||||
// needs to notify waiters on these transitions, update this to use TryTransitionFast().
|
||||
func (cn *Conn) TryAcquire() bool {
|
||||
// The || operator short-circuits, so only 1 CAS in the common case
|
||||
return cn.stateMachine.state.CompareAndSwap(uint32(StateIdle), uint32(StateInUse)) ||
|
||||
cn.stateMachine.state.CompareAndSwap(uint32(StateCreated), uint32(StateCreated))
|
||||
}
|
||||
|
||||
// Release releases the connection back to the pool.
|
||||
// This is an optimized inline method for the hot path (Put operation).
|
||||
//
|
||||
// It tries to transition from IN_USE -> IDLE.
|
||||
// Returns true if the connection was successfully released, false otherwise.
|
||||
//
|
||||
// Performance: This is faster than calling GetStateMachine() + TryTransitionFast().
|
||||
//
|
||||
// NOTE: We directly access cn.stateMachine.state here instead of using the state machine's
|
||||
// methods. This breaks encapsulation but is necessary for performance.
|
||||
// If the state machine ever needs to notify waiters
|
||||
// on this transition, update this to use TryTransitionFast().
|
||||
func (cn *Conn) Release() bool {
|
||||
// Inline the hot path - single CAS operation
|
||||
return cn.stateMachine.state.CompareAndSwap(uint32(StateInUse), uint32(StateIdle))
|
||||
}
|
||||
|
||||
// ClearHandoffState clears the handoff state after successful handoff.
|
||||
// Makes the connection usable again.
|
||||
func (cn *Conn) ClearHandoffState() {
|
||||
// Clear handoff metadata
|
||||
cn.handoffStateAtomic.Store(&HandoffState{
|
||||
ShouldHandoff: false,
|
||||
Endpoint: "",
|
||||
SeqID: 0,
|
||||
})
|
||||
|
||||
// Reset retry counter
|
||||
cn.handoffRetriesAtomic.Store(0)
|
||||
|
||||
// Mark connection as usable again
|
||||
// Use state machine directly instead of deprecated SetUsable
|
||||
// probably done by initConn
|
||||
cn.stateMachine.Transition(StateIdle)
|
||||
}
|
||||
|
||||
// HasBufferedData safely checks if the connection has buffered data.
|
||||
// This method is used to avoid data races when checking for push notifications.
|
||||
func (cn *Conn) HasBufferedData() bool {
|
||||
// Use read lock for concurrent access to reader state
|
||||
cn.readerMu.RLock()
|
||||
defer cn.readerMu.RUnlock()
|
||||
return cn.rd.Buffered() > 0
|
||||
}
|
||||
|
||||
// PeekReplyTypeSafe safely peeks at the reply type.
|
||||
// This method is used to avoid data races when checking for push notifications.
|
||||
func (cn *Conn) PeekReplyTypeSafe() (byte, error) {
|
||||
// Use read lock for concurrent access to reader state
|
||||
cn.readerMu.RLock()
|
||||
defer cn.readerMu.RUnlock()
|
||||
|
||||
if cn.rd.Buffered() <= 0 {
|
||||
return 0, fmt.Errorf("redis: can't peek reply type, no data available")
|
||||
}
|
||||
return cn.rd.PeekReplyType()
|
||||
}
|
||||
|
||||
func (cn *Conn) Write(b []byte) (int, error) {
|
||||
// Lock-free netConn access for better performance
|
||||
if netConn := cn.getNetConn(); netConn != nil {
|
||||
return netConn.Write(b)
|
||||
}
|
||||
return 0, net.ErrClosed
|
||||
}
|
||||
|
||||
func (cn *Conn) RemoteAddr() net.Addr {
|
||||
// Lock-free netConn access for better performance
|
||||
if netConn := cn.getNetConn(); netConn != nil {
|
||||
return netConn.RemoteAddr()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (cn *Conn) WithReader(
|
||||
ctx context.Context, timeout time.Duration, fn func(rd *proto.Reader) error,
|
||||
) error {
|
||||
if timeout >= 0 {
|
||||
// Use relaxed timeout if set, otherwise use provided timeout
|
||||
effectiveTimeout := cn.getEffectiveReadTimeout(timeout)
|
||||
|
||||
// Get the connection directly from atomic storage
|
||||
netConn := cn.getNetConn()
|
||||
if netConn == nil {
|
||||
return errConnectionNotAvailable
|
||||
}
|
||||
|
||||
if err := netConn.SetReadDeadline(cn.deadline(ctx, effectiveTimeout)); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return fn(cn.rd)
|
||||
}
|
||||
|
||||
func (cn *Conn) WithWriter(
|
||||
ctx context.Context, timeout time.Duration, fn func(wr *proto.Writer) error,
|
||||
) error {
|
||||
if timeout >= 0 {
|
||||
// Use relaxed timeout if set, otherwise use provided timeout
|
||||
effectiveTimeout := cn.getEffectiveWriteTimeout(timeout)
|
||||
|
||||
// Set write deadline on the connection
|
||||
if netConn := cn.getNetConn(); netConn != nil {
|
||||
if err := netConn.SetWriteDeadline(cn.deadline(ctx, effectiveTimeout)); err != nil {
|
||||
return err
|
||||
}
|
||||
} else {
|
||||
// Connection is not available - return preallocated error
|
||||
return errConnNotAvailableForWrite
|
||||
}
|
||||
}
|
||||
|
||||
// Reset the buffered writer if needed, should not happen
|
||||
if cn.bw.Buffered() > 0 {
|
||||
if netConn := cn.getNetConn(); netConn != nil {
|
||||
cn.bw.Reset(netConn)
|
||||
}
|
||||
}
|
||||
|
||||
if err := fn(cn.wr); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return cn.bw.Flush()
|
||||
}
|
||||
|
||||
func (cn *Conn) IsClosed() bool {
|
||||
return cn.stateMachine.GetState() == StateClosed
|
||||
}
|
||||
|
||||
func (cn *Conn) Close() error {
|
||||
if cn.IsClosed() {
|
||||
return nil
|
||||
}
|
||||
// Transition to CLOSED state
|
||||
cn.stateMachine.Transition(StateClosed)
|
||||
|
||||
if cn.onClose != nil {
|
||||
// ignore error
|
||||
_ = cn.onClose()
|
||||
cn.onClose = nil
|
||||
}
|
||||
|
||||
// Lock-free netConn access for better performance
|
||||
if netConn := cn.getNetConn(); netConn != nil {
|
||||
return netConn.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// MaybeHasData tries to peek at the next byte in the socket without consuming it
|
||||
// This is used to check if there are push notifications available
|
||||
// Important: This will work on Linux, but not on Windows
|
||||
func (cn *Conn) MaybeHasData() bool {
|
||||
// Lock-free netConn access for better performance
|
||||
if netConn := cn.getNetConn(); netConn != nil {
|
||||
return maybeHasData(netConn)
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// deadline computes the effective deadline time based on context and timeout.
|
||||
// It updates the usedAt timestamp to now.
|
||||
// Uses cached time to avoid expensive syscall (max 50ms staleness is acceptable for deadline calculation).
|
||||
func (cn *Conn) deadline(ctx context.Context, timeout time.Duration) time.Time {
|
||||
// Use cached time for deadline calculation (called 2x per command: read + write)
|
||||
nowNs := getCachedTimeNs()
|
||||
cn.SetUsedAtNs(nowNs)
|
||||
tm := time.Unix(0, nowNs)
|
||||
|
||||
if timeout > 0 {
|
||||
tm = tm.Add(timeout)
|
||||
}
|
||||
|
||||
if ctx != nil {
|
||||
deadline, ok := ctx.Deadline()
|
||||
if ok {
|
||||
if timeout == 0 {
|
||||
return deadline
|
||||
}
|
||||
if deadline.Before(tm) {
|
||||
return deadline
|
||||
}
|
||||
return tm
|
||||
}
|
||||
}
|
||||
|
||||
if timeout > 0 {
|
||||
return tm
|
||||
}
|
||||
|
||||
return noDeadline
|
||||
}
|
||||
+59
@@ -0,0 +1,59 @@
|
||||
//go:build linux || darwin || dragonfly || freebsd || netbsd || openbsd || solaris || illumos
|
||||
|
||||
package pool
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"io"
|
||||
"net"
|
||||
"syscall"
|
||||
"time"
|
||||
)
|
||||
|
||||
var errUnexpectedRead = errors.New("unexpected read from socket")
|
||||
|
||||
// connCheck checks if the connection is still alive and if there is data in the socket
|
||||
// it will try to peek at the next byte without consuming it since we may want to work with it
|
||||
// later on (e.g. push notifications)
|
||||
func connCheck(conn net.Conn) error {
|
||||
// Reset previous timeout.
|
||||
_ = conn.SetDeadline(time.Time{})
|
||||
|
||||
sysConn, ok := conn.(syscall.Conn)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
rawConn, err := sysConn.SyscallConn()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
var sysErr error
|
||||
|
||||
if err := rawConn.Read(func(fd uintptr) bool {
|
||||
var buf [1]byte
|
||||
// Use MSG_PEEK to peek at data without consuming it
|
||||
n, _, err := syscall.Recvfrom(int(fd), buf[:], syscall.MSG_PEEK|syscall.MSG_DONTWAIT)
|
||||
|
||||
switch {
|
||||
case n == 0 && err == nil:
|
||||
sysErr = io.EOF
|
||||
case n > 0:
|
||||
sysErr = errUnexpectedRead
|
||||
case err == syscall.EAGAIN || err == syscall.EWOULDBLOCK:
|
||||
sysErr = nil
|
||||
default:
|
||||
sysErr = err
|
||||
}
|
||||
return true
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return sysErr
|
||||
}
|
||||
|
||||
// maybeHasData checks if there is data in the socket without consuming it
|
||||
func maybeHasData(conn net.Conn) bool {
|
||||
return connCheck(conn) == errUnexpectedRead
|
||||
}
|
||||
+20
@@ -0,0 +1,20 @@
|
||||
//go:build !linux && !darwin && !dragonfly && !freebsd && !netbsd && !openbsd && !solaris && !illumos
|
||||
|
||||
package pool
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net"
|
||||
)
|
||||
|
||||
// errUnexpectedRead is placeholder error variable for non-unix build constraints
|
||||
var errUnexpectedRead = errors.New("unexpected read from socket")
|
||||
|
||||
func connCheck(_ net.Conn) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// since we can't check for data on the socket, we just assume there is some
|
||||
func maybeHasData(_ net.Conn) bool {
|
||||
return true
|
||||
}
|
||||
+336
@@ -0,0 +1,336 @@
|
||||
package pool
|
||||
|
||||
import (
|
||||
"container/list"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
// ConnState represents the connection state in the state machine.
|
||||
// States are designed to be lightweight and fast to check.
|
||||
//
|
||||
// State Transitions:
|
||||
//
|
||||
// CREATED → INITIALIZING → IDLE ⇄ IN_USE
|
||||
// ↓
|
||||
// UNUSABLE (handoff/reauth)
|
||||
// ↓
|
||||
// IDLE/CLOSED
|
||||
type ConnState uint32
|
||||
|
||||
const (
|
||||
// StateCreated - Connection just created, not yet initialized
|
||||
StateCreated ConnState = iota
|
||||
|
||||
// StateInitializing - Connection initialization in progress
|
||||
StateInitializing
|
||||
|
||||
// StateIdle - Connection initialized and idle in pool, ready to be acquired
|
||||
StateIdle
|
||||
|
||||
// StateInUse - Connection actively processing a command (retrieved from pool)
|
||||
StateInUse
|
||||
|
||||
// StateUnusable - Connection temporarily unusable due to background operation
|
||||
// (handoff, reauth, etc.). Cannot be acquired from pool.
|
||||
StateUnusable
|
||||
|
||||
// StateClosed - Connection closed
|
||||
StateClosed
|
||||
)
|
||||
|
||||
// Predefined state slices to avoid allocations in hot paths
|
||||
var (
|
||||
validFromInUse = []ConnState{StateInUse}
|
||||
validFromCreatedOrIdle = []ConnState{StateCreated, StateIdle}
|
||||
validFromCreatedInUseOrIdle = []ConnState{StateCreated, StateInUse, StateIdle}
|
||||
// For AwaitAndTransition calls
|
||||
validFromCreatedIdleOrUnusable = []ConnState{StateCreated, StateIdle, StateUnusable}
|
||||
validFromIdle = []ConnState{StateIdle}
|
||||
// For CompareAndSwapUsable
|
||||
validFromInitializingOrUnusable = []ConnState{StateInitializing, StateUnusable}
|
||||
)
|
||||
|
||||
// Accessor functions for predefined slices to avoid allocations in external packages
|
||||
// These return the same slice instance, so they're zero-allocation
|
||||
|
||||
// ValidFromIdle returns a predefined slice containing only StateIdle.
|
||||
// Use this to avoid allocations when calling AwaitAndTransition or TryTransition.
|
||||
func ValidFromIdle() []ConnState {
|
||||
return validFromIdle
|
||||
}
|
||||
|
||||
// ValidFromCreatedIdleOrUnusable returns a predefined slice for initialization transitions.
|
||||
// Use this to avoid allocations when calling AwaitAndTransition or TryTransition.
|
||||
func ValidFromCreatedIdleOrUnusable() []ConnState {
|
||||
return validFromCreatedIdleOrUnusable
|
||||
}
|
||||
|
||||
// String returns a human-readable string representation of the state.
|
||||
func (s ConnState) String() string {
|
||||
switch s {
|
||||
case StateCreated:
|
||||
return "CREATED"
|
||||
case StateInitializing:
|
||||
return "INITIALIZING"
|
||||
case StateIdle:
|
||||
return "IDLE"
|
||||
case StateInUse:
|
||||
return "IN_USE"
|
||||
case StateUnusable:
|
||||
return "UNUSABLE"
|
||||
case StateClosed:
|
||||
return "CLOSED"
|
||||
default:
|
||||
return fmt.Sprintf("UNKNOWN(%d)", s)
|
||||
}
|
||||
}
|
||||
|
||||
var (
|
||||
// ErrInvalidStateTransition is returned when a state transition is not allowed
|
||||
ErrInvalidStateTransition = errors.New("invalid state transition")
|
||||
|
||||
// ErrStateMachineClosed is returned when operating on a closed state machine
|
||||
ErrStateMachineClosed = errors.New("state machine is closed")
|
||||
|
||||
// ErrTimeout is returned when a state transition times out
|
||||
ErrTimeout = errors.New("state transition timeout")
|
||||
)
|
||||
|
||||
// waiter represents a goroutine waiting for a state transition.
|
||||
// Designed for minimal allocations and fast processing.
|
||||
type waiter struct {
|
||||
validStates map[ConnState]struct{} // States we're waiting for
|
||||
targetState ConnState // State to transition to
|
||||
done chan error // Signaled when transition completes or times out
|
||||
}
|
||||
|
||||
// ConnStateMachine manages connection state transitions with FIFO waiting queue.
|
||||
// Optimized for:
|
||||
// - Lock-free reads (hot path)
|
||||
// - Minimal allocations
|
||||
// - Fast state transitions
|
||||
// - FIFO fairness for waiters
|
||||
// Note: Handoff metadata (endpoint, seqID, retries) is managed separately in the Conn struct.
|
||||
type ConnStateMachine struct {
|
||||
// Current state - atomic for lock-free reads
|
||||
state atomic.Uint32
|
||||
|
||||
// FIFO queue for waiters - only locked during waiter add/remove/notify
|
||||
mu sync.Mutex
|
||||
waiters *list.List // List of *waiter
|
||||
waiterCount atomic.Int32 // Fast lock-free check for waiters (avoids mutex in hot path)
|
||||
}
|
||||
|
||||
// NewConnStateMachine creates a new connection state machine.
|
||||
// Initial state is StateCreated.
|
||||
func NewConnStateMachine() *ConnStateMachine {
|
||||
sm := &ConnStateMachine{
|
||||
waiters: list.New(),
|
||||
}
|
||||
sm.state.Store(uint32(StateCreated))
|
||||
return sm
|
||||
}
|
||||
|
||||
// GetState returns the current state (lock-free read).
|
||||
// This is the hot path - optimized for zero allocations and minimal overhead.
|
||||
// Note: Zero allocations applies to state reads; converting the returned state to a string
|
||||
// (via String()) may allocate if the state is unknown.
|
||||
func (sm *ConnStateMachine) GetState() ConnState {
|
||||
return ConnState(sm.state.Load())
|
||||
}
|
||||
|
||||
// TryTransitionFast is an optimized version for the hot path (Get/Put operations).
|
||||
// It only handles simple state transitions without waiter notification.
|
||||
// This is safe because:
|
||||
// 1. Get/Put don't need to wait for state changes
|
||||
// 2. Background operations (handoff/reauth) use UNUSABLE state, which this won't match
|
||||
// 3. If a background operation is in progress (state is UNUSABLE), this fails fast
|
||||
//
|
||||
// Returns true if transition succeeded, false otherwise.
|
||||
// Use this for performance-critical paths where you don't need error details.
|
||||
//
|
||||
// Performance: Single CAS operation - as fast as the old atomic bool!
|
||||
// For multiple from states, use: sm.TryTransitionFast(State1, Target) || sm.TryTransitionFast(State2, Target)
|
||||
// The || operator short-circuits, so only 1 CAS is executed in the common case.
|
||||
func (sm *ConnStateMachine) TryTransitionFast(fromState, targetState ConnState) bool {
|
||||
return sm.state.CompareAndSwap(uint32(fromState), uint32(targetState))
|
||||
}
|
||||
|
||||
// TryTransition attempts an immediate state transition without waiting.
|
||||
// Returns the current state after the transition attempt and an error if the transition failed.
|
||||
// The returned state is the CURRENT state (after the attempt), not the previous state.
|
||||
// This is faster than AwaitAndTransition when you don't need to wait.
|
||||
// Uses compare-and-swap to atomically transition, preventing concurrent transitions.
|
||||
// This method does NOT wait - it fails immediately if the transition cannot be performed.
|
||||
//
|
||||
// Performance: Zero allocations on success path (hot path).
|
||||
func (sm *ConnStateMachine) TryTransition(validFromStates []ConnState, targetState ConnState) (ConnState, error) {
|
||||
// Try each valid from state with CAS
|
||||
// This ensures only ONE goroutine can successfully transition at a time
|
||||
for _, fromState := range validFromStates {
|
||||
// Try to atomically swap from fromState to targetState
|
||||
// If successful, we won the race and can proceed
|
||||
if sm.state.CompareAndSwap(uint32(fromState), uint32(targetState)) {
|
||||
// Success! We transitioned atomically
|
||||
// Hot path optimization: only check for waiters if transition succeeded
|
||||
// This avoids atomic load on every Get/Put when no waiters exist
|
||||
if sm.waiterCount.Load() > 0 {
|
||||
sm.notifyWaiters()
|
||||
}
|
||||
return targetState, nil
|
||||
}
|
||||
}
|
||||
|
||||
// All CAS attempts failed - state is not valid for this transition
|
||||
// Return the current state so caller can decide what to do
|
||||
// Note: This error path allocates, but it's the exceptional case
|
||||
currentState := sm.GetState()
|
||||
return currentState, fmt.Errorf("%w: cannot transition from %s to %s (valid from: %v)",
|
||||
ErrInvalidStateTransition, currentState, targetState, validFromStates)
|
||||
}
|
||||
|
||||
// Transition unconditionally transitions to the target state.
|
||||
// Use with caution - prefer AwaitAndTransition or TryTransition for safety.
|
||||
// This is useful for error paths or when you know the transition is valid.
|
||||
func (sm *ConnStateMachine) Transition(targetState ConnState) {
|
||||
sm.state.Store(uint32(targetState))
|
||||
sm.notifyWaiters()
|
||||
}
|
||||
|
||||
// AwaitAndTransition waits for the connection to reach one of the valid states,
|
||||
// then atomically transitions to the target state.
|
||||
// Returns the current state after the transition attempt and an error if the operation failed.
|
||||
// The returned state is the CURRENT state (after the attempt), not the previous state.
|
||||
// Returns error if timeout expires or context is cancelled.
|
||||
//
|
||||
// This method implements FIFO fairness - the first caller to wait gets priority
|
||||
// when the state becomes available.
|
||||
//
|
||||
// Performance notes:
|
||||
// - If already in a valid state, this is very fast (no allocation, no waiting)
|
||||
// - If waiting is required, allocates one waiter struct and one channel
|
||||
func (sm *ConnStateMachine) AwaitAndTransition(
|
||||
ctx context.Context,
|
||||
validFromStates []ConnState,
|
||||
targetState ConnState,
|
||||
) (ConnState, error) {
|
||||
// Fast path: try immediate transition with CAS to prevent race conditions
|
||||
// BUT: only if there are no waiters in the queue (to maintain FIFO ordering)
|
||||
if sm.waiterCount.Load() == 0 {
|
||||
for _, fromState := range validFromStates {
|
||||
// Check if we're already in target state
|
||||
if fromState == targetState && sm.GetState() == targetState {
|
||||
return targetState, nil
|
||||
}
|
||||
|
||||
// Try to atomically swap from fromState to targetState
|
||||
if sm.state.CompareAndSwap(uint32(fromState), uint32(targetState)) {
|
||||
// Success! We transitioned atomically
|
||||
sm.notifyWaiters()
|
||||
return targetState, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fast path failed - check if we should wait or fail
|
||||
currentState := sm.GetState()
|
||||
|
||||
// Check if closed
|
||||
if currentState == StateClosed {
|
||||
return currentState, ErrStateMachineClosed
|
||||
}
|
||||
|
||||
// Slow path: need to wait for state change
|
||||
// Create waiter with valid states map for fast lookup
|
||||
validStatesMap := make(map[ConnState]struct{}, len(validFromStates))
|
||||
for _, s := range validFromStates {
|
||||
validStatesMap[s] = struct{}{}
|
||||
}
|
||||
|
||||
w := &waiter{
|
||||
validStates: validStatesMap,
|
||||
targetState: targetState,
|
||||
done: make(chan error, 1), // Buffered to avoid goroutine leak
|
||||
}
|
||||
|
||||
// Add to FIFO queue
|
||||
sm.mu.Lock()
|
||||
elem := sm.waiters.PushBack(w)
|
||||
sm.waiterCount.Add(1)
|
||||
sm.mu.Unlock()
|
||||
|
||||
// Wait for state change or timeout
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
// Timeout or cancellation - remove from queue
|
||||
sm.mu.Lock()
|
||||
sm.waiters.Remove(elem)
|
||||
sm.waiterCount.Add(-1)
|
||||
sm.mu.Unlock()
|
||||
return sm.GetState(), ctx.Err()
|
||||
case err := <-w.done:
|
||||
// Transition completed (or failed)
|
||||
// Note: waiterCount is decremented either in notifyWaiters (when the waiter is notified and removed)
|
||||
// or here (on timeout/cancellation).
|
||||
return sm.GetState(), err
|
||||
}
|
||||
}
|
||||
|
||||
// notifyWaiters checks if any waiters can proceed and notifies them in FIFO order.
|
||||
// This is called after every state transition.
|
||||
func (sm *ConnStateMachine) notifyWaiters() {
|
||||
// Fast path: check atomic counter without acquiring lock
|
||||
// This eliminates mutex overhead in the common case (no waiters)
|
||||
if sm.waiterCount.Load() == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
sm.mu.Lock()
|
||||
defer sm.mu.Unlock()
|
||||
|
||||
// Double-check after acquiring lock (waiters might have been processed)
|
||||
if sm.waiters.Len() == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// Track state locally so we only consider transitions made within this
|
||||
// call, not concurrent transitions from woken goroutines. Re-reading the
|
||||
// atomic would let a fast goroutine's Transition(StateIdle) leak into our
|
||||
// view, causing us to wake multiple waiters at once and breaking FIFO
|
||||
// execution ordering.
|
||||
currentState := sm.GetState()
|
||||
|
||||
for {
|
||||
processed := false
|
||||
|
||||
for elem := sm.waiters.Front(); elem != nil; elem = elem.Next() {
|
||||
w := elem.Value.(*waiter)
|
||||
|
||||
if _, valid := w.validStates[currentState]; valid {
|
||||
sm.waiters.Remove(elem)
|
||||
sm.waiterCount.Add(-1)
|
||||
|
||||
if sm.state.CompareAndSwap(uint32(currentState), uint32(w.targetState)) {
|
||||
w.done <- nil
|
||||
currentState = w.targetState
|
||||
processed = true
|
||||
break
|
||||
} else {
|
||||
sm.waiters.PushFront(w)
|
||||
sm.waiterCount.Add(1)
|
||||
currentState = sm.GetState()
|
||||
processed = true
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !processed {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
+165
@@ -0,0 +1,165 @@
|
||||
package pool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// PoolHook defines the interface for connection lifecycle hooks.
|
||||
type PoolHook interface {
|
||||
// OnGet is called when a connection is retrieved from the pool.
|
||||
// It can modify the connection or return an error to prevent its use.
|
||||
// The accept flag can be used to prevent the connection from being used.
|
||||
// On Accept = false the connection is rejected and returned to the pool.
|
||||
// The error can be used to prevent the connection from being used and returned to the pool.
|
||||
// On Errors, the connection is removed from the pool.
|
||||
// It has isNewConn flag to indicate if this is a new connection (rather than idle from the pool)
|
||||
// The flag can be used for gathering metrics on pool hit/miss ratio.
|
||||
OnGet(ctx context.Context, conn *Conn, isNewConn bool) (accept bool, err error)
|
||||
|
||||
// OnPut is called when a connection is returned to the pool.
|
||||
// It returns whether the connection should be pooled and whether it should be removed.
|
||||
OnPut(ctx context.Context, conn *Conn) (shouldPool bool, shouldRemove bool, err error)
|
||||
|
||||
// OnRemove is called when a connection is removed from the pool.
|
||||
// This happens when:
|
||||
// - Connection fails health check
|
||||
// - Connection exceeds max lifetime
|
||||
// - Pool is being closed
|
||||
// - Connection encounters an error
|
||||
// Implementations should clean up any per-connection state.
|
||||
// The reason parameter indicates why the connection was removed.
|
||||
OnRemove(ctx context.Context, conn *Conn, reason error)
|
||||
}
|
||||
|
||||
// PoolHookManager manages multiple pool hooks.
|
||||
type PoolHookManager struct {
|
||||
hooks []PoolHook
|
||||
hooksMu sync.RWMutex
|
||||
}
|
||||
|
||||
// NewPoolHookManager creates a new pool hook manager.
|
||||
func NewPoolHookManager() *PoolHookManager {
|
||||
return &PoolHookManager{
|
||||
hooks: make([]PoolHook, 0),
|
||||
}
|
||||
}
|
||||
|
||||
// AddHook adds a pool hook to the manager.
|
||||
// Hooks are called in the order they were added.
|
||||
func (phm *PoolHookManager) AddHook(hook PoolHook) {
|
||||
phm.hooksMu.Lock()
|
||||
defer phm.hooksMu.Unlock()
|
||||
phm.hooks = append(phm.hooks, hook)
|
||||
}
|
||||
|
||||
// RemoveHook removes a pool hook from the manager.
|
||||
func (phm *PoolHookManager) RemoveHook(hook PoolHook) {
|
||||
phm.hooksMu.Lock()
|
||||
defer phm.hooksMu.Unlock()
|
||||
|
||||
for i, h := range phm.hooks {
|
||||
if h == hook {
|
||||
// Remove hook by swapping with last element and truncating
|
||||
phm.hooks[i] = phm.hooks[len(phm.hooks)-1]
|
||||
phm.hooks = phm.hooks[:len(phm.hooks)-1]
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ProcessOnGet calls all OnGet hooks in order.
|
||||
// If any hook returns an error, processing stops and the error is returned.
|
||||
func (phm *PoolHookManager) ProcessOnGet(ctx context.Context, conn *Conn, isNewConn bool) (acceptConn bool, err error) {
|
||||
// Copy slice reference while holding lock (fast)
|
||||
phm.hooksMu.RLock()
|
||||
hooks := phm.hooks
|
||||
phm.hooksMu.RUnlock()
|
||||
|
||||
// Call hooks without holding lock (slow operations)
|
||||
for _, hook := range hooks {
|
||||
acceptConn, err := hook.OnGet(ctx, conn, isNewConn)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
if !acceptConn {
|
||||
return false, nil
|
||||
}
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// ProcessOnPut calls all OnPut hooks in order.
|
||||
// The first hook that returns shouldRemove=true or shouldPool=false will stop processing.
|
||||
func (phm *PoolHookManager) ProcessOnPut(ctx context.Context, conn *Conn) (shouldPool bool, shouldRemove bool, err error) {
|
||||
// Copy slice reference while holding lock (fast)
|
||||
phm.hooksMu.RLock()
|
||||
hooks := phm.hooks
|
||||
phm.hooksMu.RUnlock()
|
||||
|
||||
shouldPool = true // Default to pooling the connection
|
||||
|
||||
// Call hooks without holding lock (slow operations)
|
||||
for _, hook := range hooks {
|
||||
hookShouldPool, hookShouldRemove, hookErr := hook.OnPut(ctx, conn)
|
||||
|
||||
if hookErr != nil {
|
||||
return false, true, hookErr
|
||||
}
|
||||
|
||||
// If any hook says to remove or not pool, respect that decision
|
||||
if hookShouldRemove {
|
||||
return false, true, nil
|
||||
}
|
||||
|
||||
if !hookShouldPool {
|
||||
shouldPool = false
|
||||
}
|
||||
}
|
||||
|
||||
return shouldPool, false, nil
|
||||
}
|
||||
|
||||
// ProcessOnRemove calls all OnRemove hooks in order.
|
||||
func (phm *PoolHookManager) ProcessOnRemove(ctx context.Context, conn *Conn, reason error) {
|
||||
// Copy slice reference while holding lock (fast)
|
||||
phm.hooksMu.RLock()
|
||||
hooks := phm.hooks
|
||||
phm.hooksMu.RUnlock()
|
||||
|
||||
// Call hooks without holding lock (slow operations)
|
||||
for _, hook := range hooks {
|
||||
hook.OnRemove(ctx, conn, reason)
|
||||
}
|
||||
}
|
||||
|
||||
// GetHookCount returns the number of registered hooks (for testing).
|
||||
func (phm *PoolHookManager) GetHookCount() int {
|
||||
phm.hooksMu.RLock()
|
||||
defer phm.hooksMu.RUnlock()
|
||||
return len(phm.hooks)
|
||||
}
|
||||
|
||||
// GetHooks returns a copy of all registered hooks.
|
||||
func (phm *PoolHookManager) GetHooks() []PoolHook {
|
||||
phm.hooksMu.RLock()
|
||||
defer phm.hooksMu.RUnlock()
|
||||
|
||||
hooks := make([]PoolHook, len(phm.hooks))
|
||||
copy(hooks, phm.hooks)
|
||||
return hooks
|
||||
}
|
||||
|
||||
// Clone creates a copy of the hook manager with the same hooks.
|
||||
// This is used for lock-free atomic updates of the hook manager.
|
||||
func (phm *PoolHookManager) Clone() *PoolHookManager {
|
||||
phm.hooksMu.RLock()
|
||||
defer phm.hooksMu.RUnlock()
|
||||
|
||||
newManager := &PoolHookManager{
|
||||
hooks: make([]PoolHook, len(phm.hooks)),
|
||||
}
|
||||
copy(newManager.hooks, phm.hooks)
|
||||
return newManager
|
||||
}
|
||||
+1697
File diff suppressed because it is too large
Load Diff
+104
@@ -0,0 +1,104 @@
|
||||
package pool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
)
|
||||
|
||||
// SingleConnPool is a pool that always returns the same connection.
|
||||
// Note: This pool is not thread-safe.
|
||||
// It is intended to be used by clients that need a single connection.
|
||||
type SingleConnPool struct {
|
||||
pool Pooler
|
||||
cn *Conn
|
||||
stickyErr error
|
||||
}
|
||||
|
||||
var _ Pooler = (*SingleConnPool)(nil)
|
||||
|
||||
// NewSingleConnPool creates a new single connection pool.
|
||||
// The pool will always return the same connection.
|
||||
// The pool will not:
|
||||
// - Close the connection
|
||||
// - Reconnect the connection
|
||||
// - Track the connection in any way
|
||||
func NewSingleConnPool(pool Pooler, cn *Conn) *SingleConnPool {
|
||||
return &SingleConnPool{
|
||||
pool: pool,
|
||||
cn: cn,
|
||||
}
|
||||
}
|
||||
|
||||
func (p *SingleConnPool) NewConn(ctx context.Context) (*Conn, error) {
|
||||
return p.pool.NewConn(ctx)
|
||||
}
|
||||
|
||||
func (p *SingleConnPool) CloseConn(ctx context.Context, cn *Conn, reason string, fromState string) error {
|
||||
return p.pool.CloseConn(ctx, cn, reason, fromState)
|
||||
}
|
||||
|
||||
func (p *SingleConnPool) Get(_ context.Context) (*Conn, error) {
|
||||
if p.stickyErr != nil {
|
||||
return nil, p.stickyErr
|
||||
}
|
||||
if p.cn == nil {
|
||||
return nil, ErrClosed
|
||||
}
|
||||
|
||||
// NOTE: SingleConnPool is NOT thread-safe by design and is used in special scenarios:
|
||||
// - During initialization (connection is in INITIALIZING state)
|
||||
// - During re-authentication (connection is in UNUSABLE state)
|
||||
// - For transactions (connection might be in various states)
|
||||
// We use SetUsed() which forces the transition, rather than TryTransition() which
|
||||
// would fail if the connection is not in IDLE/CREATED state.
|
||||
p.cn.SetUsed(true)
|
||||
p.cn.SetUsedAt(time.Now())
|
||||
return p.cn, nil
|
||||
}
|
||||
|
||||
func (p *SingleConnPool) Put(_ context.Context, cn *Conn) {
|
||||
if p.cn == nil {
|
||||
return
|
||||
}
|
||||
if p.cn != cn {
|
||||
return
|
||||
}
|
||||
p.cn.SetUsed(false)
|
||||
}
|
||||
|
||||
func (p *SingleConnPool) Remove(_ context.Context, cn *Conn, reason error) {
|
||||
cn.SetUsed(false)
|
||||
p.cn = nil
|
||||
p.stickyErr = reason
|
||||
}
|
||||
|
||||
// RemoveWithoutTurn has the same behavior as Remove for SingleConnPool
|
||||
// since SingleConnPool doesn't use a turn-based queue system.
|
||||
func (p *SingleConnPool) RemoveWithoutTurn(ctx context.Context, cn *Conn, reason error) {
|
||||
p.Remove(ctx, cn, reason)
|
||||
}
|
||||
|
||||
func (p *SingleConnPool) Close() error {
|
||||
p.cn = nil
|
||||
p.stickyErr = ErrClosed
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *SingleConnPool) Len() int {
|
||||
return 0
|
||||
}
|
||||
|
||||
func (p *SingleConnPool) IdleLen() int {
|
||||
return 0
|
||||
}
|
||||
|
||||
// Size returns the maximum pool size, which is always 1 for SingleConnPool.
|
||||
func (p *SingleConnPool) Size() int { return 1 }
|
||||
|
||||
func (p *SingleConnPool) Stats() *Stats {
|
||||
return &Stats{}
|
||||
}
|
||||
|
||||
func (p *SingleConnPool) AddPoolHook(_ PoolHook) {}
|
||||
|
||||
func (p *SingleConnPool) RemovePoolHook(_ PoolHook) {}
|
||||
+214
@@ -0,0 +1,214 @@
|
||||
package pool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
const (
|
||||
stateDefault = 0
|
||||
stateInited = 1
|
||||
stateClosed = 2
|
||||
)
|
||||
|
||||
type BadConnError struct {
|
||||
wrapped error
|
||||
}
|
||||
|
||||
var _ error = (*BadConnError)(nil)
|
||||
|
||||
func (e BadConnError) Error() string {
|
||||
s := "redis: Conn is in a bad state"
|
||||
if e.wrapped != nil {
|
||||
s += ": " + e.wrapped.Error()
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func (e BadConnError) Unwrap() error {
|
||||
return e.wrapped
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------------
|
||||
|
||||
type StickyConnPool struct {
|
||||
pool Pooler
|
||||
shared int32 // atomic
|
||||
|
||||
state uint32 // atomic
|
||||
ch chan *Conn
|
||||
|
||||
_badConnError atomic.Value
|
||||
}
|
||||
|
||||
var _ Pooler = (*StickyConnPool)(nil)
|
||||
|
||||
func NewStickyConnPool(pool Pooler) *StickyConnPool {
|
||||
p, ok := pool.(*StickyConnPool)
|
||||
if !ok {
|
||||
p = &StickyConnPool{
|
||||
pool: pool,
|
||||
ch: make(chan *Conn, 1),
|
||||
}
|
||||
}
|
||||
atomic.AddInt32(&p.shared, 1)
|
||||
return p
|
||||
}
|
||||
|
||||
func (p *StickyConnPool) NewConn(ctx context.Context) (*Conn, error) {
|
||||
return p.pool.NewConn(ctx)
|
||||
}
|
||||
|
||||
func (p *StickyConnPool) CloseConn(ctx context.Context, cn *Conn, reason string, fromState string) error {
|
||||
return p.pool.CloseConn(ctx, cn, reason, fromState)
|
||||
}
|
||||
|
||||
func (p *StickyConnPool) Get(ctx context.Context) (*Conn, error) {
|
||||
// In worst case this races with Close which is not a very common operation.
|
||||
for i := 0; i < 1000; i++ {
|
||||
switch atomic.LoadUint32(&p.state) {
|
||||
case stateDefault:
|
||||
cn, err := p.pool.Get(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if atomic.CompareAndSwapUint32(&p.state, stateDefault, stateInited) {
|
||||
return cn, nil
|
||||
}
|
||||
p.pool.Remove(ctx, cn, ErrClosed)
|
||||
case stateInited:
|
||||
if err := p.badConnError(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
cn, ok := <-p.ch
|
||||
if !ok {
|
||||
return nil, ErrClosed
|
||||
}
|
||||
return cn, nil
|
||||
case stateClosed:
|
||||
return nil, ErrClosed
|
||||
default:
|
||||
panic("not reached")
|
||||
}
|
||||
}
|
||||
return nil, fmt.Errorf("redis: StickyConnPool.Get: infinite loop")
|
||||
}
|
||||
|
||||
func (p *StickyConnPool) Put(ctx context.Context, cn *Conn) {
|
||||
defer func() {
|
||||
if recover() != nil {
|
||||
p.freeConn(ctx, cn)
|
||||
}
|
||||
}()
|
||||
p.ch <- cn
|
||||
}
|
||||
|
||||
func (p *StickyConnPool) freeConn(ctx context.Context, cn *Conn) {
|
||||
if err := p.badConnError(); err != nil {
|
||||
p.pool.Remove(ctx, cn, err)
|
||||
} else {
|
||||
p.pool.Put(ctx, cn)
|
||||
}
|
||||
}
|
||||
|
||||
func (p *StickyConnPool) Remove(ctx context.Context, cn *Conn, reason error) {
|
||||
defer func() {
|
||||
if recover() != nil {
|
||||
p.pool.Remove(ctx, cn, ErrClosed)
|
||||
}
|
||||
}()
|
||||
p._badConnError.Store(BadConnError{wrapped: reason})
|
||||
p.ch <- cn
|
||||
}
|
||||
|
||||
// RemoveWithoutTurn has the same behavior as Remove for StickyConnPool
|
||||
// since StickyConnPool doesn't use a turn-based queue system.
|
||||
func (p *StickyConnPool) RemoveWithoutTurn(ctx context.Context, cn *Conn, reason error) {
|
||||
p.Remove(ctx, cn, reason)
|
||||
}
|
||||
|
||||
func (p *StickyConnPool) Close() error {
|
||||
if shared := atomic.AddInt32(&p.shared, -1); shared > 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
for i := 0; i < 1000; i++ {
|
||||
state := atomic.LoadUint32(&p.state)
|
||||
if state == stateClosed {
|
||||
return ErrClosed
|
||||
}
|
||||
if atomic.CompareAndSwapUint32(&p.state, state, stateClosed) {
|
||||
close(p.ch)
|
||||
cn, ok := <-p.ch
|
||||
if ok {
|
||||
p.freeConn(context.TODO(), cn)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
return errors.New("redis: StickyConnPool.Close: infinite loop")
|
||||
}
|
||||
|
||||
func (p *StickyConnPool) Reset(ctx context.Context) error {
|
||||
if p.badConnError() == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
select {
|
||||
case cn, ok := <-p.ch:
|
||||
if !ok {
|
||||
return ErrClosed
|
||||
}
|
||||
p.pool.Remove(ctx, cn, ErrClosed)
|
||||
p._badConnError.Store(BadConnError{wrapped: nil})
|
||||
default:
|
||||
return errors.New("redis: StickyConnPool does not have a Conn")
|
||||
}
|
||||
|
||||
if !atomic.CompareAndSwapUint32(&p.state, stateInited, stateDefault) {
|
||||
state := atomic.LoadUint32(&p.state)
|
||||
return fmt.Errorf("redis: invalid StickyConnPool state: %d", state)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *StickyConnPool) badConnError() error {
|
||||
if v := p._badConnError.Load(); v != nil {
|
||||
if err := v.(BadConnError); err.wrapped != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *StickyConnPool) Len() int {
|
||||
switch atomic.LoadUint32(&p.state) {
|
||||
case stateDefault:
|
||||
return 0
|
||||
case stateInited:
|
||||
return 1
|
||||
case stateClosed:
|
||||
return 0
|
||||
default:
|
||||
panic("not reached")
|
||||
}
|
||||
}
|
||||
|
||||
func (p *StickyConnPool) IdleLen() int {
|
||||
return len(p.ch)
|
||||
}
|
||||
|
||||
// Size returns the maximum pool size, which is always 1 for StickyConnPool.
|
||||
func (p *StickyConnPool) Size() int { return 1 }
|
||||
|
||||
func (p *StickyConnPool) Stats() *Stats {
|
||||
return &Stats{}
|
||||
}
|
||||
|
||||
func (p *StickyConnPool) AddPoolHook(hook PoolHook) {}
|
||||
|
||||
func (p *StickyConnPool) RemovePoolHook(hook PoolHook) {}
|
||||
+105
@@ -0,0 +1,105 @@
|
||||
package pool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
type PubSubStats struct {
|
||||
Created uint32
|
||||
Untracked uint32
|
||||
Active uint32
|
||||
}
|
||||
|
||||
// PubSubPool manages a pool of PubSub connections.
|
||||
type PubSubPool struct {
|
||||
opt *Options
|
||||
netDialer func(ctx context.Context, network, addr string) (net.Conn, error)
|
||||
|
||||
// Map to track active PubSub connections
|
||||
activeConns sync.Map // map[uint64]*Conn (connID -> conn)
|
||||
closed atomic.Bool
|
||||
stats PubSubStats
|
||||
}
|
||||
|
||||
// NewPubSubPool implements a pool for PubSub connections.
|
||||
// It intentionally does not implement the Pooler interface
|
||||
func NewPubSubPool(opt *Options, netDialer func(ctx context.Context, network, addr string) (net.Conn, error)) *PubSubPool {
|
||||
return &PubSubPool{
|
||||
opt: opt,
|
||||
netDialer: netDialer,
|
||||
}
|
||||
}
|
||||
|
||||
func (p *PubSubPool) NewConn(ctx context.Context, network string, addr string, channels []string) (*Conn, error) {
|
||||
if p.closed.Load() {
|
||||
return nil, ErrClosed
|
||||
}
|
||||
|
||||
netConn, err := p.netDialer(ctx, network, addr)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
cn := NewConnWithBufferSize(netConn, p.opt.ReadBufferSize, p.opt.WriteBufferSize)
|
||||
cn.pubsub = true
|
||||
// Set pool name for metrics
|
||||
cn.SetPoolName(p.opt.Name)
|
||||
atomic.AddUint32(&p.stats.Created, 1)
|
||||
return cn, nil
|
||||
}
|
||||
|
||||
func (p *PubSubPool) TrackConn(cn *Conn) {
|
||||
atomic.AddUint32(&p.stats.Active, 1)
|
||||
p.activeConns.Store(cn.GetID(), cn)
|
||||
// Emit +1 used for PubSub connection
|
||||
if cb := getMetricConnectionCountCallback(); cb != nil {
|
||||
cb(context.Background(), 1, cn, "used", true)
|
||||
}
|
||||
}
|
||||
|
||||
func (p *PubSubPool) UntrackConn(cn *Conn) {
|
||||
// LoadAndDelete ensures each connection is only decremented once,
|
||||
// guarding against double-decrement if Close() already untracked it.
|
||||
if _, loaded := p.activeConns.LoadAndDelete(cn.GetID()); !loaded {
|
||||
return
|
||||
}
|
||||
atomic.AddUint32(&p.stats.Active, ^uint32(0))
|
||||
atomic.AddUint32(&p.stats.Untracked, 1)
|
||||
// Emit -1 used for PubSub connection
|
||||
if cb := getMetricConnectionCountCallback(); cb != nil {
|
||||
cb(context.Background(), -1, cn, "used", true)
|
||||
}
|
||||
}
|
||||
|
||||
func (p *PubSubPool) Close() error {
|
||||
p.closed.Store(true)
|
||||
cb := getMetricConnectionCountCallback()
|
||||
p.activeConns.Range(func(key, value interface{}) bool {
|
||||
cn := value.(*Conn)
|
||||
// Use LoadAndDelete to atomically claim ownership of this entry.
|
||||
// If a concurrent UntrackConn already removed it, skip to avoid double-decrement.
|
||||
if _, loaded := p.activeConns.LoadAndDelete(key); !loaded {
|
||||
return true
|
||||
}
|
||||
atomic.AddUint32(&p.stats.Active, ^uint32(0))
|
||||
atomic.AddUint32(&p.stats.Untracked, 1)
|
||||
// Emit -1 used for each PubSub connection being closed
|
||||
if cb != nil {
|
||||
cb(context.Background(), -1, cn, "used", true)
|
||||
}
|
||||
_ = cn.Close()
|
||||
return true
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *PubSubPool) Stats() *PubSubStats {
|
||||
// load stats atomically
|
||||
return &PubSubStats{
|
||||
Created: atomic.LoadUint32(&p.stats.Created),
|
||||
Untracked: atomic.LoadUint32(&p.stats.Untracked),
|
||||
Active: atomic.LoadUint32(&p.stats.Active),
|
||||
}
|
||||
}
|
||||
+115
@@ -0,0 +1,115 @@
|
||||
package pool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
)
|
||||
|
||||
type wantConn struct {
|
||||
mu sync.RWMutex // protects ctx, done and sending of the result
|
||||
ctx context.Context // context for dial, cleared after delivered or canceled
|
||||
cancelCtx context.CancelFunc
|
||||
done bool // true after delivered or canceled
|
||||
result chan wantConnResult // channel to deliver connection or error
|
||||
}
|
||||
|
||||
// getCtxForDial returns context for dial or nil if connection was delivered or canceled.
|
||||
func (w *wantConn) getCtxForDial() context.Context {
|
||||
w.mu.RLock()
|
||||
defer w.mu.RUnlock()
|
||||
|
||||
return w.ctx
|
||||
}
|
||||
|
||||
func (w *wantConn) tryDeliver(cn *Conn, err error) bool {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
if w.done {
|
||||
return false
|
||||
}
|
||||
|
||||
w.done = true
|
||||
w.ctx = nil
|
||||
|
||||
w.result <- wantConnResult{cn: cn, err: err}
|
||||
close(w.result)
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
func (w *wantConn) cancel() *Conn {
|
||||
w.mu.Lock()
|
||||
var cn *Conn
|
||||
if w.done {
|
||||
select {
|
||||
case result := <-w.result:
|
||||
cn = result.cn
|
||||
default:
|
||||
}
|
||||
} else {
|
||||
close(w.result)
|
||||
}
|
||||
|
||||
w.done = true
|
||||
w.ctx = nil
|
||||
w.mu.Unlock()
|
||||
|
||||
return cn
|
||||
}
|
||||
|
||||
func (w *wantConn) isOngoing() bool {
|
||||
w.mu.RLock()
|
||||
defer w.mu.RUnlock()
|
||||
return !w.done
|
||||
}
|
||||
|
||||
type wantConnResult struct {
|
||||
cn *Conn
|
||||
err error
|
||||
}
|
||||
|
||||
type wantConnQueue struct {
|
||||
mu sync.RWMutex
|
||||
items []*wantConn
|
||||
}
|
||||
|
||||
func newWantConnQueue() *wantConnQueue {
|
||||
return &wantConnQueue{
|
||||
items: make([]*wantConn, 0),
|
||||
}
|
||||
}
|
||||
|
||||
func (q *wantConnQueue) enqueue(w *wantConn) {
|
||||
q.mu.Lock()
|
||||
defer q.mu.Unlock()
|
||||
q.items = append(q.items, w)
|
||||
}
|
||||
|
||||
func (q *wantConnQueue) dequeue() (*wantConn, bool) {
|
||||
q.mu.Lock()
|
||||
defer q.mu.Unlock()
|
||||
|
||||
if len(q.items) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
item := q.items[0]
|
||||
q.items = q.items[1:]
|
||||
return item, true
|
||||
}
|
||||
|
||||
func (q *wantConnQueue) discardDoneAtFront() int {
|
||||
q.mu.Lock()
|
||||
defer q.mu.Unlock()
|
||||
count := 0
|
||||
for len(q.items) > 0 {
|
||||
if q.items[0].isOngoing() {
|
||||
break
|
||||
}
|
||||
|
||||
q.items = q.items[1:]
|
||||
count++
|
||||
}
|
||||
|
||||
return count
|
||||
}
|
||||
Reference in New Issue
Block a user