fix(go.sum): update ResolveSpec dependency to v1.0.87
This commit is contained in:
+58
@@ -0,0 +1,58 @@
|
||||
package internal
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9/internal/util"
|
||||
)
|
||||
|
||||
func AppendArg(b []byte, v interface{}) []byte {
|
||||
switch v := v.(type) {
|
||||
case nil:
|
||||
return append(b, "<nil>"...)
|
||||
case string:
|
||||
return appendUTF8String(b, util.StringToBytes(v))
|
||||
case []byte:
|
||||
return appendUTF8String(b, v)
|
||||
case int:
|
||||
return strconv.AppendInt(b, int64(v), 10)
|
||||
case int8:
|
||||
return strconv.AppendInt(b, int64(v), 10)
|
||||
case int16:
|
||||
return strconv.AppendInt(b, int64(v), 10)
|
||||
case int32:
|
||||
return strconv.AppendInt(b, int64(v), 10)
|
||||
case int64:
|
||||
return strconv.AppendInt(b, v, 10)
|
||||
case uint:
|
||||
return strconv.AppendUint(b, uint64(v), 10)
|
||||
case uint8:
|
||||
return strconv.AppendUint(b, uint64(v), 10)
|
||||
case uint16:
|
||||
return strconv.AppendUint(b, uint64(v), 10)
|
||||
case uint32:
|
||||
return strconv.AppendUint(b, uint64(v), 10)
|
||||
case uint64:
|
||||
return strconv.AppendUint(b, v, 10)
|
||||
case float32:
|
||||
return strconv.AppendFloat(b, float64(v), 'f', -1, 64)
|
||||
case float64:
|
||||
return strconv.AppendFloat(b, v, 'f', -1, 64)
|
||||
case bool:
|
||||
if v {
|
||||
return append(b, "true"...)
|
||||
}
|
||||
return append(b, "false"...)
|
||||
case time.Time:
|
||||
return v.AppendFormat(b, time.RFC3339Nano)
|
||||
default:
|
||||
return append(b, fmt.Sprint(v)...)
|
||||
}
|
||||
}
|
||||
|
||||
func appendUTF8String(dst []byte, src []byte) []byte {
|
||||
dst = append(dst, src...)
|
||||
return dst
|
||||
}
|
||||
Generated
Vendored
+100
@@ -0,0 +1,100 @@
|
||||
package streaming
|
||||
|
||||
import (
|
||||
"github.com/redis/go-redis/v9/auth"
|
||||
"github.com/redis/go-redis/v9/internal/pool"
|
||||
)
|
||||
|
||||
// ConnReAuthCredentialsListener is a credentials listener for a specific connection
|
||||
// that triggers re-authentication when credentials change.
|
||||
//
|
||||
// This listener implements the auth.CredentialsListener interface and is subscribed
|
||||
// to a StreamingCredentialsProvider. When new credentials are received via OnNext,
|
||||
// it marks the connection for re-authentication through the manager.
|
||||
//
|
||||
// The re-authentication is always performed asynchronously to avoid blocking the
|
||||
// credentials provider and to prevent potential deadlocks with the pool semaphore.
|
||||
// The actual re-auth happens when the connection is returned to the pool in an idle state.
|
||||
//
|
||||
// Lifecycle:
|
||||
// - Created during connection initialization via Manager.Listener()
|
||||
// - Subscribed to the StreamingCredentialsProvider
|
||||
// - Receives credential updates via OnNext()
|
||||
// - Cleaned up when connection is removed from pool via Manager.RemoveListener()
|
||||
type ConnReAuthCredentialsListener struct {
|
||||
// reAuth is the function to re-authenticate the connection with new credentials
|
||||
reAuth func(conn *pool.Conn, credentials auth.Credentials) error
|
||||
|
||||
// onErr is the function to call when re-authentication or acquisition fails
|
||||
onErr func(conn *pool.Conn, err error)
|
||||
|
||||
// conn is the connection this listener is associated with
|
||||
conn *pool.Conn
|
||||
|
||||
// manager is the streaming credentials manager for coordinating re-auth
|
||||
manager *Manager
|
||||
}
|
||||
|
||||
// OnNext is called when new credentials are received from the StreamingCredentialsProvider.
|
||||
//
|
||||
// This method marks the connection for asynchronous re-authentication. The actual
|
||||
// re-authentication happens in the background when the connection is returned to the
|
||||
// pool and is in an idle state.
|
||||
//
|
||||
// Asynchronous re-auth is used to:
|
||||
// - Avoid blocking the credentials provider's notification goroutine
|
||||
// - Prevent deadlocks with the pool's semaphore (especially with small pool sizes)
|
||||
// - Ensure re-auth happens when the connection is safe to use (not processing commands)
|
||||
//
|
||||
// The reAuthFn callback receives:
|
||||
// - nil if the connection was successfully acquired for re-auth
|
||||
// - error if acquisition timed out or failed
|
||||
//
|
||||
// Thread-safe: Called by the credentials provider's notification goroutine.
|
||||
func (c *ConnReAuthCredentialsListener) OnNext(credentials auth.Credentials) {
|
||||
if c.conn == nil || c.conn.IsClosed() || c.manager == nil || c.reAuth == nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Always use async reauth to avoid complex pool semaphore issues
|
||||
// The synchronous path can cause deadlocks in the pool's semaphore mechanism
|
||||
// when called from the Subscribe goroutine, especially with small pool sizes.
|
||||
// The connection pool hook will re-authenticate the connection when it is
|
||||
// returned to the pool in a clean, idle state.
|
||||
c.manager.MarkForReAuth(c.conn, func(err error) {
|
||||
// err is from connection acquisition (timeout, etc.)
|
||||
if err != nil {
|
||||
// Log the error
|
||||
c.OnError(err)
|
||||
return
|
||||
}
|
||||
// err is from reauth command execution
|
||||
err = c.reAuth(c.conn, credentials)
|
||||
if err != nil {
|
||||
// Log the error
|
||||
c.OnError(err)
|
||||
return
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// OnError is called when an error occurs during credential streaming or re-authentication.
|
||||
//
|
||||
// This method can be called from:
|
||||
// - The StreamingCredentialsProvider when there's an error in the credentials stream
|
||||
// - The re-auth process when connection acquisition times out
|
||||
// - The re-auth process when the AUTH command fails
|
||||
//
|
||||
// The error is delegated to the onErr callback provided during listener creation.
|
||||
//
|
||||
// Thread-safe: Can be called from multiple goroutines (provider, re-auth worker).
|
||||
func (c *ConnReAuthCredentialsListener) OnError(err error) {
|
||||
if c.onErr == nil {
|
||||
return
|
||||
}
|
||||
|
||||
c.onErr(c.conn, err)
|
||||
}
|
||||
|
||||
// Ensure ConnReAuthCredentialsListener implements the CredentialsListener interface.
|
||||
var _ auth.CredentialsListener = (*ConnReAuthCredentialsListener)(nil)
|
||||
+77
@@ -0,0 +1,77 @@
|
||||
package streaming
|
||||
|
||||
import (
|
||||
"sync"
|
||||
|
||||
"github.com/redis/go-redis/v9/auth"
|
||||
)
|
||||
|
||||
// CredentialsListeners is a thread-safe collection of credentials listeners
|
||||
// indexed by connection ID.
|
||||
//
|
||||
// This collection is used by the Manager to maintain a registry of listeners
|
||||
// for each connection in the pool. Listeners are reused when connections are
|
||||
// reinitialized (e.g., after a handoff) to avoid creating duplicate subscriptions
|
||||
// to the StreamingCredentialsProvider.
|
||||
//
|
||||
// The collection supports concurrent access from multiple goroutines during
|
||||
// connection initialization, credential updates, and connection removal.
|
||||
type CredentialsListeners struct {
|
||||
// listeners maps connection ID to credentials listener
|
||||
listeners map[uint64]auth.CredentialsListener
|
||||
|
||||
// lock protects concurrent access to the listeners map
|
||||
lock sync.RWMutex
|
||||
}
|
||||
|
||||
// NewCredentialsListeners creates a new thread-safe credentials listeners collection.
|
||||
func NewCredentialsListeners() *CredentialsListeners {
|
||||
return &CredentialsListeners{
|
||||
listeners: make(map[uint64]auth.CredentialsListener),
|
||||
}
|
||||
}
|
||||
|
||||
// Add adds or updates a credentials listener for a connection.
|
||||
//
|
||||
// If a listener already exists for the connection ID, it is replaced.
|
||||
// This is safe because the old listener should have been unsubscribed
|
||||
// before the connection was reinitialized.
|
||||
//
|
||||
// Thread-safe: Can be called concurrently from multiple goroutines.
|
||||
func (c *CredentialsListeners) Add(connID uint64, listener auth.CredentialsListener) {
|
||||
c.lock.Lock()
|
||||
defer c.lock.Unlock()
|
||||
if c.listeners == nil {
|
||||
c.listeners = make(map[uint64]auth.CredentialsListener)
|
||||
}
|
||||
c.listeners[connID] = listener
|
||||
}
|
||||
|
||||
// Get retrieves the credentials listener for a connection.
|
||||
//
|
||||
// Returns:
|
||||
// - listener: The credentials listener for the connection, or nil if not found
|
||||
// - ok: true if a listener exists for the connection ID, false otherwise
|
||||
//
|
||||
// Thread-safe: Can be called concurrently from multiple goroutines.
|
||||
func (c *CredentialsListeners) Get(connID uint64) (auth.CredentialsListener, bool) {
|
||||
c.lock.RLock()
|
||||
defer c.lock.RUnlock()
|
||||
if len(c.listeners) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
listener, ok := c.listeners[connID]
|
||||
return listener, ok
|
||||
}
|
||||
|
||||
// Remove removes the credentials listener for a connection.
|
||||
//
|
||||
// This is called when a connection is removed from the pool to prevent
|
||||
// memory leaks. If no listener exists for the connection ID, this is a no-op.
|
||||
//
|
||||
// Thread-safe: Can be called concurrently from multiple goroutines.
|
||||
func (c *CredentialsListeners) Remove(connID uint64) {
|
||||
c.lock.Lock()
|
||||
defer c.lock.Unlock()
|
||||
delete(c.listeners, connID)
|
||||
}
|
||||
+137
@@ -0,0 +1,137 @@
|
||||
package streaming
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9/auth"
|
||||
"github.com/redis/go-redis/v9/internal/pool"
|
||||
)
|
||||
|
||||
// Manager coordinates streaming credentials and re-authentication for a connection pool.
|
||||
//
|
||||
// The manager is responsible for:
|
||||
// - Creating and managing per-connection credentials listeners
|
||||
// - Providing the pool hook for re-authentication
|
||||
// - Coordinating between credentials updates and pool operations
|
||||
//
|
||||
// When credentials change via a StreamingCredentialsProvider:
|
||||
// 1. The credentials listener (ConnReAuthCredentialsListener) receives the update
|
||||
// 2. It calls MarkForReAuth on the manager
|
||||
// 3. The manager delegates to the pool hook
|
||||
// 4. The pool hook schedules background re-authentication
|
||||
//
|
||||
// The manager maintains a registry of credentials listeners indexed by connection ID,
|
||||
// allowing listener reuse when connections are reinitialized (e.g., after handoff).
|
||||
type Manager struct {
|
||||
// credentialsListeners maps connection ID to credentials listener
|
||||
credentialsListeners *CredentialsListeners
|
||||
|
||||
// pool is the connection pool being managed
|
||||
pool pool.Pooler
|
||||
|
||||
// poolHookRef is the re-authentication pool hook
|
||||
poolHookRef *ReAuthPoolHook
|
||||
}
|
||||
|
||||
// NewManager creates a new streaming credentials manager.
|
||||
//
|
||||
// Parameters:
|
||||
// - pl: The connection pool to manage
|
||||
// - reAuthTimeout: Maximum time to wait for acquiring a connection for re-authentication
|
||||
//
|
||||
// The manager creates a ReAuthPoolHook sized to match the pool size, ensuring that
|
||||
// re-auth operations don't exhaust the connection pool.
|
||||
func NewManager(pl pool.Pooler, reAuthTimeout time.Duration) *Manager {
|
||||
m := &Manager{
|
||||
pool: pl,
|
||||
poolHookRef: NewReAuthPoolHook(pl.Size(), reAuthTimeout),
|
||||
credentialsListeners: NewCredentialsListeners(),
|
||||
}
|
||||
m.poolHookRef.manager = m
|
||||
return m
|
||||
}
|
||||
|
||||
// PoolHook returns the pool hook for re-authentication.
|
||||
//
|
||||
// This hook should be registered with the connection pool to enable
|
||||
// automatic re-authentication when credentials change.
|
||||
func (m *Manager) PoolHook() pool.PoolHook {
|
||||
return m.poolHookRef
|
||||
}
|
||||
|
||||
// Listener returns or creates a credentials listener for a connection.
|
||||
//
|
||||
// This method is called during connection initialization to set up the
|
||||
// credentials listener. If a listener already exists for the connection ID
|
||||
// (e.g., after a handoff), it is reused.
|
||||
//
|
||||
// Parameters:
|
||||
// - poolCn: The connection to create/get a listener for
|
||||
// - reAuth: Function to re-authenticate the connection with new credentials
|
||||
// - onErr: Function to call when re-authentication fails
|
||||
//
|
||||
// Returns:
|
||||
// - auth.CredentialsListener: The listener to subscribe to the credentials provider
|
||||
// - error: Non-nil if poolCn is nil
|
||||
//
|
||||
// Note: The reAuth and onErr callbacks are captured once when the listener is
|
||||
// created and reused for the connection's lifetime. They should not change.
|
||||
//
|
||||
// Thread-safe: Can be called concurrently during connection initialization.
|
||||
func (m *Manager) Listener(
|
||||
poolCn *pool.Conn,
|
||||
reAuth func(*pool.Conn, auth.Credentials) error,
|
||||
onErr func(*pool.Conn, error),
|
||||
) (auth.CredentialsListener, error) {
|
||||
if poolCn == nil {
|
||||
return nil, errors.New("poolCn cannot be nil")
|
||||
}
|
||||
connID := poolCn.GetID()
|
||||
// if we reconnect the underlying network connection, the streaming credentials listener will continue to work
|
||||
// so we can get the old listener from the cache and use it.
|
||||
// subscribing the same (an already subscribed) listener for a StreamingCredentialsProvider SHOULD be a no-op
|
||||
listener, ok := m.credentialsListeners.Get(connID)
|
||||
if !ok || listener == nil {
|
||||
// Create new listener for this connection
|
||||
// Note: Callbacks (reAuth, onErr) are captured once and reused for the connection's lifetime
|
||||
newCredListener := &ConnReAuthCredentialsListener{
|
||||
conn: poolCn,
|
||||
reAuth: reAuth,
|
||||
onErr: onErr,
|
||||
manager: m,
|
||||
}
|
||||
|
||||
m.credentialsListeners.Add(connID, newCredListener)
|
||||
listener = newCredListener
|
||||
}
|
||||
return listener, nil
|
||||
}
|
||||
|
||||
// MarkForReAuth marks a connection for re-authentication.
|
||||
//
|
||||
// This method is called by the credentials listener when new credentials are
|
||||
// received. It delegates to the pool hook to schedule background re-authentication.
|
||||
//
|
||||
// Parameters:
|
||||
// - poolCn: The connection to re-authenticate
|
||||
// - reAuthFn: Function to call for re-authentication, receives error if acquisition fails
|
||||
//
|
||||
// Thread-safe: Called by credentials listeners when credentials change.
|
||||
func (m *Manager) MarkForReAuth(poolCn *pool.Conn, reAuthFn func(error)) {
|
||||
connID := poolCn.GetID()
|
||||
m.poolHookRef.MarkForReAuth(connID, reAuthFn)
|
||||
}
|
||||
|
||||
// RemoveListener removes the credentials listener for a connection.
|
||||
//
|
||||
// This method is called by the pool hook's OnRemove to clean up listeners
|
||||
// when connections are removed from the pool.
|
||||
//
|
||||
// Parameters:
|
||||
// - connID: The connection ID whose listener should be removed
|
||||
//
|
||||
// Thread-safe: Called during connection removal.
|
||||
func (m *Manager) RemoveListener(connID uint64) {
|
||||
m.credentialsListeners.Remove(connID)
|
||||
}
|
||||
+241
@@ -0,0 +1,241 @@
|
||||
package streaming
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9/internal"
|
||||
"github.com/redis/go-redis/v9/internal/pool"
|
||||
)
|
||||
|
||||
// ReAuthPoolHook is a pool hook that manages background re-authentication of connections
|
||||
// when credentials change via a streaming credentials provider.
|
||||
//
|
||||
// The hook uses a semaphore-based worker pool to limit concurrent re-authentication
|
||||
// operations and prevent pool exhaustion. When credentials change, connections are
|
||||
// marked for re-authentication and processed asynchronously in the background.
|
||||
//
|
||||
// The re-authentication process:
|
||||
// 1. OnPut: When a connection is returned to the pool, check if it needs re-auth
|
||||
// 2. If yes, schedule it for background processing (move from shouldReAuth to scheduledReAuth)
|
||||
// 3. A worker goroutine acquires the connection (waits until it's not in use)
|
||||
// 4. Executes the re-auth function while holding the connection
|
||||
// 5. Releases the connection back to the pool
|
||||
//
|
||||
// The hook ensures that:
|
||||
// - Only one re-auth operation runs per connection at a time
|
||||
// - Connections are not used for commands during re-authentication
|
||||
// - Re-auth operations timeout if they can't acquire the connection
|
||||
// - Resources are properly cleaned up on connection removal
|
||||
type ReAuthPoolHook struct {
|
||||
// shouldReAuth maps connection ID to re-auth function
|
||||
// Connections in this map need re-authentication but haven't been scheduled yet
|
||||
shouldReAuth map[uint64]func(error)
|
||||
shouldReAuthLock sync.RWMutex
|
||||
|
||||
// workers is a semaphore limiting concurrent re-auth operations
|
||||
// Initialized with poolSize tokens to prevent pool exhaustion
|
||||
// Uses FastSemaphore for better performance with eventual fairness
|
||||
workers *internal.FastSemaphore
|
||||
|
||||
// reAuthTimeout is the maximum time to wait for acquiring a connection for re-auth
|
||||
reAuthTimeout time.Duration
|
||||
|
||||
// scheduledReAuth maps connection ID to scheduled status
|
||||
// Connections in this map have a background worker attempting re-authentication
|
||||
scheduledReAuth map[uint64]bool
|
||||
scheduledLock sync.RWMutex
|
||||
|
||||
// manager is a back-reference for cleanup operations
|
||||
manager *Manager
|
||||
}
|
||||
|
||||
// NewReAuthPoolHook creates a new re-authentication pool hook.
|
||||
//
|
||||
// Parameters:
|
||||
// - poolSize: Maximum number of concurrent re-auth operations (typically matches pool size)
|
||||
// - reAuthTimeout: Maximum time to wait for acquiring a connection for re-authentication
|
||||
//
|
||||
// The poolSize parameter is used to initialize the worker semaphore, ensuring that
|
||||
// re-auth operations don't exhaust the connection pool.
|
||||
func NewReAuthPoolHook(poolSize int, reAuthTimeout time.Duration) *ReAuthPoolHook {
|
||||
return &ReAuthPoolHook{
|
||||
shouldReAuth: make(map[uint64]func(error)),
|
||||
scheduledReAuth: make(map[uint64]bool),
|
||||
workers: internal.NewFastSemaphore(int32(poolSize)),
|
||||
reAuthTimeout: reAuthTimeout,
|
||||
}
|
||||
}
|
||||
|
||||
// MarkForReAuth marks a connection for re-authentication.
|
||||
//
|
||||
// This method is called when credentials change and a connection needs to be
|
||||
// re-authenticated. The actual re-authentication happens asynchronously when
|
||||
// the connection is returned to the pool (in OnPut).
|
||||
//
|
||||
// Parameters:
|
||||
// - connID: The connection ID to mark for re-authentication
|
||||
// - reAuthFn: Function to call for re-authentication, receives error if acquisition fails
|
||||
//
|
||||
// Thread-safe: Can be called concurrently from multiple goroutines.
|
||||
func (r *ReAuthPoolHook) MarkForReAuth(connID uint64, reAuthFn func(error)) {
|
||||
r.shouldReAuthLock.Lock()
|
||||
defer r.shouldReAuthLock.Unlock()
|
||||
r.shouldReAuth[connID] = reAuthFn
|
||||
}
|
||||
|
||||
// OnGet is called when a connection is retrieved from the pool.
|
||||
//
|
||||
// This hook checks if the connection needs re-authentication or has a scheduled
|
||||
// re-auth operation. If so, it rejects the connection (returns accept=false),
|
||||
// causing the pool to try another connection.
|
||||
//
|
||||
// Returns:
|
||||
// - accept: false if connection needs re-auth, true otherwise
|
||||
// - err: always nil (errors are not used in this hook)
|
||||
//
|
||||
// Thread-safe: Called concurrently by multiple goroutines getting connections.
|
||||
func (r *ReAuthPoolHook) OnGet(_ context.Context, conn *pool.Conn, _ bool) (accept bool, err error) {
|
||||
connID := conn.GetID()
|
||||
r.shouldReAuthLock.RLock()
|
||||
_, shouldReAuth := r.shouldReAuth[connID]
|
||||
r.shouldReAuthLock.RUnlock()
|
||||
// This connection was marked for reauth while in the pool,
|
||||
// reject the connection
|
||||
if shouldReAuth {
|
||||
// simply reject the connection, it will be re-authenticated in OnPut
|
||||
return false, nil
|
||||
}
|
||||
r.scheduledLock.RLock()
|
||||
_, hasScheduled := r.scheduledReAuth[connID]
|
||||
r.scheduledLock.RUnlock()
|
||||
// has scheduled reauth, reject the connection
|
||||
if hasScheduled {
|
||||
// simply reject the connection, it currently has a reauth scheduled
|
||||
// and the worker is waiting for slot to execute the reauth
|
||||
return false, nil
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// OnPut is called when a connection is returned to the pool.
|
||||
//
|
||||
// This hook checks if the connection needs re-authentication. If so, it schedules
|
||||
// a background goroutine to perform the re-auth asynchronously. The goroutine:
|
||||
// 1. Waits for a worker slot (semaphore)
|
||||
// 2. Acquires the connection (waits until not in use)
|
||||
// 3. Executes the re-auth function
|
||||
// 4. Releases the connection and worker slot
|
||||
//
|
||||
// The connection is always pooled (not removed) since re-auth happens in background.
|
||||
//
|
||||
// Returns:
|
||||
// - shouldPool: always true (connection stays in pool during background re-auth)
|
||||
// - shouldRemove: always false
|
||||
// - err: always nil
|
||||
//
|
||||
// Thread-safe: Called concurrently by multiple goroutines returning connections.
|
||||
func (r *ReAuthPoolHook) OnPut(_ context.Context, conn *pool.Conn) (bool, bool, error) {
|
||||
if conn == nil {
|
||||
// noop
|
||||
return true, false, nil
|
||||
}
|
||||
connID := conn.GetID()
|
||||
// Check if reauth is needed and get the function with proper locking
|
||||
r.shouldReAuthLock.RLock()
|
||||
reAuthFn, ok := r.shouldReAuth[connID]
|
||||
r.shouldReAuthLock.RUnlock()
|
||||
|
||||
if ok {
|
||||
// Acquire both locks to atomically move from shouldReAuth to scheduledReAuth
|
||||
// This prevents race conditions where OnGet might miss the transition
|
||||
r.shouldReAuthLock.Lock()
|
||||
r.scheduledLock.Lock()
|
||||
r.scheduledReAuth[connID] = true
|
||||
delete(r.shouldReAuth, connID)
|
||||
r.scheduledLock.Unlock()
|
||||
r.shouldReAuthLock.Unlock()
|
||||
go func() {
|
||||
r.workers.AcquireBlocking()
|
||||
// safety first
|
||||
if conn == nil || (conn != nil && conn.IsClosed()) {
|
||||
r.workers.Release()
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
if rec := recover(); rec != nil {
|
||||
// once again - safety first
|
||||
internal.Logger.Printf(context.Background(), "panic in reauth worker: %v", rec)
|
||||
}
|
||||
r.scheduledLock.Lock()
|
||||
delete(r.scheduledReAuth, connID)
|
||||
r.scheduledLock.Unlock()
|
||||
r.workers.Release()
|
||||
}()
|
||||
|
||||
// Create timeout context for connection acquisition
|
||||
// This prevents indefinite waiting if the connection is stuck
|
||||
ctx, cancel := context.WithTimeout(context.Background(), r.reAuthTimeout)
|
||||
defer cancel()
|
||||
|
||||
// Try to acquire the connection for re-authentication
|
||||
// We need to ensure the connection is IDLE (not IN_USE) before transitioning to UNUSABLE
|
||||
// This prevents re-authentication from interfering with active commands
|
||||
// Use AwaitAndTransition to wait for the connection to become IDLE
|
||||
stateMachine := conn.GetStateMachine()
|
||||
if stateMachine == nil {
|
||||
// No state machine - should not happen, but handle gracefully
|
||||
reAuthFn(pool.ErrConnUnusableTimeout)
|
||||
return
|
||||
}
|
||||
|
||||
// Use predefined slice to avoid allocation
|
||||
_, err := stateMachine.AwaitAndTransition(ctx, pool.ValidFromIdle(), pool.StateUnusable)
|
||||
if err != nil {
|
||||
// Timeout or other error occurred, cannot acquire connection
|
||||
reAuthFn(err)
|
||||
return
|
||||
}
|
||||
|
||||
// safety first
|
||||
if !conn.IsClosed() {
|
||||
// Successfully acquired the connection, perform reauth
|
||||
reAuthFn(nil)
|
||||
}
|
||||
|
||||
// Release the connection: transition from UNUSABLE back to IDLE
|
||||
stateMachine.Transition(pool.StateIdle)
|
||||
}()
|
||||
}
|
||||
|
||||
// the reauth will happen in background, as far as the pool is concerned:
|
||||
// pool the connection, don't remove it, no error
|
||||
return true, false, nil
|
||||
}
|
||||
|
||||
// OnRemove is called when a connection is removed from the pool.
|
||||
//
|
||||
// This hook cleans up all state associated with the connection:
|
||||
// - Removes from shouldReAuth map (pending re-auth)
|
||||
// - Removes from scheduledReAuth map (active re-auth)
|
||||
// - Removes credentials listener from manager
|
||||
//
|
||||
// This prevents memory leaks and ensures that removed connections don't have
|
||||
// lingering re-auth operations or listeners.
|
||||
//
|
||||
// Thread-safe: Called when connections are removed due to errors, timeouts, or pool closure.
|
||||
func (r *ReAuthPoolHook) OnRemove(_ context.Context, conn *pool.Conn, _ error) {
|
||||
connID := conn.GetID()
|
||||
r.shouldReAuthLock.Lock()
|
||||
r.scheduledLock.Lock()
|
||||
delete(r.scheduledReAuth, connID)
|
||||
delete(r.shouldReAuth, connID)
|
||||
r.scheduledLock.Unlock()
|
||||
r.shouldReAuthLock.Unlock()
|
||||
if r.manager != nil {
|
||||
r.manager.RemoveListener(connID)
|
||||
}
|
||||
}
|
||||
|
||||
var _ pool.PoolHook = (*ReAuthPoolHook)(nil)
|
||||
+90
@@ -0,0 +1,90 @@
|
||||
package hashtag
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/redis/go-redis/v9/internal/rand"
|
||||
)
|
||||
|
||||
const slotNumber = 16384
|
||||
|
||||
// CRC16 implementation according to CCITT standards.
|
||||
// Copyright 2001-2010 Georges Menie (www.menie.org)
|
||||
// Copyright 2013 The Go Authors. All rights reserved.
|
||||
// https://redis.io/docs/latest/operate/oss_and_stack/reference/cluster-spec#appendix-a-crc16-reference-implementation-in-ansi-c.
|
||||
var crc16tab = [256]uint16{
|
||||
0x0000, 0x1021, 0x2042, 0x3063, 0x4084, 0x50a5, 0x60c6, 0x70e7,
|
||||
0x8108, 0x9129, 0xa14a, 0xb16b, 0xc18c, 0xd1ad, 0xe1ce, 0xf1ef,
|
||||
0x1231, 0x0210, 0x3273, 0x2252, 0x52b5, 0x4294, 0x72f7, 0x62d6,
|
||||
0x9339, 0x8318, 0xb37b, 0xa35a, 0xd3bd, 0xc39c, 0xf3ff, 0xe3de,
|
||||
0x2462, 0x3443, 0x0420, 0x1401, 0x64e6, 0x74c7, 0x44a4, 0x5485,
|
||||
0xa56a, 0xb54b, 0x8528, 0x9509, 0xe5ee, 0xf5cf, 0xc5ac, 0xd58d,
|
||||
0x3653, 0x2672, 0x1611, 0x0630, 0x76d7, 0x66f6, 0x5695, 0x46b4,
|
||||
0xb75b, 0xa77a, 0x9719, 0x8738, 0xf7df, 0xe7fe, 0xd79d, 0xc7bc,
|
||||
0x48c4, 0x58e5, 0x6886, 0x78a7, 0x0840, 0x1861, 0x2802, 0x3823,
|
||||
0xc9cc, 0xd9ed, 0xe98e, 0xf9af, 0x8948, 0x9969, 0xa90a, 0xb92b,
|
||||
0x5af5, 0x4ad4, 0x7ab7, 0x6a96, 0x1a71, 0x0a50, 0x3a33, 0x2a12,
|
||||
0xdbfd, 0xcbdc, 0xfbbf, 0xeb9e, 0x9b79, 0x8b58, 0xbb3b, 0xab1a,
|
||||
0x6ca6, 0x7c87, 0x4ce4, 0x5cc5, 0x2c22, 0x3c03, 0x0c60, 0x1c41,
|
||||
0xedae, 0xfd8f, 0xcdec, 0xddcd, 0xad2a, 0xbd0b, 0x8d68, 0x9d49,
|
||||
0x7e97, 0x6eb6, 0x5ed5, 0x4ef4, 0x3e13, 0x2e32, 0x1e51, 0x0e70,
|
||||
0xff9f, 0xefbe, 0xdfdd, 0xcffc, 0xbf1b, 0xaf3a, 0x9f59, 0x8f78,
|
||||
0x9188, 0x81a9, 0xb1ca, 0xa1eb, 0xd10c, 0xc12d, 0xf14e, 0xe16f,
|
||||
0x1080, 0x00a1, 0x30c2, 0x20e3, 0x5004, 0x4025, 0x7046, 0x6067,
|
||||
0x83b9, 0x9398, 0xa3fb, 0xb3da, 0xc33d, 0xd31c, 0xe37f, 0xf35e,
|
||||
0x02b1, 0x1290, 0x22f3, 0x32d2, 0x4235, 0x5214, 0x6277, 0x7256,
|
||||
0xb5ea, 0xa5cb, 0x95a8, 0x8589, 0xf56e, 0xe54f, 0xd52c, 0xc50d,
|
||||
0x34e2, 0x24c3, 0x14a0, 0x0481, 0x7466, 0x6447, 0x5424, 0x4405,
|
||||
0xa7db, 0xb7fa, 0x8799, 0x97b8, 0xe75f, 0xf77e, 0xc71d, 0xd73c,
|
||||
0x26d3, 0x36f2, 0x0691, 0x16b0, 0x6657, 0x7676, 0x4615, 0x5634,
|
||||
0xd94c, 0xc96d, 0xf90e, 0xe92f, 0x99c8, 0x89e9, 0xb98a, 0xa9ab,
|
||||
0x5844, 0x4865, 0x7806, 0x6827, 0x18c0, 0x08e1, 0x3882, 0x28a3,
|
||||
0xcb7d, 0xdb5c, 0xeb3f, 0xfb1e, 0x8bf9, 0x9bd8, 0xabbb, 0xbb9a,
|
||||
0x4a75, 0x5a54, 0x6a37, 0x7a16, 0x0af1, 0x1ad0, 0x2ab3, 0x3a92,
|
||||
0xfd2e, 0xed0f, 0xdd6c, 0xcd4d, 0xbdaa, 0xad8b, 0x9de8, 0x8dc9,
|
||||
0x7c26, 0x6c07, 0x5c64, 0x4c45, 0x3ca2, 0x2c83, 0x1ce0, 0x0cc1,
|
||||
0xef1f, 0xff3e, 0xcf5d, 0xdf7c, 0xaf9b, 0xbfba, 0x8fd9, 0x9ff8,
|
||||
0x6e17, 0x7e36, 0x4e55, 0x5e74, 0x2e93, 0x3eb2, 0x0ed1, 0x1ef0,
|
||||
}
|
||||
|
||||
func Key(key string) string {
|
||||
if s := strings.IndexByte(key, '{'); s > -1 {
|
||||
if e := strings.IndexByte(key[s+1:], '}'); e > 0 {
|
||||
return key[s+1 : s+e+1]
|
||||
}
|
||||
}
|
||||
return key
|
||||
}
|
||||
|
||||
func Present(key string) bool {
|
||||
if key == "" {
|
||||
return false
|
||||
}
|
||||
if s := strings.IndexByte(key, '{'); s > -1 {
|
||||
if e := strings.IndexByte(key[s+1:], '}'); e > 0 {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func RandomSlot() int {
|
||||
return rand.Intn(slotNumber)
|
||||
}
|
||||
|
||||
// Slot returns a consistent slot number between 0 and 16383
|
||||
// for any given string key.
|
||||
func Slot(key string) int {
|
||||
if key == "" {
|
||||
return RandomSlot()
|
||||
}
|
||||
key = Key(key)
|
||||
return int(crc16sum(key)) % slotNumber
|
||||
}
|
||||
|
||||
func crc16sum(key string) (crc uint16) {
|
||||
for i := 0; i < len(key); i++ {
|
||||
crc = (crc << 8) ^ crc16tab[(byte(crc>>8)^key[i])&0x00ff]
|
||||
}
|
||||
return
|
||||
}
|
||||
+54
@@ -0,0 +1,54 @@
|
||||
package hashtag
|
||||
|
||||
import "github.com/cespare/xxhash/v2"
|
||||
|
||||
// RendezvousHash implements HRW (Highest Random Weight) hashing.
|
||||
type RendezvousHash struct {
|
||||
nodes []node
|
||||
}
|
||||
|
||||
type node struct {
|
||||
name string
|
||||
hash uint64
|
||||
}
|
||||
|
||||
// NewRendezvousHash builds a hash from shard names.
|
||||
func NewRendezvousHash(shards []string) *RendezvousHash {
|
||||
n := make([]node, len(shards))
|
||||
for i, s := range shards {
|
||||
n[i] = node{
|
||||
name: s,
|
||||
hash: xxhash.Sum64String(s),
|
||||
}
|
||||
}
|
||||
return &RendezvousHash{nodes: n}
|
||||
}
|
||||
|
||||
// Get returns the shard name for the given key.
|
||||
func (r *RendezvousHash) Get(key string) string {
|
||||
if len(r.nodes) == 0 {
|
||||
return ""
|
||||
}
|
||||
|
||||
kh := xxhash.Sum64String(key)
|
||||
|
||||
bestIdx := 0
|
||||
bestScore := mix64(kh ^ r.nodes[0].hash)
|
||||
|
||||
for i := 1; i < len(r.nodes); i++ {
|
||||
if score := mix64(kh ^ r.nodes[i].hash); score > bestScore {
|
||||
bestScore = score
|
||||
bestIdx = i
|
||||
}
|
||||
}
|
||||
|
||||
return r.nodes[bestIdx].name
|
||||
}
|
||||
|
||||
// mix64 is a xorshift-based mixing function.
|
||||
func mix64(x uint64) uint64 {
|
||||
x ^= x >> 12
|
||||
x ^= x << 25
|
||||
x ^= x >> 27
|
||||
return x * 2685821657736338717
|
||||
}
|
||||
+207
@@ -0,0 +1,207 @@
|
||||
package hscan
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"reflect"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
// decoderFunc represents decoding functions for default built-in types.
|
||||
type decoderFunc func(reflect.Value, string) error
|
||||
|
||||
// Scanner is the interface implemented by themselves,
|
||||
// which will override the decoding behavior of decoderFunc.
|
||||
type Scanner interface {
|
||||
ScanRedis(s string) error
|
||||
}
|
||||
|
||||
var (
|
||||
// List of built-in decoders indexed by their numeric constant values (eg: reflect.Bool = 1).
|
||||
decoders = []decoderFunc{
|
||||
reflect.Bool: decodeBool,
|
||||
reflect.Int: decodeInt,
|
||||
reflect.Int8: decodeInt8,
|
||||
reflect.Int16: decodeInt16,
|
||||
reflect.Int32: decodeInt32,
|
||||
reflect.Int64: decodeInt64,
|
||||
reflect.Uint: decodeUint,
|
||||
reflect.Uint8: decodeUint8,
|
||||
reflect.Uint16: decodeUint16,
|
||||
reflect.Uint32: decodeUint32,
|
||||
reflect.Uint64: decodeUint64,
|
||||
reflect.Float32: decodeFloat32,
|
||||
reflect.Float64: decodeFloat64,
|
||||
reflect.Complex64: decodeUnsupported,
|
||||
reflect.Complex128: decodeUnsupported,
|
||||
reflect.Array: decodeUnsupported,
|
||||
reflect.Chan: decodeUnsupported,
|
||||
reflect.Func: decodeUnsupported,
|
||||
reflect.Interface: decodeUnsupported,
|
||||
reflect.Map: decodeUnsupported,
|
||||
reflect.Ptr: decodeUnsupported,
|
||||
reflect.Slice: decodeSlice,
|
||||
reflect.String: decodeString,
|
||||
reflect.Struct: decodeUnsupported,
|
||||
reflect.UnsafePointer: decodeUnsupported,
|
||||
}
|
||||
|
||||
// Global map of struct field specs that is populated once for every new
|
||||
// struct type that is scanned. This caches the field types and the corresponding
|
||||
// decoder functions to avoid iterating through struct fields on subsequent scans.
|
||||
globalStructMap = newStructMap()
|
||||
)
|
||||
|
||||
func Struct(dst interface{}) (StructValue, error) {
|
||||
v := reflect.ValueOf(dst)
|
||||
|
||||
// The destination to scan into should be a struct pointer.
|
||||
if v.Kind() != reflect.Ptr || v.IsNil() {
|
||||
return StructValue{}, fmt.Errorf("redis.Scan(non-pointer %T)", dst)
|
||||
}
|
||||
|
||||
v = v.Elem()
|
||||
if v.Kind() != reflect.Struct {
|
||||
return StructValue{}, fmt.Errorf("redis.Scan(non-struct %T)", dst)
|
||||
}
|
||||
|
||||
return StructValue{
|
||||
spec: globalStructMap.get(v.Type()),
|
||||
value: v,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Scan scans the results from a key-value Redis map result set to a destination struct.
|
||||
// The Redis keys are matched to the struct's field with the `redis` tag.
|
||||
func Scan(dst interface{}, keys []interface{}, vals []interface{}) error {
|
||||
if len(keys) != len(vals) {
|
||||
return errors.New("args should have the same number of keys and vals")
|
||||
}
|
||||
|
||||
strct, err := Struct(dst)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Iterate through the (key, value) sequence.
|
||||
for i := 0; i < len(vals); i++ {
|
||||
key, ok := keys[i].(string)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
val, ok := vals[i].(string)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
if err := strct.Scan(key, val); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func decodeBool(f reflect.Value, s string) error {
|
||||
b, err := strconv.ParseBool(s)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
f.SetBool(b)
|
||||
return nil
|
||||
}
|
||||
|
||||
func decodeInt8(f reflect.Value, s string) error {
|
||||
return decodeNumber(f, s, 8)
|
||||
}
|
||||
|
||||
func decodeInt16(f reflect.Value, s string) error {
|
||||
return decodeNumber(f, s, 16)
|
||||
}
|
||||
|
||||
func decodeInt32(f reflect.Value, s string) error {
|
||||
return decodeNumber(f, s, 32)
|
||||
}
|
||||
|
||||
func decodeInt64(f reflect.Value, s string) error {
|
||||
return decodeNumber(f, s, 64)
|
||||
}
|
||||
|
||||
func decodeInt(f reflect.Value, s string) error {
|
||||
return decodeNumber(f, s, 0)
|
||||
}
|
||||
|
||||
func decodeNumber(f reflect.Value, s string, bitSize int) error {
|
||||
v, err := strconv.ParseInt(s, 10, bitSize)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
f.SetInt(v)
|
||||
return nil
|
||||
}
|
||||
|
||||
func decodeUint8(f reflect.Value, s string) error {
|
||||
return decodeUnsignedNumber(f, s, 8)
|
||||
}
|
||||
|
||||
func decodeUint16(f reflect.Value, s string) error {
|
||||
return decodeUnsignedNumber(f, s, 16)
|
||||
}
|
||||
|
||||
func decodeUint32(f reflect.Value, s string) error {
|
||||
return decodeUnsignedNumber(f, s, 32)
|
||||
}
|
||||
|
||||
func decodeUint64(f reflect.Value, s string) error {
|
||||
return decodeUnsignedNumber(f, s, 64)
|
||||
}
|
||||
|
||||
func decodeUint(f reflect.Value, s string) error {
|
||||
return decodeUnsignedNumber(f, s, 0)
|
||||
}
|
||||
|
||||
func decodeUnsignedNumber(f reflect.Value, s string, bitSize int) error {
|
||||
v, err := strconv.ParseUint(s, 10, bitSize)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
f.SetUint(v)
|
||||
return nil
|
||||
}
|
||||
|
||||
func decodeFloat32(f reflect.Value, s string) error {
|
||||
v, err := strconv.ParseFloat(s, 32)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
f.SetFloat(v)
|
||||
return nil
|
||||
}
|
||||
|
||||
// although the default is float64, but we better define it.
|
||||
func decodeFloat64(f reflect.Value, s string) error {
|
||||
v, err := strconv.ParseFloat(s, 64)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
f.SetFloat(v)
|
||||
return nil
|
||||
}
|
||||
|
||||
func decodeString(f reflect.Value, s string) error {
|
||||
f.SetString(s)
|
||||
return nil
|
||||
}
|
||||
|
||||
func decodeSlice(f reflect.Value, s string) error {
|
||||
// []byte slice ([]uint8).
|
||||
if f.Type().Elem().Kind() == reflect.Uint8 {
|
||||
f.SetBytes([]byte(s))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func decodeUnsupported(v reflect.Value, s string) error {
|
||||
return fmt.Errorf("redis.Scan(unsupported %s)", v.Type())
|
||||
}
|
||||
+127
@@ -0,0 +1,127 @@
|
||||
package hscan
|
||||
|
||||
import (
|
||||
"encoding"
|
||||
"fmt"
|
||||
"reflect"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/redis/go-redis/v9/internal/util"
|
||||
)
|
||||
|
||||
// structMap contains the map of struct fields for target structs
|
||||
// indexed by the struct type.
|
||||
type structMap struct {
|
||||
m sync.Map
|
||||
}
|
||||
|
||||
func newStructMap() *structMap {
|
||||
return new(structMap)
|
||||
}
|
||||
|
||||
func (s *structMap) get(t reflect.Type) *structSpec {
|
||||
if v, ok := s.m.Load(t); ok {
|
||||
return v.(*structSpec)
|
||||
}
|
||||
|
||||
spec := newStructSpec(t, "redis")
|
||||
s.m.Store(t, spec)
|
||||
return spec
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------------
|
||||
|
||||
// structSpec contains the list of all fields in a target struct.
|
||||
type structSpec struct {
|
||||
m map[string]*structField
|
||||
}
|
||||
|
||||
func (s *structSpec) set(tag string, sf *structField) {
|
||||
s.m[tag] = sf
|
||||
}
|
||||
|
||||
func newStructSpec(t reflect.Type, fieldTag string) *structSpec {
|
||||
numField := t.NumField()
|
||||
out := &structSpec{
|
||||
m: make(map[string]*structField, numField),
|
||||
}
|
||||
|
||||
for i := 0; i < numField; i++ {
|
||||
f := t.Field(i)
|
||||
|
||||
tag := f.Tag.Get(fieldTag)
|
||||
if tag == "" || tag == "-" {
|
||||
continue
|
||||
}
|
||||
|
||||
tag = strings.Split(tag, ",")[0]
|
||||
if tag == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
// Use the built-in decoder.
|
||||
kind := f.Type.Kind()
|
||||
if kind == reflect.Pointer {
|
||||
kind = f.Type.Elem().Kind()
|
||||
}
|
||||
out.set(tag, &structField{index: i, fn: decoders[kind]})
|
||||
}
|
||||
|
||||
return out
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------------
|
||||
|
||||
// structField represents a single field in a target struct.
|
||||
type structField struct {
|
||||
index int
|
||||
fn decoderFunc
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------------
|
||||
|
||||
type StructValue struct {
|
||||
spec *structSpec
|
||||
value reflect.Value
|
||||
}
|
||||
|
||||
func (s StructValue) Scan(key string, value string) error {
|
||||
field, ok := s.spec.m[key]
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
v := s.value.Field(field.index)
|
||||
isPtr := v.Kind() == reflect.Ptr
|
||||
|
||||
if isPtr && v.IsNil() {
|
||||
v.Set(reflect.New(v.Type().Elem()))
|
||||
}
|
||||
if !isPtr && v.Type().Name() != "" && v.CanAddr() {
|
||||
v = v.Addr()
|
||||
isPtr = true
|
||||
}
|
||||
|
||||
if isPtr && v.Type().NumMethod() > 0 && v.CanInterface() {
|
||||
switch scan := v.Interface().(type) {
|
||||
case Scanner:
|
||||
return scan.ScanRedis(value)
|
||||
case encoding.TextUnmarshaler:
|
||||
return scan.UnmarshalText(util.StringToBytes(value))
|
||||
case encoding.BinaryUnmarshaler:
|
||||
return scan.UnmarshalBinary(util.StringToBytes(value))
|
||||
}
|
||||
}
|
||||
|
||||
if isPtr {
|
||||
v = v.Elem()
|
||||
}
|
||||
|
||||
if err := field.fn(v, value); err != nil {
|
||||
t := s.value.Type()
|
||||
return fmt.Errorf("cannot scan redis.result %s into struct field %s.%s of type %s, error-%s",
|
||||
value, t.Name(), t.Field(field.index).Name, t.Field(field.index).Type, err.Error())
|
||||
}
|
||||
return nil
|
||||
}
|
||||
+59
@@ -0,0 +1,59 @@
|
||||
// Package interfaces provides shared interfaces used by both the main redis package
|
||||
// and the maintnotifications upgrade package to avoid circular dependencies.
|
||||
package interfaces
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"time"
|
||||
)
|
||||
|
||||
// NotificationProcessor is (most probably) a push.NotificationProcessor
|
||||
// forward declaration to avoid circular imports
|
||||
type NotificationProcessor interface {
|
||||
RegisterHandler(pushNotificationName string, handler interface{}, protected bool) error
|
||||
UnregisterHandler(pushNotificationName string) error
|
||||
GetHandler(pushNotificationName string) interface{}
|
||||
}
|
||||
|
||||
// ClientInterface defines the interface that clients must implement for maintnotifications upgrades.
|
||||
type ClientInterface interface {
|
||||
// GetOptions returns the client options.
|
||||
GetOptions() OptionsInterface
|
||||
|
||||
// GetPushProcessor returns the client's push notification processor.
|
||||
GetPushProcessor() NotificationProcessor
|
||||
}
|
||||
|
||||
// OptionsInterface defines the interface for client options.
|
||||
// Uses an adapter pattern to avoid circular dependencies.
|
||||
type OptionsInterface interface {
|
||||
// GetReadTimeout returns the read timeout.
|
||||
GetReadTimeout() time.Duration
|
||||
|
||||
// GetWriteTimeout returns the write timeout.
|
||||
GetWriteTimeout() time.Duration
|
||||
|
||||
// GetNetwork returns the network type.
|
||||
GetNetwork() string
|
||||
|
||||
// GetAddr returns the connection address.
|
||||
GetAddr() string
|
||||
|
||||
// GetNodeAddress returns the address of the Redis node as reported by the server.
|
||||
// For cluster clients, this is the endpoint from CLUSTER SLOTS before any transformation.
|
||||
// For standalone clients, this defaults to Addr.
|
||||
GetNodeAddress() string
|
||||
|
||||
// IsTLSEnabled returns true if TLS is enabled.
|
||||
IsTLSEnabled() bool
|
||||
|
||||
// GetProtocol returns the protocol version.
|
||||
GetProtocol() int
|
||||
|
||||
// GetPoolSize returns the connection pool size.
|
||||
GetPoolSize() int
|
||||
|
||||
// NewDialer returns a new dialer function for the connection.
|
||||
NewDialer() func(context.Context) (net.Conn, error)
|
||||
}
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
package internal
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9/internal/rand"
|
||||
)
|
||||
|
||||
func RetryBackoff(retry int, minBackoff, maxBackoff time.Duration) time.Duration {
|
||||
if retry < 0 {
|
||||
panic("not reached")
|
||||
}
|
||||
if minBackoff == 0 {
|
||||
return 0
|
||||
}
|
||||
|
||||
d := minBackoff << uint(retry)
|
||||
if d < minBackoff {
|
||||
return maxBackoff
|
||||
}
|
||||
|
||||
d = minBackoff + time.Duration(rand.Int63n(int64(d)))
|
||||
|
||||
if d > maxBackoff || d < minBackoff {
|
||||
d = maxBackoff
|
||||
}
|
||||
|
||||
return d
|
||||
}
|
||||
+79
@@ -0,0 +1,79 @@
|
||||
package internal
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
)
|
||||
|
||||
// TODO (ned): Revisit logging
|
||||
// Add more standardized approach with log levels and configurability
|
||||
|
||||
type Logging interface {
|
||||
Printf(ctx context.Context, format string, v ...interface{})
|
||||
}
|
||||
|
||||
type DefaultLogger struct {
|
||||
log *log.Logger
|
||||
}
|
||||
|
||||
func (l *DefaultLogger) Printf(ctx context.Context, format string, v ...interface{}) {
|
||||
_ = l.log.Output(2, fmt.Sprintf(format, v...))
|
||||
}
|
||||
|
||||
func NewDefaultLogger() Logging {
|
||||
return &DefaultLogger{
|
||||
log: log.New(os.Stderr, "redis: ", log.LstdFlags|log.Lshortfile),
|
||||
}
|
||||
}
|
||||
|
||||
// Logger calls Output to print to the stderr.
|
||||
// Arguments are handled in the manner of fmt.Print.
|
||||
var Logger Logging = NewDefaultLogger()
|
||||
|
||||
var LogLevel LogLevelT = LogLevelError
|
||||
|
||||
// LogLevelT represents the logging level
|
||||
type LogLevelT int
|
||||
|
||||
// Log level constants for the entire go-redis library
|
||||
const (
|
||||
LogLevelError LogLevelT = iota // 0 - errors only
|
||||
LogLevelWarn // 1 - warnings and errors
|
||||
LogLevelInfo // 2 - info, warnings, and errors
|
||||
LogLevelDebug // 3 - debug, info, warnings, and errors
|
||||
)
|
||||
|
||||
// String returns the string representation of the log level
|
||||
func (l LogLevelT) String() string {
|
||||
switch l {
|
||||
case LogLevelError:
|
||||
return "ERROR"
|
||||
case LogLevelWarn:
|
||||
return "WARN"
|
||||
case LogLevelInfo:
|
||||
return "INFO"
|
||||
case LogLevelDebug:
|
||||
return "DEBUG"
|
||||
default:
|
||||
return "UNKNOWN"
|
||||
}
|
||||
}
|
||||
|
||||
// IsValid returns true if the log level is valid
|
||||
func (l LogLevelT) IsValid() bool {
|
||||
return l >= LogLevelError && l <= LogLevelDebug
|
||||
}
|
||||
|
||||
func (l LogLevelT) WarnOrAbove() bool {
|
||||
return l >= LogLevelWarn
|
||||
}
|
||||
|
||||
func (l LogLevelT) InfoOrAbove() bool {
|
||||
return l >= LogLevelInfo
|
||||
}
|
||||
|
||||
func (l LogLevelT) DebugOrAbove() bool {
|
||||
return l >= LogLevelDebug
|
||||
}
|
||||
Generated
Vendored
+663
@@ -0,0 +1,663 @@
|
||||
package logs
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"regexp"
|
||||
|
||||
"github.com/redis/go-redis/v9/internal"
|
||||
)
|
||||
|
||||
// appendJSONIfDebug appends JSON data to a message only if the global log level is Debug
|
||||
func appendJSONIfDebug(message string, data map[string]interface{}) string {
|
||||
if internal.LogLevel.DebugOrAbove() {
|
||||
jsonData, _ := json.Marshal(data)
|
||||
return fmt.Sprintf("%s %s", message, string(jsonData))
|
||||
}
|
||||
return message
|
||||
}
|
||||
|
||||
const (
|
||||
// ========================================
|
||||
// CIRCUIT_BREAKER.GO - Circuit breaker management
|
||||
// ========================================
|
||||
CircuitBreakerTransitioningToHalfOpenMessage = "circuit breaker transitioning to half-open"
|
||||
CircuitBreakerOpenedMessage = "circuit breaker opened"
|
||||
CircuitBreakerReopenedMessage = "circuit breaker reopened"
|
||||
CircuitBreakerClosedMessage = "circuit breaker closed"
|
||||
CircuitBreakerCleanupMessage = "circuit breaker cleanup"
|
||||
CircuitBreakerOpenMessage = "circuit breaker is open, failing fast"
|
||||
|
||||
// ========================================
|
||||
// CONFIG.GO - Configuration and debug
|
||||
// ========================================
|
||||
DebugLoggingEnabledMessage = "debug logging enabled"
|
||||
ConfigDebugMessage = "config debug"
|
||||
|
||||
// ========================================
|
||||
// ERRORS.GO - Error message constants
|
||||
// ========================================
|
||||
InvalidRelaxedTimeoutErrorMessage = "relaxed timeout must be greater than 0"
|
||||
InvalidHandoffTimeoutErrorMessage = "handoff timeout must be greater than 0"
|
||||
InvalidHandoffWorkersErrorMessage = "MaxWorkers must be greater than or equal to 0"
|
||||
InvalidHandoffQueueSizeErrorMessage = "handoff queue size must be greater than 0"
|
||||
InvalidPostHandoffRelaxedDurationErrorMessage = "post-handoff relaxed duration must be greater than or equal to 0"
|
||||
InvalidEndpointTypeErrorMessage = "invalid endpoint type"
|
||||
InvalidMaintNotificationsErrorMessage = "invalid maintenance notifications setting (must be 'disabled', 'enabled', or 'auto')"
|
||||
InvalidHandoffRetriesErrorMessage = "MaxHandoffRetries must be between 1 and 10"
|
||||
InvalidClientErrorMessage = "invalid client type"
|
||||
InvalidNotificationErrorMessage = "invalid notification format"
|
||||
MaxHandoffRetriesReachedErrorMessage = "max handoff retries reached"
|
||||
HandoffQueueFullErrorMessage = "handoff queue is full, cannot queue new handoff requests - consider increasing HandoffQueueSize or MaxWorkers in configuration"
|
||||
InvalidCircuitBreakerFailureThresholdErrorMessage = "circuit breaker failure threshold must be >= 1"
|
||||
InvalidCircuitBreakerResetTimeoutErrorMessage = "circuit breaker reset timeout must be >= 0"
|
||||
InvalidCircuitBreakerMaxRequestsErrorMessage = "circuit breaker max requests must be >= 1"
|
||||
ConnectionMarkedForHandoffErrorMessage = "connection marked for handoff"
|
||||
ConnectionInvalidHandoffStateErrorMessage = "connection is in invalid state for handoff"
|
||||
ShutdownErrorMessage = "shutdown"
|
||||
CircuitBreakerOpenErrorMessage = "circuit breaker is open, failing fast"
|
||||
|
||||
// ========================================
|
||||
// EXAMPLE_HOOKS.GO - Example metrics hooks
|
||||
// ========================================
|
||||
MetricsHookProcessingNotificationMessage = "metrics hook processing"
|
||||
MetricsHookRecordedErrorMessage = "metrics hook recorded error"
|
||||
|
||||
// ========================================
|
||||
// HANDOFF_WORKER.GO - Connection handoff processing
|
||||
// ========================================
|
||||
HandoffStartedMessage = "handoff started"
|
||||
HandoffFailedMessage = "handoff failed"
|
||||
ConnectionNotMarkedForHandoffMessage = "is not marked for handoff and has no retries"
|
||||
ConnectionNotMarkedForHandoffErrorMessage = "is not marked for handoff"
|
||||
HandoffRetryAttemptMessage = "Performing handoff"
|
||||
CannotQueueHandoffForRetryMessage = "can't queue handoff for retry"
|
||||
HandoffQueueFullMessage = "handoff queue is full"
|
||||
FailedToDialNewEndpointMessage = "failed to dial new endpoint"
|
||||
ApplyingRelaxedTimeoutDueToPostHandoffMessage = "applying relaxed timeout due to post-handoff"
|
||||
HandoffSuccessMessage = "handoff succeeded"
|
||||
RemovingConnectionFromPoolMessage = "removing connection from pool"
|
||||
NoPoolProvidedMessageCannotRemoveMessage = "no pool provided, cannot remove connection, closing it"
|
||||
WorkerExitingDueToShutdownMessage = "worker exiting due to shutdown"
|
||||
WorkerExitingDueToShutdownWhileProcessingMessage = "worker exiting due to shutdown while processing request"
|
||||
WorkerPanicRecoveredMessage = "worker panic recovered"
|
||||
WorkerExitingDueToInactivityTimeoutMessage = "worker exiting due to inactivity timeout"
|
||||
ReachedMaxHandoffRetriesMessage = "reached max handoff retries"
|
||||
|
||||
// ========================================
|
||||
// MANAGER.GO - Moving operation tracking and handler registration
|
||||
// ========================================
|
||||
DuplicateMovingOperationMessage = "duplicate MOVING operation ignored"
|
||||
TrackingMovingOperationMessage = "tracking MOVING operation"
|
||||
UntrackingMovingOperationMessage = "untracking MOVING operation"
|
||||
OperationNotTrackedMessage = "operation not tracked"
|
||||
FailedToRegisterHandlerMessage = "failed to register handler"
|
||||
|
||||
// ========================================
|
||||
// HOOKS.GO - Notification processing hooks
|
||||
// ========================================
|
||||
ProcessingNotificationMessage = "processing notification started"
|
||||
ProcessingNotificationFailedMessage = "proccessing notification failed"
|
||||
ProcessingNotificationSucceededMessage = "processing notification succeeded"
|
||||
|
||||
// ========================================
|
||||
// POOL_HOOK.GO - Pool connection management
|
||||
// ========================================
|
||||
FailedToQueueHandoffMessage = "failed to queue handoff"
|
||||
MarkedForHandoffMessage = "connection marked for handoff"
|
||||
|
||||
// ========================================
|
||||
// PUSH_NOTIFICATION_HANDLER.GO - Push notification validation and processing
|
||||
// ========================================
|
||||
InvalidNotificationFormatMessage = "invalid notification format"
|
||||
InvalidNotificationTypeFormatMessage = "invalid notification type format"
|
||||
InvalidSeqIDInMovingNotificationMessage = "invalid seqID in MOVING notification"
|
||||
InvalidTimeSInMovingNotificationMessage = "invalid timeS in MOVING notification"
|
||||
InvalidNewEndpointInMovingNotificationMessage = "invalid newEndpoint in MOVING notification"
|
||||
NoConnectionInHandlerContextMessage = "no connection in handler context"
|
||||
InvalidConnectionTypeInHandlerContextMessage = "invalid connection type in handler context"
|
||||
SchedulingHandoffToCurrentEndpointMessage = "scheduling handoff to current endpoint"
|
||||
RelaxedTimeoutDueToNotificationMessage = "applying relaxed timeout due to notification"
|
||||
UnrelaxedTimeoutMessage = "clearing relaxed timeout"
|
||||
ManagerNotInitializedMessage = "manager not initialized"
|
||||
FailedToMarkForHandoffMessage = "failed to mark connection for handoff"
|
||||
InvalidSeqIDInSMigratingNotificationMessage = "invalid SeqID in SMIGRATING notification"
|
||||
InvalidSeqIDInSMigratedNotificationMessage = "invalid SeqID in SMIGRATED notification"
|
||||
TriggeringClusterStateReloadMessage = "triggering cluster state reload"
|
||||
|
||||
// ========================================
|
||||
// used in pool/conn
|
||||
// ========================================
|
||||
UnrelaxedTimeoutAfterDeadlineMessage = "clearing relaxed timeout after deadline"
|
||||
)
|
||||
|
||||
func HandoffStarted(connID uint64, newEndpoint string) string {
|
||||
message := fmt.Sprintf("conn[%d] %s to %s", connID, HandoffStartedMessage, newEndpoint)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
"endpoint": newEndpoint,
|
||||
})
|
||||
}
|
||||
|
||||
func HandoffFailed(connID uint64, newEndpoint string, attempt int, maxAttempts int, err error) string {
|
||||
message := fmt.Sprintf("conn[%d] %s to %s (attempt %d/%d): %v", connID, HandoffFailedMessage, newEndpoint, attempt, maxAttempts, err)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
"endpoint": newEndpoint,
|
||||
"attempt": attempt,
|
||||
"maxAttempts": maxAttempts,
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
|
||||
func HandoffSucceeded(connID uint64, newEndpoint string) string {
|
||||
message := fmt.Sprintf("conn[%d] %s to %s", connID, HandoffSuccessMessage, newEndpoint)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
"endpoint": newEndpoint,
|
||||
})
|
||||
}
|
||||
|
||||
// Timeout-related log functions
|
||||
func RelaxedTimeoutDueToNotification(connID uint64, notificationType string, timeout interface{}) string {
|
||||
message := fmt.Sprintf("conn[%d] %s %s (%v)", connID, RelaxedTimeoutDueToNotificationMessage, notificationType, timeout)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
"notificationType": notificationType,
|
||||
"timeout": fmt.Sprintf("%v", timeout),
|
||||
})
|
||||
}
|
||||
|
||||
func UnrelaxedTimeout(connID uint64) string {
|
||||
message := fmt.Sprintf("conn[%d] %s", connID, UnrelaxedTimeoutMessage)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
})
|
||||
}
|
||||
|
||||
func UnrelaxedTimeoutAfterDeadline(connID uint64) string {
|
||||
message := fmt.Sprintf("conn[%d] %s", connID, UnrelaxedTimeoutAfterDeadlineMessage)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
})
|
||||
}
|
||||
|
||||
// Handoff queue and marking functions
|
||||
func HandoffQueueFull(queueLen, queueCap int) string {
|
||||
message := fmt.Sprintf("%s (%d/%d), cannot queue new handoff requests - consider increasing HandoffQueueSize or MaxWorkers in configuration", HandoffQueueFullMessage, queueLen, queueCap)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"queueLen": queueLen,
|
||||
"queueCap": queueCap,
|
||||
})
|
||||
}
|
||||
|
||||
func FailedToQueueHandoff(connID uint64, err error) string {
|
||||
message := fmt.Sprintf("conn[%d] %s: %v", connID, FailedToQueueHandoffMessage, err)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
|
||||
func FailedToMarkForHandoff(connID uint64, err error) string {
|
||||
message := fmt.Sprintf("conn[%d] %s: %v", connID, FailedToMarkForHandoffMessage, err)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
|
||||
func FailedToDialNewEndpoint(connID uint64, endpoint string, err error) string {
|
||||
message := fmt.Sprintf("conn[%d] %s %s: %v", connID, FailedToDialNewEndpointMessage, endpoint, err)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
"endpoint": endpoint,
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
|
||||
func ReachedMaxHandoffRetries(connID uint64, endpoint string, maxRetries int) string {
|
||||
message := fmt.Sprintf("conn[%d] %s to %s (max retries: %d)", connID, ReachedMaxHandoffRetriesMessage, endpoint, maxRetries)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
"endpoint": endpoint,
|
||||
"maxRetries": maxRetries,
|
||||
})
|
||||
}
|
||||
|
||||
// Notification processing functions
|
||||
func ProcessingNotification(connID uint64, seqID int64, notificationType string, notification interface{}) string {
|
||||
message := fmt.Sprintf("conn[%d] seqID[%d] %s %s: %v", connID, seqID, ProcessingNotificationMessage, notificationType, notification)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
"seqID": seqID,
|
||||
"notificationType": notificationType,
|
||||
"notification": fmt.Sprintf("%v", notification),
|
||||
})
|
||||
}
|
||||
|
||||
func ProcessingNotificationFailed(connID uint64, notificationType string, err error, notification interface{}) string {
|
||||
message := fmt.Sprintf("conn[%d] %s %s: %v - %v", connID, ProcessingNotificationFailedMessage, notificationType, err, notification)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
"notificationType": notificationType,
|
||||
"error": err.Error(),
|
||||
"notification": fmt.Sprintf("%v", notification),
|
||||
})
|
||||
}
|
||||
|
||||
func ProcessingNotificationSucceeded(connID uint64, notificationType string) string {
|
||||
message := fmt.Sprintf("conn[%d] %s %s", connID, ProcessingNotificationSucceededMessage, notificationType)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
"notificationType": notificationType,
|
||||
})
|
||||
}
|
||||
|
||||
// Moving operation tracking functions
|
||||
func DuplicateMovingOperation(connID uint64, endpoint string, seqID int64) string {
|
||||
message := fmt.Sprintf("conn[%d] %s for %s seqID[%d]", connID, DuplicateMovingOperationMessage, endpoint, seqID)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
"endpoint": endpoint,
|
||||
"seqID": seqID,
|
||||
})
|
||||
}
|
||||
|
||||
func TrackingMovingOperation(connID uint64, endpoint string, seqID int64) string {
|
||||
message := fmt.Sprintf("conn[%d] %s for %s seqID[%d]", connID, TrackingMovingOperationMessage, endpoint, seqID)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
"endpoint": endpoint,
|
||||
"seqID": seqID,
|
||||
})
|
||||
}
|
||||
|
||||
func UntrackingMovingOperation(connID uint64, seqID int64) string {
|
||||
message := fmt.Sprintf("conn[%d] %s seqID[%d]", connID, UntrackingMovingOperationMessage, seqID)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
"seqID": seqID,
|
||||
})
|
||||
}
|
||||
|
||||
func OperationNotTracked(connID uint64, seqID int64) string {
|
||||
message := fmt.Sprintf("conn[%d] %s seqID[%d]", connID, OperationNotTrackedMessage, seqID)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
"seqID": seqID,
|
||||
})
|
||||
}
|
||||
|
||||
// Connection pool functions
|
||||
func RemovingConnectionFromPool(connID uint64, reason error) string {
|
||||
metadata := map[string]interface{}{
|
||||
"connID": connID,
|
||||
"reason": "unknown", // this will be overwritten if reason is not nil
|
||||
}
|
||||
if reason != nil {
|
||||
metadata["reason"] = reason.Error()
|
||||
}
|
||||
|
||||
message := fmt.Sprintf("conn[%d] %s due to: %v", connID, RemovingConnectionFromPoolMessage, reason)
|
||||
return appendJSONIfDebug(message, metadata)
|
||||
}
|
||||
|
||||
func NoPoolProvidedCannotRemove(connID uint64, reason error) string {
|
||||
metadata := map[string]interface{}{
|
||||
"connID": connID,
|
||||
"reason": "unknown", // this will be overwritten if reason is not nil
|
||||
}
|
||||
if reason != nil {
|
||||
metadata["reason"] = reason.Error()
|
||||
}
|
||||
|
||||
message := fmt.Sprintf("conn[%d] %s due to: %v", connID, NoPoolProvidedMessageCannotRemoveMessage, reason)
|
||||
return appendJSONIfDebug(message, metadata)
|
||||
}
|
||||
|
||||
// Circuit breaker functions
|
||||
func CircuitBreakerOpen(connID uint64, endpoint string) string {
|
||||
message := fmt.Sprintf("conn[%d] %s for %s", connID, CircuitBreakerOpenMessage, endpoint)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
"endpoint": endpoint,
|
||||
})
|
||||
}
|
||||
|
||||
// Additional handoff functions for specific cases
|
||||
func ConnectionNotMarkedForHandoff(connID uint64) string {
|
||||
message := fmt.Sprintf("conn[%d] %s", connID, ConnectionNotMarkedForHandoffMessage)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
})
|
||||
}
|
||||
|
||||
func ConnectionNotMarkedForHandoffError(connID uint64) string {
|
||||
return fmt.Sprintf("conn[%d] %s", connID, ConnectionNotMarkedForHandoffErrorMessage)
|
||||
}
|
||||
|
||||
func HandoffRetryAttempt(connID uint64, retries int, newEndpoint string, oldEndpoint string) string {
|
||||
message := fmt.Sprintf("conn[%d] Retry %d: %s to %s(was %s)", connID, retries, HandoffRetryAttemptMessage, newEndpoint, oldEndpoint)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
"retries": retries,
|
||||
"newEndpoint": newEndpoint,
|
||||
"oldEndpoint": oldEndpoint,
|
||||
})
|
||||
}
|
||||
|
||||
func CannotQueueHandoffForRetry(err error) string {
|
||||
message := fmt.Sprintf("%s: %v", CannotQueueHandoffForRetryMessage, err)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
|
||||
// Validation and error functions
|
||||
func InvalidNotificationFormat(notification interface{}) string {
|
||||
message := fmt.Sprintf("%s: %v", InvalidNotificationFormatMessage, notification)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"notification": fmt.Sprintf("%v", notification),
|
||||
})
|
||||
}
|
||||
|
||||
func InvalidNotificationTypeFormat(notificationType interface{}) string {
|
||||
message := fmt.Sprintf("%s: %v", InvalidNotificationTypeFormatMessage, notificationType)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"notificationType": fmt.Sprintf("%v", notificationType),
|
||||
})
|
||||
}
|
||||
|
||||
// InvalidNotification creates a log message for invalid notifications of any type
|
||||
func InvalidNotification(notificationType string, notification interface{}) string {
|
||||
message := fmt.Sprintf("invalid %s notification: %v", notificationType, notification)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"notificationType": notificationType,
|
||||
"notification": fmt.Sprintf("%v", notification),
|
||||
})
|
||||
}
|
||||
|
||||
func InvalidSeqIDInMovingNotification(seqID interface{}) string {
|
||||
message := fmt.Sprintf("%s: %v", InvalidSeqIDInMovingNotificationMessage, seqID)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"seqID": fmt.Sprintf("%v", seqID),
|
||||
})
|
||||
}
|
||||
|
||||
func InvalidTimeSInMovingNotification(timeS interface{}) string {
|
||||
message := fmt.Sprintf("%s: %v", InvalidTimeSInMovingNotificationMessage, timeS)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"timeS": fmt.Sprintf("%v", timeS),
|
||||
})
|
||||
}
|
||||
|
||||
func InvalidNewEndpointInMovingNotification(newEndpoint interface{}) string {
|
||||
message := fmt.Sprintf("%s: %v", InvalidNewEndpointInMovingNotificationMessage, newEndpoint)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"newEndpoint": fmt.Sprintf("%v", newEndpoint),
|
||||
})
|
||||
}
|
||||
|
||||
func NoConnectionInHandlerContext(notificationType string) string {
|
||||
message := fmt.Sprintf("%s for %s notification", NoConnectionInHandlerContextMessage, notificationType)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"notificationType": notificationType,
|
||||
})
|
||||
}
|
||||
|
||||
func InvalidConnectionTypeInHandlerContext(notificationType string, conn interface{}, handlerCtx interface{}) string {
|
||||
message := fmt.Sprintf("%s for %s notification - %T %#v", InvalidConnectionTypeInHandlerContextMessage, notificationType, conn, handlerCtx)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"notificationType": notificationType,
|
||||
"connType": fmt.Sprintf("%T", conn),
|
||||
})
|
||||
}
|
||||
|
||||
func SchedulingHandoffToCurrentEndpoint(connID uint64, seconds float64) string {
|
||||
message := fmt.Sprintf("conn[%d] %s in %v seconds", connID, SchedulingHandoffToCurrentEndpointMessage, seconds)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
"seconds": seconds,
|
||||
})
|
||||
}
|
||||
|
||||
func ManagerNotInitialized() string {
|
||||
return appendJSONIfDebug(ManagerNotInitializedMessage, map[string]interface{}{})
|
||||
}
|
||||
|
||||
func FailedToRegisterHandler(notificationType string, err error) string {
|
||||
message := fmt.Sprintf("%s for %s: %v", FailedToRegisterHandlerMessage, notificationType, err)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"notificationType": notificationType,
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
|
||||
func ShutdownError() string {
|
||||
return appendJSONIfDebug(ShutdownErrorMessage, map[string]interface{}{})
|
||||
}
|
||||
|
||||
// Configuration validation error functions
|
||||
func InvalidRelaxedTimeoutError() string {
|
||||
return appendJSONIfDebug(InvalidRelaxedTimeoutErrorMessage, map[string]interface{}{})
|
||||
}
|
||||
|
||||
func InvalidHandoffTimeoutError() string {
|
||||
return appendJSONIfDebug(InvalidHandoffTimeoutErrorMessage, map[string]interface{}{})
|
||||
}
|
||||
|
||||
func InvalidHandoffWorkersError() string {
|
||||
return appendJSONIfDebug(InvalidHandoffWorkersErrorMessage, map[string]interface{}{})
|
||||
}
|
||||
|
||||
func InvalidHandoffQueueSizeError() string {
|
||||
return appendJSONIfDebug(InvalidHandoffQueueSizeErrorMessage, map[string]interface{}{})
|
||||
}
|
||||
|
||||
func InvalidPostHandoffRelaxedDurationError() string {
|
||||
return appendJSONIfDebug(InvalidPostHandoffRelaxedDurationErrorMessage, map[string]interface{}{})
|
||||
}
|
||||
|
||||
func InvalidEndpointTypeError() string {
|
||||
return appendJSONIfDebug(InvalidEndpointTypeErrorMessage, map[string]interface{}{})
|
||||
}
|
||||
|
||||
func InvalidMaintNotificationsError() string {
|
||||
return appendJSONIfDebug(InvalidMaintNotificationsErrorMessage, map[string]interface{}{})
|
||||
}
|
||||
|
||||
func InvalidHandoffRetriesError() string {
|
||||
return appendJSONIfDebug(InvalidHandoffRetriesErrorMessage, map[string]interface{}{})
|
||||
}
|
||||
|
||||
func InvalidClientError() string {
|
||||
return appendJSONIfDebug(InvalidClientErrorMessage, map[string]interface{}{})
|
||||
}
|
||||
|
||||
func InvalidNotificationError() string {
|
||||
return appendJSONIfDebug(InvalidNotificationErrorMessage, map[string]interface{}{})
|
||||
}
|
||||
|
||||
func MaxHandoffRetriesReachedError() string {
|
||||
return appendJSONIfDebug(MaxHandoffRetriesReachedErrorMessage, map[string]interface{}{})
|
||||
}
|
||||
|
||||
func HandoffQueueFullError() string {
|
||||
return appendJSONIfDebug(HandoffQueueFullErrorMessage, map[string]interface{}{})
|
||||
}
|
||||
|
||||
func InvalidCircuitBreakerFailureThresholdError() string {
|
||||
return appendJSONIfDebug(InvalidCircuitBreakerFailureThresholdErrorMessage, map[string]interface{}{})
|
||||
}
|
||||
|
||||
func InvalidCircuitBreakerResetTimeoutError() string {
|
||||
return appendJSONIfDebug(InvalidCircuitBreakerResetTimeoutErrorMessage, map[string]interface{}{})
|
||||
}
|
||||
|
||||
func InvalidCircuitBreakerMaxRequestsError() string {
|
||||
return appendJSONIfDebug(InvalidCircuitBreakerMaxRequestsErrorMessage, map[string]interface{}{})
|
||||
}
|
||||
|
||||
// Configuration and debug functions
|
||||
func DebugLoggingEnabled() string {
|
||||
return appendJSONIfDebug(DebugLoggingEnabledMessage, map[string]interface{}{})
|
||||
}
|
||||
|
||||
func ConfigDebug(config interface{}) string {
|
||||
message := fmt.Sprintf("%s: %+v", ConfigDebugMessage, config)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"config": fmt.Sprintf("%+v", config),
|
||||
})
|
||||
}
|
||||
|
||||
// Handoff worker functions
|
||||
func WorkerExitingDueToShutdown() string {
|
||||
return appendJSONIfDebug(WorkerExitingDueToShutdownMessage, map[string]interface{}{})
|
||||
}
|
||||
|
||||
func WorkerExitingDueToShutdownWhileProcessing() string {
|
||||
return appendJSONIfDebug(WorkerExitingDueToShutdownWhileProcessingMessage, map[string]interface{}{})
|
||||
}
|
||||
|
||||
func WorkerPanicRecovered(panicValue interface{}) string {
|
||||
message := fmt.Sprintf("%s: %v", WorkerPanicRecoveredMessage, panicValue)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"panic": fmt.Sprintf("%v", panicValue),
|
||||
})
|
||||
}
|
||||
|
||||
func WorkerExitingDueToInactivityTimeout(timeout interface{}) string {
|
||||
message := fmt.Sprintf("%s (%v)", WorkerExitingDueToInactivityTimeoutMessage, timeout)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"timeout": fmt.Sprintf("%v", timeout),
|
||||
})
|
||||
}
|
||||
|
||||
func ApplyingRelaxedTimeoutDueToPostHandoff(connID uint64, timeout interface{}, until string) string {
|
||||
message := fmt.Sprintf("conn[%d] %s (%v) until %s", connID, ApplyingRelaxedTimeoutDueToPostHandoffMessage, timeout, until)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
"timeout": fmt.Sprintf("%v", timeout),
|
||||
"until": until,
|
||||
})
|
||||
}
|
||||
|
||||
// Example hooks functions
|
||||
func MetricsHookProcessingNotification(notificationType string, connID uint64) string {
|
||||
message := fmt.Sprintf("%s %s notification on conn[%d]", MetricsHookProcessingNotificationMessage, notificationType, connID)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"notificationType": notificationType,
|
||||
"connID": connID,
|
||||
})
|
||||
}
|
||||
|
||||
func MetricsHookRecordedError(notificationType string, connID uint64, err error) string {
|
||||
message := fmt.Sprintf("%s for %s notification on conn[%d]: %v", MetricsHookRecordedErrorMessage, notificationType, connID, err)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"notificationType": notificationType,
|
||||
"connID": connID,
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
|
||||
// Pool hook functions
|
||||
func MarkedForHandoff(connID uint64) string {
|
||||
message := fmt.Sprintf("conn[%d] %s", connID, MarkedForHandoffMessage)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"connID": connID,
|
||||
})
|
||||
}
|
||||
|
||||
// Circuit breaker additional functions
|
||||
func CircuitBreakerTransitioningToHalfOpen(endpoint string) string {
|
||||
message := fmt.Sprintf("%s for %s", CircuitBreakerTransitioningToHalfOpenMessage, endpoint)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"endpoint": endpoint,
|
||||
})
|
||||
}
|
||||
|
||||
func CircuitBreakerOpened(endpoint string, failures int64) string {
|
||||
message := fmt.Sprintf("%s for endpoint %s after %d failures", CircuitBreakerOpenedMessage, endpoint, failures)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"endpoint": endpoint,
|
||||
"failures": failures,
|
||||
})
|
||||
}
|
||||
|
||||
func CircuitBreakerReopened(endpoint string) string {
|
||||
message := fmt.Sprintf("%s for endpoint %s due to failure in half-open state", CircuitBreakerReopenedMessage, endpoint)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"endpoint": endpoint,
|
||||
})
|
||||
}
|
||||
|
||||
func CircuitBreakerClosed(endpoint string, successes int64) string {
|
||||
message := fmt.Sprintf("%s for endpoint %s after %d successful requests", CircuitBreakerClosedMessage, endpoint, successes)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"endpoint": endpoint,
|
||||
"successes": successes,
|
||||
})
|
||||
}
|
||||
|
||||
func CircuitBreakerCleanup(removed int, total int) string {
|
||||
message := fmt.Sprintf("%s removed %d/%d entries", CircuitBreakerCleanupMessage, removed, total)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"removed": removed,
|
||||
"total": total,
|
||||
})
|
||||
}
|
||||
|
||||
// ExtractDataFromLogMessage extracts structured data from maintnotifications log messages
|
||||
// Returns a map containing the parsed key-value pairs from the structured data section
|
||||
// Example: "conn[123] handoff started to localhost:6379 {"connID":123,"endpoint":"localhost:6379"}"
|
||||
// Returns: map[string]interface{}{"connID": 123, "endpoint": "localhost:6379"}
|
||||
func ExtractDataFromLogMessage(logMessage string) map[string]interface{} {
|
||||
result := make(map[string]interface{})
|
||||
|
||||
// Find the JSON data section at the end of the message
|
||||
re := regexp.MustCompile(`(\{.*\})$`)
|
||||
matches := re.FindStringSubmatch(logMessage)
|
||||
if len(matches) < 2 {
|
||||
return result
|
||||
}
|
||||
|
||||
jsonStr := matches[1]
|
||||
if jsonStr == "" {
|
||||
return result
|
||||
}
|
||||
|
||||
// Parse the JSON directly
|
||||
var jsonResult map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(jsonStr), &jsonResult); err == nil {
|
||||
return jsonResult
|
||||
}
|
||||
|
||||
// If JSON parsing fails, return empty map
|
||||
return result
|
||||
}
|
||||
|
||||
// Cluster notification functions
|
||||
func InvalidSeqIDInSMigratingNotification(seqID interface{}) string {
|
||||
message := fmt.Sprintf("%s: %v", InvalidSeqIDInSMigratingNotificationMessage, seqID)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"seqID": fmt.Sprintf("%v", seqID),
|
||||
})
|
||||
}
|
||||
|
||||
func InvalidSeqIDInSMigratedNotification(seqID interface{}) string {
|
||||
message := fmt.Sprintf("%s: %v", InvalidSeqIDInSMigratedNotificationMessage, seqID)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"seqID": fmt.Sprintf("%v", seqID),
|
||||
})
|
||||
}
|
||||
|
||||
// TriggeringClusterStateReload logs when cluster state reload is triggered (deduplicated, once per seqID)
|
||||
func TriggeringClusterStateReload(seqID int64, hostPort string, slotRanges []string) string {
|
||||
message := fmt.Sprintf("%s seqID=%d host:port=%s slots=%v", TriggeringClusterStateReloadMessage, seqID, hostPort, slotRanges)
|
||||
return appendJSONIfDebug(message, map[string]interface{}{
|
||||
"seqID": seqID,
|
||||
"hostPort": hostPort,
|
||||
"slotRanges": slotRanges,
|
||||
})
|
||||
}
|
||||
+63
@@ -0,0 +1,63 @@
|
||||
/*
|
||||
Copyright 2014 The Camlistore 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 internal
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
// A Once will perform a successful action exactly once.
|
||||
//
|
||||
// Unlike a sync.Once, this Once's func returns an error
|
||||
// and is re-armed on failure.
|
||||
type Once struct {
|
||||
m sync.Mutex
|
||||
done uint32
|
||||
}
|
||||
|
||||
// Do calls the function f if and only if Do has not been invoked
|
||||
// without error for this instance of Once. In other words, given
|
||||
//
|
||||
// var once Once
|
||||
//
|
||||
// if once.Do(f) is called multiple times, only the first call will
|
||||
// invoke f, even if f has a different value in each invocation unless
|
||||
// f returns an error. A new instance of Once is required for each
|
||||
// function to execute.
|
||||
//
|
||||
// Do is intended for initialization that must be run exactly once. Since f
|
||||
// is niladic, it may be necessary to use a function literal to capture the
|
||||
// arguments to a function to be invoked by Do:
|
||||
//
|
||||
// err := config.once.Do(func() error { return config.init(filename) })
|
||||
func (o *Once) Do(f func() error) error {
|
||||
if atomic.LoadUint32(&o.done) == 1 {
|
||||
return nil
|
||||
}
|
||||
// Slow-path.
|
||||
o.m.Lock()
|
||||
defer o.m.Unlock()
|
||||
var err error
|
||||
if o.done == 0 {
|
||||
err = f()
|
||||
if err == nil {
|
||||
atomic.StoreUint32(&o.done, 1)
|
||||
}
|
||||
}
|
||||
return err
|
||||
}
|
||||
+298
@@ -0,0 +1,298 @@
|
||||
package otel
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9/internal/pool"
|
||||
)
|
||||
|
||||
// generateUniqueID generates a short unique identifier for pool names.
|
||||
func generateUniqueID() string {
|
||||
b := make([]byte, 4)
|
||||
if _, err := rand.Read(b); err != nil {
|
||||
return ""
|
||||
}
|
||||
return hex.EncodeToString(b)
|
||||
}
|
||||
|
||||
// Cmder is a minimal interface for command information needed for metrics.
|
||||
// This avoids circular dependencies with the main redis package.
|
||||
type Cmder interface {
|
||||
Name() string
|
||||
FullName() string
|
||||
Args() []interface{}
|
||||
Err() error
|
||||
}
|
||||
|
||||
// Recorder is the interface for recording metrics.
|
||||
type Recorder interface {
|
||||
// RecordOperationDuration records the total operation duration (including all retries)
|
||||
// dbIndex is the Redis database index (0-15)
|
||||
RecordOperationDuration(ctx context.Context, duration time.Duration, cmd Cmder, attempts int, err error, cn *pool.Conn, dbIndex int)
|
||||
|
||||
// RecordPipelineOperationDuration records the total pipeline/transaction duration.
|
||||
// operationName should be "PIPELINE" for regular pipelines or "MULTI" for transactions.
|
||||
// cmdCount is the number of commands in the pipeline.
|
||||
// err is the error from the pipeline execution (can be nil).
|
||||
// dbIndex is the Redis database index (0-15)
|
||||
RecordPipelineOperationDuration(ctx context.Context, duration time.Duration, operationName string, cmdCount int, attempts int, err error, cn *pool.Conn, dbIndex int)
|
||||
|
||||
// RecordConnectionCreateTime records the time it took to create a new connection
|
||||
RecordConnectionCreateTime(ctx context.Context, duration time.Duration, cn *pool.Conn)
|
||||
|
||||
// RecordConnectionRelaxedTimeout records when connection timeout is relaxed/unrelaxed
|
||||
// delta: +1 for relaxed, -1 for unrelaxed
|
||||
// poolName: name of the connection pool (e.g., "main", "pubsub")
|
||||
// notificationType: the notification type that triggered the timeout relaxation (e.g., "MOVING")
|
||||
RecordConnectionRelaxedTimeout(ctx context.Context, delta int, cn *pool.Conn, poolName, notificationType string)
|
||||
|
||||
// RecordConnectionHandoff records when a connection is handed off to another node
|
||||
// poolName: name of the connection pool (e.g., "main", "pubsub")
|
||||
RecordConnectionHandoff(ctx context.Context, cn *pool.Conn, poolName string)
|
||||
|
||||
// RecordError records client errors (ASK, MOVED, handshake failures, etc.)
|
||||
// errorType: type of error (e.g., "ASK", "MOVED", "HANDSHAKE_FAILED")
|
||||
// statusCode: Redis response status code if available (e.g., "MOVED", "ASK")
|
||||
// isInternal: whether this is an internal error
|
||||
// retryAttempts: number of retry attempts made
|
||||
RecordError(ctx context.Context, errorType string, cn *pool.Conn, statusCode string, isInternal bool, retryAttempts int)
|
||||
|
||||
// RecordMaintenanceNotification records when a maintenance notification is received
|
||||
// notificationType: the type of notification (e.g., "MOVING", "MIGRATING", etc.)
|
||||
RecordMaintenanceNotification(ctx context.Context, cn *pool.Conn, notificationType string)
|
||||
|
||||
// RecordConnectionWaitTime records the time spent waiting for a connection from the pool
|
||||
RecordConnectionWaitTime(ctx context.Context, duration time.Duration, cn *pool.Conn)
|
||||
|
||||
// RecordConnectionClosed records when a connection is closed
|
||||
// reason: reason for closing (e.g., "idle", "max_lifetime", "error", "pool_closed")
|
||||
// err: the error that caused the close (nil for non-error closures)
|
||||
RecordConnectionClosed(ctx context.Context, cn *pool.Conn, reason string, err error)
|
||||
|
||||
// RecordPubSubMessage records a Pub/Sub message
|
||||
// direction: "sent" or "received"
|
||||
// channel: channel name (may be hidden for cardinality reduction)
|
||||
// sharded: true for sharded pub/sub (SPUBLISH/SSUBSCRIBE)
|
||||
RecordPubSubMessage(ctx context.Context, cn *pool.Conn, direction, channel string, sharded bool)
|
||||
|
||||
// RecordStreamLag records the lag for stream consumer group processing
|
||||
// lag: time difference between message creation and consumption
|
||||
// streamName: name of the stream (may be hidden for cardinality reduction)
|
||||
// consumerGroup: name of the consumer group
|
||||
// consumerName: name of the consumer
|
||||
RecordStreamLag(ctx context.Context, lag time.Duration, cn *pool.Conn, streamName, consumerGroup, consumerName string)
|
||||
|
||||
// RecordConnectionCount records a change in connection count (UpDownCounter)
|
||||
// delta: +1 when connection added, -1 when connection removed
|
||||
// state: connection state (e.g., "idle", "used")
|
||||
// isPubSub: true if this is a PubSub connection
|
||||
RecordConnectionCount(ctx context.Context, delta int, cn *pool.Conn, state string, isPubSub bool)
|
||||
|
||||
// RecordPendingRequests records a change in pending requests (UpDownCounter)
|
||||
// delta: +1 when request starts waiting, -1 when request stops waiting
|
||||
// poolName is passed explicitly because we may not have a connection yet when request starts
|
||||
RecordPendingRequests(ctx context.Context, delta int, cn *pool.Conn, poolName string)
|
||||
}
|
||||
|
||||
type PubSubPooler interface {
|
||||
Stats() *pool.PubSubStats
|
||||
}
|
||||
|
||||
type PoolRegistrar interface {
|
||||
// RegisterPool is called when a new client is created with its connection pools.
|
||||
// poolName: identifier for the pool (e.g., "main_abc123")
|
||||
// pool: the connection pool
|
||||
RegisterPool(poolName string, pool pool.Pooler)
|
||||
// UnregisterPool is called when a client is closed to remove its pool from the registry.
|
||||
// pool: the connection pool to unregister
|
||||
UnregisterPool(pool pool.Pooler)
|
||||
// RegisterPubSubPool is called when a new client is created with a PubSub pool.
|
||||
// poolName: identifier for the pool (e.g., "main_abc123_pubsub")
|
||||
// pool: the PubSub connection pool
|
||||
RegisterPubSubPool(poolName string, pool PubSubPooler)
|
||||
// UnregisterPubSubPool is called when a PubSub client is closed to remove its pool.
|
||||
// pool: the PubSub connection pool to unregister
|
||||
UnregisterPubSubPool(pool PubSubPooler)
|
||||
}
|
||||
|
||||
var (
|
||||
// recorderMu protects globalRecorder and operation duration callbacks
|
||||
recorderMu sync.RWMutex
|
||||
|
||||
// Global recorder instance (initialized by extra/redisotel-native)
|
||||
globalRecorder Recorder = noopRecorder{}
|
||||
|
||||
// Callbacks for operation duration metrics
|
||||
operationDurationCallback func(ctx context.Context, duration time.Duration, cmd Cmder, attempts int, err error, cn *pool.Conn, dbIndex int)
|
||||
pipelineOperationDurationCallback func(ctx context.Context, duration time.Duration, operationName string, cmdCount int, attempts int, err error, cn *pool.Conn, dbIndex int)
|
||||
)
|
||||
|
||||
// GetOperationDurationCallback returns the callback for operation duration.
|
||||
func GetOperationDurationCallback() func(ctx context.Context, duration time.Duration, cmd Cmder, attempts int, err error, cn *pool.Conn, dbIndex int) {
|
||||
recorderMu.RLock()
|
||||
cb := operationDurationCallback
|
||||
recorderMu.RUnlock()
|
||||
return cb
|
||||
}
|
||||
|
||||
// GetPipelineOperationDurationCallback returns the callback for pipeline operation duration.
|
||||
func GetPipelineOperationDurationCallback() func(ctx context.Context, duration time.Duration, operationName string, cmdCount int, attempts int, err error, cn *pool.Conn, dbIndex int) {
|
||||
recorderMu.RLock()
|
||||
cb := pipelineOperationDurationCallback
|
||||
recorderMu.RUnlock()
|
||||
return cb
|
||||
}
|
||||
|
||||
// getRecorder returns the current global recorder under a read lock.
|
||||
func getRecorder() Recorder {
|
||||
recorderMu.RLock()
|
||||
r := globalRecorder
|
||||
recorderMu.RUnlock()
|
||||
return r
|
||||
}
|
||||
|
||||
// SetGlobalRecorder sets the global recorder (called by Init() in extra/redisotel-native)
|
||||
func SetGlobalRecorder(r Recorder) {
|
||||
recorderMu.Lock()
|
||||
if r == nil {
|
||||
globalRecorder = noopRecorder{}
|
||||
operationDurationCallback = nil
|
||||
pipelineOperationDurationCallback = nil
|
||||
recorderMu.Unlock()
|
||||
// Unregister all pool metric callbacks atomically
|
||||
pool.SetAllMetricCallbacks(nil)
|
||||
return
|
||||
}
|
||||
globalRecorder = r
|
||||
|
||||
// Register operation duration callbacks
|
||||
// These capture r directly since we want them to use the specific recorder
|
||||
// that was set at this point in time
|
||||
operationDurationCallback = func(ctx context.Context, duration time.Duration, cmd Cmder, attempts int, err error, cn *pool.Conn, dbIndex int) {
|
||||
getRecorder().RecordOperationDuration(ctx, duration, cmd, attempts, err, cn, dbIndex)
|
||||
}
|
||||
pipelineOperationDurationCallback = func(ctx context.Context, duration time.Duration, operationName string, cmdCount int, attempts int, err error, cn *pool.Conn, dbIndex int) {
|
||||
getRecorder().RecordPipelineOperationDuration(ctx, duration, operationName, cmdCount, attempts, err, cn, dbIndex)
|
||||
}
|
||||
recorderMu.Unlock()
|
||||
|
||||
// Register all pool metric callbacks atomically
|
||||
// These use getRecorder() to safely access the current recorder
|
||||
pool.SetAllMetricCallbacks(&pool.MetricCallbacks{
|
||||
ConnectionCreateTime: func(ctx context.Context, duration time.Duration, cn *pool.Conn) {
|
||||
getRecorder().RecordConnectionCreateTime(ctx, duration, cn)
|
||||
},
|
||||
ConnectionRelaxedTimeout: func(ctx context.Context, delta int, cn *pool.Conn, poolName, notificationType string) {
|
||||
getRecorder().RecordConnectionRelaxedTimeout(ctx, delta, cn, poolName, notificationType)
|
||||
},
|
||||
ConnectionHandoff: func(ctx context.Context, cn *pool.Conn, poolName string) {
|
||||
getRecorder().RecordConnectionHandoff(ctx, cn, poolName)
|
||||
},
|
||||
Error: func(ctx context.Context, errorType string, cn *pool.Conn, statusCode string, isInternal bool, retryAttempts int) {
|
||||
getRecorder().RecordError(ctx, errorType, cn, statusCode, isInternal, retryAttempts)
|
||||
},
|
||||
MaintenanceNotification: func(ctx context.Context, cn *pool.Conn, notificationType string) {
|
||||
getRecorder().RecordMaintenanceNotification(ctx, cn, notificationType)
|
||||
},
|
||||
ConnectionWaitTime: func(ctx context.Context, duration time.Duration, cn *pool.Conn) {
|
||||
getRecorder().RecordConnectionWaitTime(ctx, duration, cn)
|
||||
},
|
||||
ConnectionClosed: func(ctx context.Context, cn *pool.Conn, reason string, err error) {
|
||||
getRecorder().RecordConnectionClosed(ctx, cn, reason, err)
|
||||
},
|
||||
ConnectionCount: func(ctx context.Context, delta int, cn *pool.Conn, state string, isPubSub bool) {
|
||||
getRecorder().RecordConnectionCount(ctx, delta, cn, state, isPubSub)
|
||||
},
|
||||
PendingRequests: func(ctx context.Context, delta int, cn *pool.Conn, poolName string) {
|
||||
getRecorder().RecordPendingRequests(ctx, delta, cn, poolName)
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// RecordOperationDuration records the total operation duration.
|
||||
// dbIndex is the Redis database index (0-15).
|
||||
func RecordOperationDuration(ctx context.Context, duration time.Duration, cmd Cmder, attempts int, err error, cn *pool.Conn, dbIndex int) {
|
||||
getRecorder().RecordOperationDuration(ctx, duration, cmd, attempts, err, cn, dbIndex)
|
||||
}
|
||||
|
||||
// RecordPipelineOperationDuration records the total pipeline/transaction duration.
|
||||
// This is called from redis.go after pipeline/transaction execution completes.
|
||||
// operationName should be "PIPELINE" for regular pipelines or "MULTI" for transactions.
|
||||
// err is the error from the pipeline execution (can be nil).
|
||||
// dbIndex is the Redis database index (0-15).
|
||||
func RecordPipelineOperationDuration(ctx context.Context, duration time.Duration, operationName string, cmdCount int, attempts int, err error, cn *pool.Conn, dbIndex int) {
|
||||
getRecorder().RecordPipelineOperationDuration(ctx, duration, operationName, cmdCount, attempts, err, cn, dbIndex)
|
||||
}
|
||||
|
||||
// RecordConnectionCreateTime records the time it took to create a new connection.
|
||||
func RecordConnectionCreateTime(ctx context.Context, duration time.Duration, cn *pool.Conn) {
|
||||
getRecorder().RecordConnectionCreateTime(ctx, duration, cn)
|
||||
}
|
||||
|
||||
// RecordPubSubMessage records a Pub/Sub message sent or received.
|
||||
func RecordPubSubMessage(ctx context.Context, cn *pool.Conn, direction, channel string, sharded bool) {
|
||||
getRecorder().RecordPubSubMessage(ctx, cn, direction, channel, sharded)
|
||||
}
|
||||
|
||||
// RecordStreamLag records the lag between message creation and consumption in a stream.
|
||||
func RecordStreamLag(ctx context.Context, lag time.Duration, cn *pool.Conn, streamName, consumerGroup, consumerName string) {
|
||||
getRecorder().RecordStreamLag(ctx, lag, cn, streamName, consumerGroup, consumerName)
|
||||
}
|
||||
|
||||
type noopRecorder struct{}
|
||||
|
||||
func (noopRecorder) RecordOperationDuration(context.Context, time.Duration, Cmder, int, error, *pool.Conn, int) {
|
||||
}
|
||||
func (noopRecorder) RecordPipelineOperationDuration(context.Context, time.Duration, string, int, int, error, *pool.Conn, int) {
|
||||
}
|
||||
func (noopRecorder) RecordConnectionCreateTime(context.Context, time.Duration, *pool.Conn) {}
|
||||
func (noopRecorder) RecordConnectionRelaxedTimeout(context.Context, int, *pool.Conn, string, string) {
|
||||
}
|
||||
func (noopRecorder) RecordConnectionHandoff(context.Context, *pool.Conn, string) {}
|
||||
func (noopRecorder) RecordError(context.Context, string, *pool.Conn, string, bool, int) {}
|
||||
func (noopRecorder) RecordMaintenanceNotification(context.Context, *pool.Conn, string) {}
|
||||
|
||||
func (noopRecorder) RecordConnectionWaitTime(context.Context, time.Duration, *pool.Conn) {}
|
||||
func (noopRecorder) RecordConnectionClosed(context.Context, *pool.Conn, string, error) {}
|
||||
|
||||
func (noopRecorder) RecordPubSubMessage(context.Context, *pool.Conn, string, string, bool) {}
|
||||
|
||||
func (noopRecorder) RecordStreamLag(context.Context, time.Duration, *pool.Conn, string, string, string) {
|
||||
}
|
||||
func (noopRecorder) RecordConnectionCount(context.Context, int, *pool.Conn, string, bool) {}
|
||||
func (noopRecorder) RecordPendingRequests(context.Context, int, *pool.Conn, string) {}
|
||||
|
||||
// RegisterPools registers connection pools with the global recorder.
|
||||
func RegisterPools(connPool pool.Pooler, pubSubPool PubSubPooler, addr string) {
|
||||
// Check if the global recorder implements PoolRegistrar
|
||||
if registrar, ok := globalRecorder.(PoolRegistrar); ok {
|
||||
// Generate a unique ID for this client's pools
|
||||
uniqueID := generateUniqueID()
|
||||
|
||||
if connPool != nil {
|
||||
poolName := addr + "_" + uniqueID
|
||||
registrar.RegisterPool(poolName, connPool)
|
||||
}
|
||||
if pubSubPool != nil {
|
||||
poolName := addr + "_" + uniqueID + "_pubsub"
|
||||
registrar.RegisterPubSubPool(poolName, pubSubPool)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// UnregisterPools removes connection pools from the global recorder
|
||||
func UnregisterPools(connPool pool.Pooler, pubSubPool PubSubPooler) {
|
||||
// Check if the global recorder implements PoolRegistrar
|
||||
if registrar, ok := globalRecorder.(PoolRegistrar); ok {
|
||||
if connPool != nil {
|
||||
registrar.UnregisterPool(connPool)
|
||||
}
|
||||
if pubSubPool != nil {
|
||||
registrar.UnregisterPubSubPool(pubSubPool)
|
||||
}
|
||||
}
|
||||
}
|
||||
+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
|
||||
}
|
||||
+838
@@ -0,0 +1,838 @@
|
||||
package proto
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"math/big"
|
||||
"strconv"
|
||||
|
||||
"github.com/redis/go-redis/v9/internal/util"
|
||||
)
|
||||
|
||||
// DefaultBufferSize is the default size for read/write buffers (32 KiB).
|
||||
const DefaultBufferSize = 32 * 1024
|
||||
|
||||
// redis resp protocol data type.
|
||||
const (
|
||||
RespStatus = '+' // +<string>\r\n
|
||||
RespError = '-' // -<string>\r\n
|
||||
RespString = '$' // $<length>\r\n<bytes>\r\n
|
||||
RespInt = ':' // :<number>\r\n
|
||||
RespNil = '_' // _\r\n
|
||||
RespFloat = ',' // ,<floating-point-number>\r\n (golang float)
|
||||
RespBool = '#' // true: #t\r\n false: #f\r\n
|
||||
RespBlobError = '!' // !<length>\r\n<bytes>\r\n
|
||||
RespVerbatim = '=' // =<length>\r\nFORMAT:<bytes>\r\n
|
||||
RespBigInt = '(' // (<big number>\r\n
|
||||
RespArray = '*' // *<len>\r\n... (same as resp2)
|
||||
RespMap = '%' // %<len>\r\n(key)\r\n(value)\r\n... (golang map)
|
||||
RespSet = '~' // ~<len>\r\n... (same as Array)
|
||||
RespAttr = '|' // |<len>\r\n(key)\r\n(value)\r\n... + command reply
|
||||
RespPush = '>' // ><len>\r\n... (same as Array)
|
||||
)
|
||||
|
||||
// Not used temporarily.
|
||||
// Redis has not used these two data types for the time being, and will implement them later.
|
||||
// Streamed = "EOF:"
|
||||
// StreamedAggregated = '?'
|
||||
|
||||
//------------------------------------------------------------------------------
|
||||
|
||||
const Nil = RedisError("redis: nil") // nolint:errname
|
||||
|
||||
type RedisError string
|
||||
|
||||
func (e RedisError) Error() string { return string(e) }
|
||||
|
||||
func (RedisError) RedisError() {}
|
||||
|
||||
func ParseErrorReply(line []byte) error {
|
||||
msg := string(line[1:])
|
||||
return parseTypedRedisError(msg)
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------------
|
||||
|
||||
type Reader struct {
|
||||
rd *bufio.Reader
|
||||
}
|
||||
|
||||
func NewReader(rd io.Reader) *Reader {
|
||||
return &Reader{
|
||||
rd: bufio.NewReaderSize(rd, DefaultBufferSize),
|
||||
}
|
||||
}
|
||||
|
||||
func NewReaderSize(rd io.Reader, size int) *Reader {
|
||||
return &Reader{
|
||||
rd: bufio.NewReaderSize(rd, size),
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Reader) Buffered() int {
|
||||
return r.rd.Buffered()
|
||||
}
|
||||
|
||||
func (r *Reader) Peek(n int) ([]byte, error) {
|
||||
return r.rd.Peek(n)
|
||||
}
|
||||
|
||||
func (r *Reader) Reset(rd io.Reader) {
|
||||
r.rd.Reset(rd)
|
||||
}
|
||||
|
||||
// PeekReplyType returns the data type of the next response without advancing the Reader,
|
||||
// and discard the attribute type.
|
||||
func (r *Reader) PeekReplyType() (byte, error) {
|
||||
b, err := r.rd.Peek(1)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if b[0] == RespAttr {
|
||||
if err = r.DiscardNext(); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return r.PeekReplyType()
|
||||
}
|
||||
return b[0], nil
|
||||
}
|
||||
|
||||
func (r *Reader) PeekPushNotificationName() (string, error) {
|
||||
// "prime" the buffer by peeking at the next byte
|
||||
c, err := r.Peek(1)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if c[0] != RespPush {
|
||||
return "", fmt.Errorf("redis: can't peek push notification name, next reply is not a push notification")
|
||||
}
|
||||
|
||||
// peek 36 bytes at most, should be enough to read the push notification name
|
||||
toPeek := 36
|
||||
buffered := r.Buffered()
|
||||
if buffered == 0 {
|
||||
return "", fmt.Errorf("redis: can't peek push notification name, no data available")
|
||||
}
|
||||
if buffered < toPeek {
|
||||
toPeek = buffered
|
||||
}
|
||||
buf, err := r.rd.Peek(toPeek)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if buf[0] != RespPush {
|
||||
return "", fmt.Errorf("redis: can't parse push notification: %q", buf)
|
||||
}
|
||||
|
||||
if len(buf) < 3 {
|
||||
return "", fmt.Errorf("redis: can't parse push notification: %q", buf)
|
||||
}
|
||||
|
||||
// remove push notification type
|
||||
buf = buf[1:]
|
||||
// remove first line - e.g. >2\r\n
|
||||
for i := 0; i < len(buf)-1; i++ {
|
||||
if buf[i] == '\r' && buf[i+1] == '\n' {
|
||||
buf = buf[i+2:]
|
||||
break
|
||||
} else {
|
||||
if buf[i] < '0' || buf[i] > '9' {
|
||||
return "", fmt.Errorf("redis: can't parse push notification: %q", buf)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(buf) < 2 {
|
||||
return "", fmt.Errorf("redis: can't parse push notification: %q", buf)
|
||||
}
|
||||
// next line should be $<length><string>\r\n or +<length><string>\r\n
|
||||
// should have the type of the push notification name and it's length
|
||||
if buf[0] != RespString && buf[0] != RespStatus {
|
||||
return "", fmt.Errorf("redis: can't parse push notification name: %q", buf)
|
||||
}
|
||||
typeOfName := buf[0]
|
||||
// remove the type of the push notification name
|
||||
buf = buf[1:]
|
||||
if typeOfName == RespString {
|
||||
// remove the length of the string
|
||||
if len(buf) < 2 {
|
||||
return "", fmt.Errorf("redis: can't parse push notification name: %q", buf)
|
||||
}
|
||||
for i := 0; i < len(buf)-1; i++ {
|
||||
if buf[i] == '\r' && buf[i+1] == '\n' {
|
||||
buf = buf[i+2:]
|
||||
break
|
||||
} else {
|
||||
if buf[i] < '0' || buf[i] > '9' {
|
||||
return "", fmt.Errorf("redis: can't parse push notification name: %q", buf)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if len(buf) < 2 {
|
||||
return "", fmt.Errorf("redis: can't parse push notification name: %q", buf)
|
||||
}
|
||||
// keep only the notification name
|
||||
for i := 0; i < len(buf)-1; i++ {
|
||||
if buf[i] == '\r' && buf[i+1] == '\n' {
|
||||
buf = buf[:i]
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
return util.BytesToString(buf), nil
|
||||
}
|
||||
|
||||
// ReadLine Return a valid reply, it will check the protocol or redis error,
|
||||
// and discard the attribute type.
|
||||
func (r *Reader) ReadLine() ([]byte, error) {
|
||||
line, err := r.readLine()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
switch line[0] {
|
||||
case RespError:
|
||||
return nil, ParseErrorReply(line)
|
||||
case RespNil:
|
||||
return nil, Nil
|
||||
case RespBlobError:
|
||||
var blobErr string
|
||||
blobErr, err = r.readStringReply(line)
|
||||
if err == nil {
|
||||
err = parseTypedRedisError(blobErr)
|
||||
}
|
||||
return nil, err
|
||||
case RespAttr:
|
||||
if err = r.Discard(line); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return r.ReadLine()
|
||||
}
|
||||
|
||||
// Compatible with RESP2
|
||||
if IsNilReply(line) {
|
||||
return nil, Nil
|
||||
}
|
||||
|
||||
return line, nil
|
||||
}
|
||||
|
||||
// readLine returns an error if:
|
||||
// - there is a pending read error;
|
||||
// - or line does not end with \r\n.
|
||||
func (r *Reader) readLine() ([]byte, error) {
|
||||
b, err := r.rd.ReadSlice('\n')
|
||||
if err != nil {
|
||||
if err != bufio.ErrBufferFull {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
full := make([]byte, len(b))
|
||||
copy(full, b)
|
||||
|
||||
b, err = r.rd.ReadBytes('\n')
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
full = append(full, b...) //nolint:makezero
|
||||
b = full
|
||||
}
|
||||
if len(b) <= 2 || b[len(b)-1] != '\n' || b[len(b)-2] != '\r' {
|
||||
return nil, fmt.Errorf("redis: invalid reply: %q", b)
|
||||
}
|
||||
return b[:len(b)-2], nil
|
||||
}
|
||||
|
||||
func (r *Reader) ReadReply() (interface{}, error) {
|
||||
line, err := r.ReadLine()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
switch line[0] {
|
||||
case RespStatus:
|
||||
return string(line[1:]), nil
|
||||
case RespInt:
|
||||
return util.ParseInt(line[1:], 10, 64)
|
||||
case RespFloat:
|
||||
return r.readFloat(line)
|
||||
case RespBool:
|
||||
return r.readBool(line)
|
||||
case RespBigInt:
|
||||
return r.readBigInt(line)
|
||||
|
||||
case RespString:
|
||||
return r.readStringReply(line)
|
||||
case RespVerbatim:
|
||||
return r.readVerb(line)
|
||||
|
||||
case RespArray, RespSet, RespPush:
|
||||
return r.readSlice(line)
|
||||
case RespMap:
|
||||
return r.readMap(line)
|
||||
}
|
||||
return nil, fmt.Errorf("redis: can't parse %.100q", line)
|
||||
}
|
||||
|
||||
func (r *Reader) readFloat(line []byte) (float64, error) {
|
||||
v := util.BytesToString(line[1:])
|
||||
switch v {
|
||||
case "inf":
|
||||
return math.Inf(1), nil
|
||||
case "-inf":
|
||||
return math.Inf(-1), nil
|
||||
case "nan", "-nan":
|
||||
return math.NaN(), nil
|
||||
}
|
||||
return strconv.ParseFloat(v, 64)
|
||||
}
|
||||
|
||||
func (r *Reader) readBool(line []byte) (bool, error) {
|
||||
switch util.BytesToString(line[1:]) {
|
||||
case "t":
|
||||
return true, nil
|
||||
case "f":
|
||||
return false, nil
|
||||
}
|
||||
return false, fmt.Errorf("redis: can't parse bool reply: %q", line)
|
||||
}
|
||||
|
||||
func (r *Reader) readBigInt(line []byte) (*big.Int, error) {
|
||||
i := new(big.Int)
|
||||
if i, ok := i.SetString(util.BytesToString(line[1:]), 10); ok {
|
||||
return i, nil
|
||||
}
|
||||
return nil, fmt.Errorf("redis: can't parse bigInt reply: %q", line)
|
||||
}
|
||||
|
||||
func (r *Reader) readStringReply(line []byte) (string, error) {
|
||||
n, err := replyLen(line)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
b := make([]byte, n+2)
|
||||
_, err = io.ReadFull(r.rd, b)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
return util.BytesToString(b[:n]), nil
|
||||
}
|
||||
|
||||
func (r *Reader) readVerb(line []byte) (string, error) {
|
||||
s, err := r.readStringReply(line)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if len(s) < 4 || s[3] != ':' {
|
||||
return "", fmt.Errorf("redis: can't parse verbatim string reply: %q", line)
|
||||
}
|
||||
return s[4:], nil
|
||||
}
|
||||
|
||||
func (r *Reader) readSlice(line []byte) ([]interface{}, error) {
|
||||
n, err := replyLen(line)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
val := make([]interface{}, n)
|
||||
for i := 0; i < len(val); i++ {
|
||||
v, err := r.ReadReply()
|
||||
if err != nil {
|
||||
if err == Nil {
|
||||
val[i] = nil
|
||||
continue
|
||||
}
|
||||
if err, ok := err.(RedisError); ok {
|
||||
val[i] = err
|
||||
continue
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
val[i] = v
|
||||
}
|
||||
return val, nil
|
||||
}
|
||||
|
||||
func (r *Reader) readMap(line []byte) (map[interface{}]interface{}, error) {
|
||||
n, err := replyLen(line)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
m := make(map[interface{}]interface{}, n)
|
||||
for i := 0; i < n; i++ {
|
||||
k, err := r.ReadReply()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
v, err := r.ReadReply()
|
||||
if err != nil {
|
||||
if err == Nil {
|
||||
m[k] = nil
|
||||
continue
|
||||
}
|
||||
if err, ok := err.(RedisError); ok {
|
||||
m[k] = err
|
||||
continue
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
m[k] = v
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
// -------------------------------
|
||||
|
||||
func (r *Reader) ReadInt() (int64, error) {
|
||||
line, err := r.ReadLine()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
switch line[0] {
|
||||
case RespInt, RespStatus:
|
||||
return util.ParseInt(line[1:], 10, 64)
|
||||
case RespString:
|
||||
s, err := r.readStringReply(line)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return util.ParseInt([]byte(s), 10, 64)
|
||||
case RespBigInt:
|
||||
b, err := r.readBigInt(line)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if !b.IsInt64() {
|
||||
return 0, fmt.Errorf("bigInt(%s) value out of range", b.String())
|
||||
}
|
||||
return b.Int64(), nil
|
||||
}
|
||||
return 0, fmt.Errorf("redis: can't parse int reply: %.100q", line)
|
||||
}
|
||||
|
||||
func (r *Reader) ReadUint() (uint64, error) {
|
||||
line, err := r.ReadLine()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
switch line[0] {
|
||||
case RespInt, RespStatus:
|
||||
return util.ParseUint(line[1:], 10, 64)
|
||||
case RespString:
|
||||
s, err := r.readStringReply(line)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return util.ParseUint([]byte(s), 10, 64)
|
||||
case RespBigInt:
|
||||
b, err := r.readBigInt(line)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if !b.IsUint64() {
|
||||
return 0, fmt.Errorf("bigInt(%s) value out of range", b.String())
|
||||
}
|
||||
return b.Uint64(), nil
|
||||
}
|
||||
return 0, fmt.Errorf("redis: can't parse uint reply: %.100q", line)
|
||||
}
|
||||
|
||||
func (r *Reader) ReadFloat() (float64, error) {
|
||||
line, err := r.ReadLine()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
switch line[0] {
|
||||
case RespFloat:
|
||||
return r.readFloat(line)
|
||||
case RespStatus:
|
||||
return strconv.ParseFloat(util.BytesToString(line[1:]), 64)
|
||||
case RespString:
|
||||
s, err := r.readStringReply(line)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return strconv.ParseFloat(s, 64)
|
||||
}
|
||||
return 0, fmt.Errorf("redis: can't parse float reply: %.100q", line)
|
||||
}
|
||||
|
||||
func (r *Reader) ReadString() (string, error) {
|
||||
line, err := r.ReadLine()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
switch line[0] {
|
||||
case RespStatus, RespInt, RespFloat:
|
||||
return string(line[1:]), nil
|
||||
case RespString:
|
||||
return r.readStringReply(line)
|
||||
case RespBool:
|
||||
b, err := r.readBool(line)
|
||||
return strconv.FormatBool(b), err
|
||||
case RespVerbatim:
|
||||
return r.readVerb(line)
|
||||
case RespBigInt:
|
||||
b, err := r.readBigInt(line)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return b.String(), nil
|
||||
}
|
||||
return "", fmt.Errorf("redis: can't parse reply=%.100q reading string", line)
|
||||
}
|
||||
|
||||
func (r *Reader) ReadBool() (bool, error) {
|
||||
s, err := r.ReadString()
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return s == "OK" || s == "1" || s == "true", nil
|
||||
}
|
||||
|
||||
func (r *Reader) ReadSlice() ([]interface{}, error) {
|
||||
line, err := r.ReadLine()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return r.readSlice(line)
|
||||
}
|
||||
|
||||
// ReadFixedArrayLen read fixed array length.
|
||||
func (r *Reader) ReadFixedArrayLen(fixedLen int) error {
|
||||
n, err := r.ReadArrayLen()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n != fixedLen {
|
||||
return fmt.Errorf("redis: got %d elements in the array, wanted %d", n, fixedLen)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ReadArrayLen Read and return the length of the array.
|
||||
func (r *Reader) ReadArrayLen() (int, error) {
|
||||
line, err := r.ReadLine()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
switch line[0] {
|
||||
case RespArray, RespSet, RespPush:
|
||||
return replyLen(line)
|
||||
default:
|
||||
return 0, fmt.Errorf("redis: can't parse array/set/push reply: %.100q", line)
|
||||
}
|
||||
}
|
||||
|
||||
// ReadFixedMapLen reads fixed map length.
|
||||
func (r *Reader) ReadFixedMapLen(fixedLen int) error {
|
||||
n, err := r.ReadMapLen()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n != fixedLen {
|
||||
return fmt.Errorf("redis: got %d elements in the map, wanted %d", n, fixedLen)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ReadMapLen reads the length of the map type.
|
||||
// If responding to the array type (RespArray/RespSet/RespPush),
|
||||
// it must be a multiple of 2 and return n/2.
|
||||
// Other types will return an error.
|
||||
func (r *Reader) ReadMapLen() (int, error) {
|
||||
line, err := r.ReadLine()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
switch line[0] {
|
||||
case RespMap:
|
||||
return replyLen(line)
|
||||
case RespArray, RespSet, RespPush:
|
||||
// Some commands and RESP2 protocol may respond to array types.
|
||||
n, err := replyLen(line)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if n%2 != 0 {
|
||||
return 0, fmt.Errorf("redis: the length of the array must be a multiple of 2, got: %d", n)
|
||||
}
|
||||
return n / 2, nil
|
||||
default:
|
||||
return 0, fmt.Errorf("redis: can't parse map reply: %.100q", line)
|
||||
}
|
||||
}
|
||||
|
||||
// DiscardNext read and discard the data represented by the next line.
|
||||
func (r *Reader) DiscardNext() error {
|
||||
line, err := r.readLine()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return r.Discard(line)
|
||||
}
|
||||
|
||||
// Discard the data represented by line.
|
||||
func (r *Reader) Discard(line []byte) (err error) {
|
||||
if len(line) == 0 {
|
||||
return errors.New("redis: invalid line")
|
||||
}
|
||||
switch line[0] {
|
||||
case RespStatus, RespError, RespInt, RespNil, RespFloat, RespBool, RespBigInt:
|
||||
return nil
|
||||
}
|
||||
|
||||
n, err := replyLen(line)
|
||||
if err != nil && err != Nil {
|
||||
return err
|
||||
}
|
||||
|
||||
switch line[0] {
|
||||
case RespBlobError, RespString, RespVerbatim:
|
||||
// +\r\n
|
||||
_, err = r.rd.Discard(n + 2)
|
||||
return err
|
||||
case RespArray, RespSet, RespPush:
|
||||
for i := 0; i < n; i++ {
|
||||
if err = r.DiscardNext(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
case RespMap, RespAttr:
|
||||
// Read key & value.
|
||||
for i := 0; i < n*2; i++ {
|
||||
if err = r.DiscardNext(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
return fmt.Errorf("redis: can't parse %.100q", line)
|
||||
}
|
||||
|
||||
func replyLen(line []byte) (n int, err error) {
|
||||
n, err = util.Atoi(line[1:])
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
if n < -1 {
|
||||
return 0, fmt.Errorf("redis: invalid reply: %q", line)
|
||||
}
|
||||
|
||||
switch line[0] {
|
||||
case RespString, RespVerbatim, RespBlobError,
|
||||
RespArray, RespSet, RespPush, RespMap, RespAttr:
|
||||
if n == -1 {
|
||||
return 0, Nil
|
||||
}
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
|
||||
// IsNilReply detects redis.Nil of RESP2.
|
||||
func IsNilReply(line []byte) bool {
|
||||
return len(line) == 3 &&
|
||||
(line[0] == RespString || line[0] == RespArray) &&
|
||||
line[1] == '-' && line[2] == '1'
|
||||
}
|
||||
|
||||
// ReadRawReply reads the next RESP reply and returns it as raw bytes without parsing.
|
||||
func (r *Reader) ReadRawReply() ([]byte, error) {
|
||||
return r.readRawReplyBuf(nil)
|
||||
}
|
||||
|
||||
func (r *Reader) readRawReplyBuf(buf []byte) ([]byte, error) {
|
||||
line, err := r.readLine()
|
||||
if err != nil {
|
||||
return buf, err
|
||||
}
|
||||
|
||||
buf = append(buf, line...)
|
||||
buf = append(buf, '\r', '\n')
|
||||
|
||||
switch line[0] {
|
||||
case RespStatus, RespError, RespInt, RespNil, RespFloat, RespBool, RespBigInt:
|
||||
return buf, nil
|
||||
|
||||
case RespString, RespVerbatim, RespBlobError:
|
||||
n, err := replyLen(line)
|
||||
if err != nil {
|
||||
if err == Nil {
|
||||
return buf, nil
|
||||
}
|
||||
return buf, err
|
||||
}
|
||||
curLen := len(buf)
|
||||
buf = append(buf, make([]byte, n+2)...)
|
||||
_, err = io.ReadFull(r.rd, buf[curLen:])
|
||||
return buf, err
|
||||
|
||||
case RespArray, RespSet, RespPush:
|
||||
n, err := replyLen(line)
|
||||
if err != nil {
|
||||
if err == Nil {
|
||||
return buf, nil
|
||||
}
|
||||
return buf, err
|
||||
}
|
||||
for i := 0; i < n; i++ {
|
||||
buf, err = r.readRawReplyBuf(buf)
|
||||
if err != nil {
|
||||
return buf, err
|
||||
}
|
||||
}
|
||||
return buf, nil
|
||||
|
||||
case RespMap:
|
||||
n, err := replyLen(line)
|
||||
if err != nil {
|
||||
if err == Nil {
|
||||
return buf, nil
|
||||
}
|
||||
return buf, err
|
||||
}
|
||||
for i := 0; i < n*2; i++ {
|
||||
buf, err = r.readRawReplyBuf(buf)
|
||||
if err != nil {
|
||||
return buf, err
|
||||
}
|
||||
}
|
||||
return buf, nil
|
||||
|
||||
case RespAttr:
|
||||
// Per RESP3 spec, an attribute is always followed by the actual command reply.
|
||||
// We need to read the attribute's key-value pairs AND the following reply.
|
||||
n, err := replyLen(line)
|
||||
if err != nil {
|
||||
if err == Nil {
|
||||
return buf, nil
|
||||
}
|
||||
return buf, err
|
||||
}
|
||||
// Read the attribute key-value pairs
|
||||
for i := 0; i < n*2; i++ {
|
||||
buf, err = r.readRawReplyBuf(buf)
|
||||
if err != nil {
|
||||
return buf, err
|
||||
}
|
||||
}
|
||||
// Read the command reply that follows the attribute
|
||||
return r.readRawReplyBuf(buf)
|
||||
}
|
||||
|
||||
return buf, fmt.Errorf("redis: can't read raw reply: %.100q", line)
|
||||
}
|
||||
|
||||
var crlf = []byte{'\r', '\n'}
|
||||
|
||||
// ReadRawReplyWriteTo streams the next RESP reply directly to w without intermediate allocations.
|
||||
// Returns the number of bytes written and any error encountered.
|
||||
func (r *Reader) ReadRawReplyWriteTo(w io.Writer) (int64, error) {
|
||||
return r.readRawReplyWriteTo(w)
|
||||
}
|
||||
|
||||
func (r *Reader) readRawReplyWriteTo(w io.Writer) (int64, error) {
|
||||
line, err := r.readLine()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
var written int64
|
||||
n, err := w.Write(line)
|
||||
written += int64(n)
|
||||
if err != nil {
|
||||
return written, err
|
||||
}
|
||||
n, err = w.Write(crlf)
|
||||
written += int64(n)
|
||||
if err != nil {
|
||||
return written, err
|
||||
}
|
||||
|
||||
switch line[0] {
|
||||
case RespStatus, RespError, RespInt, RespNil, RespFloat, RespBool, RespBigInt:
|
||||
return written, nil
|
||||
|
||||
case RespString, RespVerbatim, RespBlobError:
|
||||
dataLen, err := replyLen(line)
|
||||
if err != nil {
|
||||
if err == Nil {
|
||||
return written, nil
|
||||
}
|
||||
return written, err
|
||||
}
|
||||
copied, err := io.CopyN(w, r.rd, int64(dataLen)+2)
|
||||
written += copied
|
||||
return written, err
|
||||
|
||||
case RespArray, RespSet, RespPush:
|
||||
count, err := replyLen(line)
|
||||
if err != nil {
|
||||
if err == Nil {
|
||||
return written, nil
|
||||
}
|
||||
return written, err
|
||||
}
|
||||
for i := 0; i < count; i++ {
|
||||
n, err := r.readRawReplyWriteTo(w)
|
||||
written += n
|
||||
if err != nil {
|
||||
return written, err
|
||||
}
|
||||
}
|
||||
return written, nil
|
||||
|
||||
case RespMap:
|
||||
count, err := replyLen(line)
|
||||
if err != nil {
|
||||
if err == Nil {
|
||||
return written, nil
|
||||
}
|
||||
return written, err
|
||||
}
|
||||
for i := 0; i < count*2; i++ {
|
||||
n, err := r.readRawReplyWriteTo(w)
|
||||
written += n
|
||||
if err != nil {
|
||||
return written, err
|
||||
}
|
||||
}
|
||||
return written, nil
|
||||
|
||||
case RespAttr:
|
||||
// Per RESP3 spec, an attribute is always followed by the actual command reply.
|
||||
// We need to read the attribute's key-value pairs AND the following reply.
|
||||
count, err := replyLen(line)
|
||||
if err != nil {
|
||||
if err == Nil {
|
||||
return written, nil
|
||||
}
|
||||
return written, err
|
||||
}
|
||||
// Read the attribute key-value pairs
|
||||
for i := 0; i < count*2; i++ {
|
||||
n, err := r.readRawReplyWriteTo(w)
|
||||
written += n
|
||||
if err != nil {
|
||||
return written, err
|
||||
}
|
||||
}
|
||||
// Read the command reply that follows the attribute
|
||||
n, err := r.readRawReplyWriteTo(w)
|
||||
written += n
|
||||
return written, err
|
||||
}
|
||||
|
||||
return written, fmt.Errorf("redis: can't read raw reply: %.100q", line)
|
||||
}
|
||||
+539
@@ -0,0 +1,539 @@
|
||||
package proto
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Typed Redis errors for better error handling with wrapping support.
|
||||
// These errors maintain backward compatibility by keeping the same error messages.
|
||||
|
||||
// LoadingError is returned when Redis is loading the dataset in memory.
|
||||
type LoadingError struct {
|
||||
msg string
|
||||
}
|
||||
|
||||
func (e *LoadingError) Error() string {
|
||||
return e.msg
|
||||
}
|
||||
|
||||
func (e *LoadingError) RedisError() {}
|
||||
|
||||
// NewLoadingError creates a new LoadingError with the given message.
|
||||
func NewLoadingError(msg string) *LoadingError {
|
||||
return &LoadingError{msg: msg}
|
||||
}
|
||||
|
||||
// ReadOnlyError is returned when trying to write to a read-only replica.
|
||||
type ReadOnlyError struct {
|
||||
msg string
|
||||
}
|
||||
|
||||
func (e *ReadOnlyError) Error() string {
|
||||
return e.msg
|
||||
}
|
||||
|
||||
func (e *ReadOnlyError) RedisError() {}
|
||||
|
||||
// NewReadOnlyError creates a new ReadOnlyError with the given message.
|
||||
func NewReadOnlyError(msg string) *ReadOnlyError {
|
||||
return &ReadOnlyError{msg: msg}
|
||||
}
|
||||
|
||||
// MovedError is returned when a key has been moved to a different node in a cluster.
|
||||
type MovedError struct {
|
||||
msg string
|
||||
addr string
|
||||
}
|
||||
|
||||
func (e *MovedError) Error() string {
|
||||
return e.msg
|
||||
}
|
||||
|
||||
func (e *MovedError) RedisError() {}
|
||||
|
||||
// Addr returns the address of the node where the key has been moved.
|
||||
func (e *MovedError) Addr() string {
|
||||
return e.addr
|
||||
}
|
||||
|
||||
// NewMovedError creates a new MovedError with the given message and address.
|
||||
func NewMovedError(msg string, addr string) *MovedError {
|
||||
return &MovedError{msg: msg, addr: addr}
|
||||
}
|
||||
|
||||
// AskError is returned when a key is being migrated and the client should ask another node.
|
||||
type AskError struct {
|
||||
msg string
|
||||
addr string
|
||||
}
|
||||
|
||||
func (e *AskError) Error() string {
|
||||
return e.msg
|
||||
}
|
||||
|
||||
func (e *AskError) RedisError() {}
|
||||
|
||||
// Addr returns the address of the node to ask.
|
||||
func (e *AskError) Addr() string {
|
||||
return e.addr
|
||||
}
|
||||
|
||||
// NewAskError creates a new AskError with the given message and address.
|
||||
func NewAskError(msg string, addr string) *AskError {
|
||||
return &AskError{msg: msg, addr: addr}
|
||||
}
|
||||
|
||||
// ClusterDownError is returned when the cluster is down.
|
||||
type ClusterDownError struct {
|
||||
msg string
|
||||
}
|
||||
|
||||
func (e *ClusterDownError) Error() string {
|
||||
return e.msg
|
||||
}
|
||||
|
||||
func (e *ClusterDownError) RedisError() {}
|
||||
|
||||
// NewClusterDownError creates a new ClusterDownError with the given message.
|
||||
func NewClusterDownError(msg string) *ClusterDownError {
|
||||
return &ClusterDownError{msg: msg}
|
||||
}
|
||||
|
||||
// TryAgainError is returned when a command cannot be processed and should be retried.
|
||||
type TryAgainError struct {
|
||||
msg string
|
||||
}
|
||||
|
||||
func (e *TryAgainError) Error() string {
|
||||
return e.msg
|
||||
}
|
||||
|
||||
func (e *TryAgainError) RedisError() {}
|
||||
|
||||
// NewTryAgainError creates a new TryAgainError with the given message.
|
||||
func NewTryAgainError(msg string) *TryAgainError {
|
||||
return &TryAgainError{msg: msg}
|
||||
}
|
||||
|
||||
// MasterDownError is returned when the master is down.
|
||||
type MasterDownError struct {
|
||||
msg string
|
||||
}
|
||||
|
||||
func (e *MasterDownError) Error() string {
|
||||
return e.msg
|
||||
}
|
||||
|
||||
func (e *MasterDownError) RedisError() {}
|
||||
|
||||
// NewMasterDownError creates a new MasterDownError with the given message.
|
||||
func NewMasterDownError(msg string) *MasterDownError {
|
||||
return &MasterDownError{msg: msg}
|
||||
}
|
||||
|
||||
// MaxClientsError is returned when the maximum number of clients has been reached.
|
||||
type MaxClientsError struct {
|
||||
msg string
|
||||
}
|
||||
|
||||
func (e *MaxClientsError) Error() string {
|
||||
return e.msg
|
||||
}
|
||||
|
||||
func (e *MaxClientsError) RedisError() {}
|
||||
|
||||
// NewMaxClientsError creates a new MaxClientsError with the given message.
|
||||
func NewMaxClientsError(msg string) *MaxClientsError {
|
||||
return &MaxClientsError{msg: msg}
|
||||
}
|
||||
|
||||
// AuthError is returned when authentication fails.
|
||||
type AuthError struct {
|
||||
msg string
|
||||
}
|
||||
|
||||
func (e *AuthError) Error() string {
|
||||
return e.msg
|
||||
}
|
||||
|
||||
func (e *AuthError) RedisError() {}
|
||||
|
||||
// NewAuthError creates a new AuthError with the given message.
|
||||
func NewAuthError(msg string) *AuthError {
|
||||
return &AuthError{msg: msg}
|
||||
}
|
||||
|
||||
// PermissionError is returned when a user lacks required permissions.
|
||||
type PermissionError struct {
|
||||
msg string
|
||||
}
|
||||
|
||||
func (e *PermissionError) Error() string {
|
||||
return e.msg
|
||||
}
|
||||
|
||||
func (e *PermissionError) RedisError() {}
|
||||
|
||||
// NewPermissionError creates a new PermissionError with the given message.
|
||||
func NewPermissionError(msg string) *PermissionError {
|
||||
return &PermissionError{msg: msg}
|
||||
}
|
||||
|
||||
// ExecAbortError is returned when a transaction is aborted.
|
||||
type ExecAbortError struct {
|
||||
msg string
|
||||
}
|
||||
|
||||
func (e *ExecAbortError) Error() string {
|
||||
return e.msg
|
||||
}
|
||||
|
||||
func (e *ExecAbortError) RedisError() {}
|
||||
|
||||
// NewExecAbortError creates a new ExecAbortError with the given message.
|
||||
func NewExecAbortError(msg string) *ExecAbortError {
|
||||
return &ExecAbortError{msg: msg}
|
||||
}
|
||||
|
||||
// OOMError is returned when Redis is out of memory.
|
||||
type OOMError struct {
|
||||
msg string
|
||||
}
|
||||
|
||||
func (e *OOMError) Error() string {
|
||||
return e.msg
|
||||
}
|
||||
|
||||
func (e *OOMError) RedisError() {}
|
||||
|
||||
// NewOOMError creates a new OOMError with the given message.
|
||||
func NewOOMError(msg string) *OOMError {
|
||||
return &OOMError{msg: msg}
|
||||
}
|
||||
|
||||
// NoReplicasError is returned when not enough replicas acknowledge a write.
|
||||
// This error occurs when using WAIT/WAITAOF commands or CLUSTER SETSLOT with
|
||||
// synchronous replication, and the required number of replicas cannot confirm
|
||||
// the write within the timeout period.
|
||||
type NoReplicasError struct {
|
||||
msg string
|
||||
}
|
||||
|
||||
func (e *NoReplicasError) Error() string {
|
||||
return e.msg
|
||||
}
|
||||
|
||||
func (e *NoReplicasError) RedisError() {}
|
||||
|
||||
// NewNoReplicasError creates a new NoReplicasError with the given message.
|
||||
func NewNoReplicasError(msg string) *NoReplicasError {
|
||||
return &NoReplicasError{msg: msg}
|
||||
}
|
||||
|
||||
// parseTypedRedisError parses a Redis error message and returns a typed error if applicable.
|
||||
// This function maintains backward compatibility by keeping the same error messages.
|
||||
func parseTypedRedisError(msg string) error {
|
||||
// Check for specific error patterns and return typed errors
|
||||
switch {
|
||||
case strings.HasPrefix(msg, "LOADING "):
|
||||
return NewLoadingError(msg)
|
||||
case strings.HasPrefix(msg, "READONLY "):
|
||||
return NewReadOnlyError(msg)
|
||||
case strings.HasPrefix(msg, "MOVED "):
|
||||
// Extract address from "MOVED <slot> <addr>"
|
||||
addr := extractAddr(msg)
|
||||
return NewMovedError(msg, addr)
|
||||
case strings.HasPrefix(msg, "ASK "):
|
||||
// Extract address from "ASK <slot> <addr>"
|
||||
addr := extractAddr(msg)
|
||||
return NewAskError(msg, addr)
|
||||
case strings.HasPrefix(msg, "CLUSTERDOWN "):
|
||||
return NewClusterDownError(msg)
|
||||
case strings.HasPrefix(msg, "TRYAGAIN "):
|
||||
return NewTryAgainError(msg)
|
||||
case strings.HasPrefix(msg, "MASTERDOWN "):
|
||||
return NewMasterDownError(msg)
|
||||
case strings.HasPrefix(msg, "NOREPLICAS "):
|
||||
return NewNoReplicasError(msg)
|
||||
case msg == "ERR max number of clients reached":
|
||||
return NewMaxClientsError(msg)
|
||||
case strings.HasPrefix(msg, "NOAUTH "), strings.HasPrefix(msg, "WRONGPASS "), strings.Contains(msg, "unauthenticated"):
|
||||
return NewAuthError(msg)
|
||||
case strings.HasPrefix(msg, "NOPERM "):
|
||||
return NewPermissionError(msg)
|
||||
case strings.HasPrefix(msg, "EXECABORT "):
|
||||
return NewExecAbortError(msg)
|
||||
case strings.HasPrefix(msg, "OOM "):
|
||||
return NewOOMError(msg)
|
||||
default:
|
||||
// Return generic RedisError for unknown error types
|
||||
return RedisError(msg)
|
||||
}
|
||||
}
|
||||
|
||||
// extractAddr extracts the address from MOVED/ASK error messages.
|
||||
// Format: "MOVED <slot> <addr>" or "ASK <slot> <addr>"
|
||||
func extractAddr(msg string) string {
|
||||
ind := strings.LastIndex(msg, " ")
|
||||
if ind == -1 {
|
||||
return ""
|
||||
}
|
||||
return msg[ind+1:]
|
||||
}
|
||||
|
||||
// IsLoadingError checks if an error is a LoadingError, even if wrapped.
|
||||
func IsLoadingError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
var loadingErr *LoadingError
|
||||
if errors.As(err, &loadingErr) {
|
||||
return true
|
||||
}
|
||||
// Check if wrapped error is a RedisError with LOADING prefix
|
||||
var redisErr RedisError
|
||||
if errors.As(err, &redisErr) && strings.HasPrefix(redisErr.Error(), "LOADING ") {
|
||||
return true
|
||||
}
|
||||
// Fallback to string checking for backward compatibility
|
||||
return strings.HasPrefix(err.Error(), "LOADING ")
|
||||
}
|
||||
|
||||
// IsReadOnlyError checks if an error is a ReadOnlyError, even if wrapped.
|
||||
func IsReadOnlyError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
var readOnlyErr *ReadOnlyError
|
||||
if errors.As(err, &readOnlyErr) {
|
||||
return true
|
||||
}
|
||||
// Check if wrapped error is a RedisError with READONLY prefix or Lua script READONLY
|
||||
var redisErr RedisError
|
||||
if errors.As(err, &redisErr) {
|
||||
s := redisErr.Error()
|
||||
if strings.HasPrefix(s, "READONLY ") {
|
||||
return true
|
||||
}
|
||||
// Lua script wrapped READONLY errors:
|
||||
// "ERR Error running script (call to f_<sha>): @user_script:N: -READONLY You can't write against a read only replica."
|
||||
if strings.Contains(s, "-READONLY You can't write against a read only replica") {
|
||||
return true
|
||||
}
|
||||
}
|
||||
// Fallback to string checking for backward compatibility
|
||||
s := err.Error()
|
||||
if strings.HasPrefix(s, "READONLY ") {
|
||||
return true
|
||||
}
|
||||
return strings.Contains(s, "-READONLY You can't write against a read only replica")
|
||||
}
|
||||
|
||||
// IsMovedError checks if an error is a MovedError, even if wrapped.
|
||||
// Returns the error and a boolean indicating if it's a MovedError.
|
||||
func IsMovedError(err error) (*MovedError, bool) {
|
||||
if err == nil {
|
||||
return nil, false
|
||||
}
|
||||
var movedErr *MovedError
|
||||
if errors.As(err, &movedErr) {
|
||||
return movedErr, true
|
||||
}
|
||||
// Fallback to string checking for backward compatibility
|
||||
s := err.Error()
|
||||
if strings.HasPrefix(s, "MOVED ") {
|
||||
// Parse: MOVED 3999 127.0.0.1:6381
|
||||
parts := strings.Split(s, " ")
|
||||
if len(parts) == 3 {
|
||||
return &MovedError{msg: s, addr: parts[2]}, true
|
||||
}
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
|
||||
// IsAskError checks if an error is an AskError, even if wrapped.
|
||||
// Returns the error and a boolean indicating if it's an AskError.
|
||||
func IsAskError(err error) (*AskError, bool) {
|
||||
if err == nil {
|
||||
return nil, false
|
||||
}
|
||||
var askErr *AskError
|
||||
if errors.As(err, &askErr) {
|
||||
return askErr, true
|
||||
}
|
||||
// Fallback to string checking for backward compatibility
|
||||
s := err.Error()
|
||||
if strings.HasPrefix(s, "ASK ") {
|
||||
// Parse: ASK 3999 127.0.0.1:6381
|
||||
parts := strings.Split(s, " ")
|
||||
if len(parts) == 3 {
|
||||
return &AskError{msg: s, addr: parts[2]}, true
|
||||
}
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
|
||||
// IsClusterDownError checks if an error is a ClusterDownError, even if wrapped.
|
||||
func IsClusterDownError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
var clusterDownErr *ClusterDownError
|
||||
if errors.As(err, &clusterDownErr) {
|
||||
return true
|
||||
}
|
||||
// Check if wrapped error is a RedisError with CLUSTERDOWN prefix
|
||||
var redisErr RedisError
|
||||
if errors.As(err, &redisErr) && strings.HasPrefix(redisErr.Error(), "CLUSTERDOWN ") {
|
||||
return true
|
||||
}
|
||||
// Fallback to string checking for backward compatibility
|
||||
return strings.HasPrefix(err.Error(), "CLUSTERDOWN ")
|
||||
}
|
||||
|
||||
// IsTryAgainError checks if an error is a TryAgainError, even if wrapped.
|
||||
func IsTryAgainError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
var tryAgainErr *TryAgainError
|
||||
if errors.As(err, &tryAgainErr) {
|
||||
return true
|
||||
}
|
||||
// Check if wrapped error is a RedisError with TRYAGAIN prefix
|
||||
var redisErr RedisError
|
||||
if errors.As(err, &redisErr) && strings.HasPrefix(redisErr.Error(), "TRYAGAIN ") {
|
||||
return true
|
||||
}
|
||||
// Fallback to string checking for backward compatibility
|
||||
return strings.HasPrefix(err.Error(), "TRYAGAIN ")
|
||||
}
|
||||
|
||||
// IsMasterDownError checks if an error is a MasterDownError, even if wrapped.
|
||||
func IsMasterDownError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
var masterDownErr *MasterDownError
|
||||
if errors.As(err, &masterDownErr) {
|
||||
return true
|
||||
}
|
||||
// Check if wrapped error is a RedisError with MASTERDOWN prefix
|
||||
var redisErr RedisError
|
||||
if errors.As(err, &redisErr) && strings.HasPrefix(redisErr.Error(), "MASTERDOWN ") {
|
||||
return true
|
||||
}
|
||||
// Fallback to string checking for backward compatibility
|
||||
return strings.HasPrefix(err.Error(), "MASTERDOWN ")
|
||||
}
|
||||
|
||||
// IsMaxClientsError checks if an error is a MaxClientsError, even if wrapped.
|
||||
func IsMaxClientsError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
var maxClientsErr *MaxClientsError
|
||||
if errors.As(err, &maxClientsErr) {
|
||||
return true
|
||||
}
|
||||
// Check if wrapped error is a RedisError with max clients prefix
|
||||
var redisErr RedisError
|
||||
if errors.As(err, &redisErr) && strings.HasPrefix(redisErr.Error(), "ERR max number of clients reached") {
|
||||
return true
|
||||
}
|
||||
// Fallback to string checking for backward compatibility
|
||||
return strings.HasPrefix(err.Error(), "ERR max number of clients reached")
|
||||
}
|
||||
|
||||
// IsAuthError checks if an error is an AuthError, even if wrapped.
|
||||
func IsAuthError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
var authErr *AuthError
|
||||
if errors.As(err, &authErr) {
|
||||
return true
|
||||
}
|
||||
// Check if wrapped error is a RedisError with auth error prefix
|
||||
var redisErr RedisError
|
||||
if errors.As(err, &redisErr) {
|
||||
s := redisErr.Error()
|
||||
return strings.HasPrefix(s, "NOAUTH ") || strings.HasPrefix(s, "WRONGPASS ") || strings.Contains(s, "unauthenticated")
|
||||
}
|
||||
// Fallback to string checking for backward compatibility
|
||||
s := err.Error()
|
||||
return strings.HasPrefix(s, "NOAUTH ") || strings.HasPrefix(s, "WRONGPASS ") || strings.Contains(s, "unauthenticated")
|
||||
}
|
||||
|
||||
// IsPermissionError checks if an error is a PermissionError, even if wrapped.
|
||||
func IsPermissionError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
var permErr *PermissionError
|
||||
if errors.As(err, &permErr) {
|
||||
return true
|
||||
}
|
||||
// Check if wrapped error is a RedisError with NOPERM prefix
|
||||
var redisErr RedisError
|
||||
if errors.As(err, &redisErr) && strings.HasPrefix(redisErr.Error(), "NOPERM ") {
|
||||
return true
|
||||
}
|
||||
// Fallback to string checking for backward compatibility
|
||||
return strings.HasPrefix(err.Error(), "NOPERM ")
|
||||
}
|
||||
|
||||
// IsExecAbortError checks if an error is an ExecAbortError, even if wrapped.
|
||||
func IsExecAbortError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
var execAbortErr *ExecAbortError
|
||||
if errors.As(err, &execAbortErr) {
|
||||
return true
|
||||
}
|
||||
// Check if wrapped error is a RedisError with EXECABORT prefix
|
||||
var redisErr RedisError
|
||||
if errors.As(err, &redisErr) && strings.HasPrefix(redisErr.Error(), "EXECABORT ") {
|
||||
return true
|
||||
}
|
||||
// Fallback to string checking for backward compatibility
|
||||
return strings.HasPrefix(err.Error(), "EXECABORT ")
|
||||
}
|
||||
|
||||
// IsOOMError checks if an error is an OOMError, even if wrapped.
|
||||
func IsOOMError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
var oomErr *OOMError
|
||||
if errors.As(err, &oomErr) {
|
||||
return true
|
||||
}
|
||||
// Check if wrapped error is a RedisError with OOM prefix
|
||||
var redisErr RedisError
|
||||
if errors.As(err, &redisErr) && strings.HasPrefix(redisErr.Error(), "OOM ") {
|
||||
return true
|
||||
}
|
||||
// Fallback to string checking for backward compatibility
|
||||
return strings.HasPrefix(err.Error(), "OOM ")
|
||||
}
|
||||
|
||||
// IsNoReplicasError checks if an error is a NoReplicasError, even if wrapped.
|
||||
func IsNoReplicasError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
var noReplicasErr *NoReplicasError
|
||||
if errors.As(err, &noReplicasErr) {
|
||||
return true
|
||||
}
|
||||
// Check if wrapped error is a RedisError with NOREPLICAS prefix
|
||||
var redisErr RedisError
|
||||
if errors.As(err, &redisErr) && strings.HasPrefix(redisErr.Error(), "NOREPLICAS ") {
|
||||
return true
|
||||
}
|
||||
// Fallback to string checking for backward compatibility
|
||||
return strings.HasPrefix(err.Error(), "NOREPLICAS ")
|
||||
}
|
||||
+185
@@ -0,0 +1,185 @@
|
||||
package proto
|
||||
|
||||
import (
|
||||
"encoding"
|
||||
"fmt"
|
||||
"net"
|
||||
"reflect"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9/internal/util"
|
||||
)
|
||||
|
||||
// Scan parses bytes `b` to `v` with appropriate type.
|
||||
//
|
||||
//nolint:gocyclo
|
||||
func Scan(b []byte, v interface{}) error {
|
||||
switch v := v.(type) {
|
||||
case nil:
|
||||
return fmt.Errorf("redis: Scan(nil)")
|
||||
case *string:
|
||||
*v = util.BytesToString(b)
|
||||
return nil
|
||||
case *[]byte:
|
||||
*v = b
|
||||
return nil
|
||||
case *int:
|
||||
var err error
|
||||
*v, err = util.Atoi(b)
|
||||
return err
|
||||
case *int8:
|
||||
n, err := util.ParseInt(b, 10, 8)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
*v = int8(n)
|
||||
return nil
|
||||
case *int16:
|
||||
n, err := util.ParseInt(b, 10, 16)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
*v = int16(n)
|
||||
return nil
|
||||
case *int32:
|
||||
n, err := util.ParseInt(b, 10, 32)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
*v = int32(n)
|
||||
return nil
|
||||
case *int64:
|
||||
n, err := util.ParseInt(b, 10, 64)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
*v = n
|
||||
return nil
|
||||
case *uint:
|
||||
n, err := util.ParseUint(b, 10, 64)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
*v = uint(n)
|
||||
return nil
|
||||
case *uint8:
|
||||
n, err := util.ParseUint(b, 10, 8)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
*v = uint8(n)
|
||||
return nil
|
||||
case *uint16:
|
||||
n, err := util.ParseUint(b, 10, 16)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
*v = uint16(n)
|
||||
return nil
|
||||
case *uint32:
|
||||
n, err := util.ParseUint(b, 10, 32)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
*v = uint32(n)
|
||||
return nil
|
||||
case *uint64:
|
||||
n, err := util.ParseUint(b, 10, 64)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
*v = n
|
||||
return nil
|
||||
case *float32:
|
||||
n, err := util.ParseFloat(b, 32)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
*v = float32(n)
|
||||
return err
|
||||
case *float64:
|
||||
var err error
|
||||
*v, err = util.ParseFloat(b, 64)
|
||||
return err
|
||||
case *bool:
|
||||
*v = len(b) == 1 && b[0] == '1'
|
||||
return nil
|
||||
case *time.Time:
|
||||
var err error
|
||||
*v, err = time.Parse(time.RFC3339Nano, util.BytesToString(b))
|
||||
return err
|
||||
case *time.Duration:
|
||||
n, err := util.ParseInt(b, 10, 64)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
*v = time.Duration(n)
|
||||
return nil
|
||||
case encoding.BinaryUnmarshaler:
|
||||
return v.UnmarshalBinary(b)
|
||||
case *net.IP:
|
||||
*v = b
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf(
|
||||
"redis: can't unmarshal %T (consider implementing BinaryUnmarshaler)", v)
|
||||
}
|
||||
}
|
||||
|
||||
func ScanSlice(data []string, slice interface{}) error {
|
||||
v := reflect.ValueOf(slice)
|
||||
if !v.IsValid() {
|
||||
return fmt.Errorf("redis: ScanSlice(nil)")
|
||||
}
|
||||
if v.Kind() != reflect.Ptr {
|
||||
return fmt.Errorf("redis: ScanSlice(non-pointer %T)", slice)
|
||||
}
|
||||
v = v.Elem()
|
||||
if v.Kind() != reflect.Slice {
|
||||
return fmt.Errorf("redis: ScanSlice(non-slice %T)", slice)
|
||||
}
|
||||
|
||||
next := makeSliceNextElemFunc(v)
|
||||
for i, s := range data {
|
||||
elem := next()
|
||||
if err := Scan([]byte(s), elem.Addr().Interface()); err != nil {
|
||||
err = fmt.Errorf("redis: ScanSlice index=%d value=%q failed: %w", i, s, err)
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func makeSliceNextElemFunc(v reflect.Value) func() reflect.Value {
|
||||
elemType := v.Type().Elem()
|
||||
|
||||
if elemType.Kind() == reflect.Ptr {
|
||||
elemType = elemType.Elem()
|
||||
return func() reflect.Value {
|
||||
if v.Len() < v.Cap() {
|
||||
v.Set(v.Slice(0, v.Len()+1))
|
||||
elem := v.Index(v.Len() - 1)
|
||||
if elem.IsNil() {
|
||||
elem.Set(reflect.New(elemType))
|
||||
}
|
||||
return elem.Elem()
|
||||
}
|
||||
|
||||
elem := reflect.New(elemType)
|
||||
v.Set(reflect.Append(v, elem))
|
||||
return elem.Elem()
|
||||
}
|
||||
}
|
||||
|
||||
zero := reflect.Zero(elemType)
|
||||
return func() reflect.Value {
|
||||
if v.Len() < v.Cap() {
|
||||
v.Set(v.Slice(0, v.Len()+1))
|
||||
return v.Index(v.Len() - 1)
|
||||
}
|
||||
|
||||
v.Set(reflect.Append(v, zero))
|
||||
return v.Index(v.Len() - 1)
|
||||
}
|
||||
}
|
||||
+242
@@ -0,0 +1,242 @@
|
||||
package proto
|
||||
|
||||
import (
|
||||
"encoding"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9/internal/util"
|
||||
)
|
||||
|
||||
type writer interface {
|
||||
io.Writer
|
||||
io.ByteWriter
|
||||
// WriteString implement io.StringWriter.
|
||||
WriteString(s string) (n int, err error)
|
||||
}
|
||||
|
||||
type Writer struct {
|
||||
writer
|
||||
|
||||
lenBuf []byte
|
||||
numBuf []byte
|
||||
}
|
||||
|
||||
func NewWriter(wr writer) *Writer {
|
||||
return &Writer{
|
||||
writer: wr,
|
||||
|
||||
lenBuf: make([]byte, 64),
|
||||
numBuf: make([]byte, 64),
|
||||
}
|
||||
}
|
||||
|
||||
func (w *Writer) WriteArgs(args []interface{}) error {
|
||||
if err := w.WriteByte(RespArray); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if err := w.writeLen(len(args)); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for _, arg := range args {
|
||||
if err := w.WriteArg(arg); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *Writer) writeLen(n int) error {
|
||||
w.lenBuf = strconv.AppendUint(w.lenBuf[:0], uint64(n), 10)
|
||||
w.lenBuf = append(w.lenBuf, '\r', '\n')
|
||||
_, err := w.Write(w.lenBuf)
|
||||
return err
|
||||
}
|
||||
|
||||
func (w *Writer) WriteArg(v interface{}) error {
|
||||
switch v := v.(type) {
|
||||
case nil:
|
||||
return w.string("")
|
||||
case string:
|
||||
return w.string(v)
|
||||
case *string:
|
||||
if v == nil {
|
||||
return w.string("")
|
||||
}
|
||||
return w.string(*v)
|
||||
case []byte:
|
||||
return w.bytes(v)
|
||||
case int:
|
||||
return w.int(int64(v))
|
||||
case *int:
|
||||
if v == nil {
|
||||
return w.int(0)
|
||||
}
|
||||
return w.int(int64(*v))
|
||||
case int8:
|
||||
return w.int(int64(v))
|
||||
case *int8:
|
||||
if v == nil {
|
||||
return w.int(0)
|
||||
}
|
||||
return w.int(int64(*v))
|
||||
case int16:
|
||||
return w.int(int64(v))
|
||||
case *int16:
|
||||
if v == nil {
|
||||
return w.int(0)
|
||||
}
|
||||
return w.int(int64(*v))
|
||||
case int32:
|
||||
return w.int(int64(v))
|
||||
case *int32:
|
||||
if v == nil {
|
||||
return w.int(0)
|
||||
}
|
||||
return w.int(int64(*v))
|
||||
case int64:
|
||||
return w.int(v)
|
||||
case *int64:
|
||||
if v == nil {
|
||||
return w.int(0)
|
||||
}
|
||||
return w.int(*v)
|
||||
case uint:
|
||||
return w.uint(uint64(v))
|
||||
case *uint:
|
||||
if v == nil {
|
||||
return w.uint(0)
|
||||
}
|
||||
return w.uint(uint64(*v))
|
||||
case uint8:
|
||||
return w.uint(uint64(v))
|
||||
case *uint8:
|
||||
if v == nil {
|
||||
return w.string("")
|
||||
}
|
||||
return w.uint(uint64(*v))
|
||||
case uint16:
|
||||
return w.uint(uint64(v))
|
||||
case *uint16:
|
||||
if v == nil {
|
||||
return w.uint(0)
|
||||
}
|
||||
return w.uint(uint64(*v))
|
||||
case uint32:
|
||||
return w.uint(uint64(v))
|
||||
case *uint32:
|
||||
if v == nil {
|
||||
return w.uint(0)
|
||||
}
|
||||
return w.uint(uint64(*v))
|
||||
case uint64:
|
||||
return w.uint(v)
|
||||
case *uint64:
|
||||
if v == nil {
|
||||
return w.uint(0)
|
||||
}
|
||||
return w.uint(*v)
|
||||
case float32:
|
||||
return w.float(float64(v))
|
||||
case *float32:
|
||||
if v == nil {
|
||||
return w.float(0)
|
||||
}
|
||||
return w.float(float64(*v))
|
||||
case float64:
|
||||
return w.float(v)
|
||||
case *float64:
|
||||
if v == nil {
|
||||
return w.float(0)
|
||||
}
|
||||
return w.float(*v)
|
||||
case bool:
|
||||
if v {
|
||||
return w.int(1)
|
||||
}
|
||||
return w.int(0)
|
||||
case *bool:
|
||||
if v == nil {
|
||||
return w.int(0)
|
||||
}
|
||||
if *v {
|
||||
return w.int(1)
|
||||
}
|
||||
return w.int(0)
|
||||
case time.Time:
|
||||
w.numBuf = v.AppendFormat(w.numBuf[:0], time.RFC3339Nano)
|
||||
return w.bytes(w.numBuf)
|
||||
case *time.Time:
|
||||
if v == nil {
|
||||
v = &time.Time{}
|
||||
}
|
||||
w.numBuf = v.AppendFormat(w.numBuf[:0], time.RFC3339Nano)
|
||||
return w.bytes(w.numBuf)
|
||||
case time.Duration:
|
||||
return w.int(v.Nanoseconds())
|
||||
case *time.Duration:
|
||||
if v == nil {
|
||||
return w.int(0)
|
||||
}
|
||||
return w.int(v.Nanoseconds())
|
||||
case encoding.BinaryMarshaler:
|
||||
b, err := v.MarshalBinary()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return w.bytes(b)
|
||||
case net.IP:
|
||||
return w.bytes(v)
|
||||
default:
|
||||
return fmt.Errorf(
|
||||
"redis: can't marshal %T (implement encoding.BinaryMarshaler)", v)
|
||||
}
|
||||
}
|
||||
|
||||
func (w *Writer) bytes(b []byte) error {
|
||||
if err := w.WriteByte(RespString); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if err := w.writeLen(len(b)); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if _, err := w.Write(b); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return w.crlf()
|
||||
}
|
||||
|
||||
func (w *Writer) string(s string) error {
|
||||
return w.bytes(util.StringToBytes(s))
|
||||
}
|
||||
|
||||
func (w *Writer) uint(n uint64) error {
|
||||
w.numBuf = strconv.AppendUint(w.numBuf[:0], n, 10)
|
||||
return w.bytes(w.numBuf)
|
||||
}
|
||||
|
||||
func (w *Writer) int(n int64) error {
|
||||
w.numBuf = strconv.AppendInt(w.numBuf[:0], n, 10)
|
||||
return w.bytes(w.numBuf)
|
||||
}
|
||||
|
||||
func (w *Writer) float(f float64) error {
|
||||
w.numBuf = strconv.AppendFloat(w.numBuf[:0], f, 'f', -1, 64)
|
||||
return w.bytes(w.numBuf)
|
||||
}
|
||||
|
||||
func (w *Writer) crlf() error {
|
||||
if err := w.WriteByte('\r'); err != nil {
|
||||
return err
|
||||
}
|
||||
return w.WriteByte('\n')
|
||||
}
|
||||
+50
@@ -0,0 +1,50 @@
|
||||
package rand
|
||||
|
||||
import (
|
||||
"math/rand"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// Int returns a non-negative pseudo-random int.
|
||||
func Int() int { return pseudo.Int() }
|
||||
|
||||
// Intn returns, as an int, a non-negative pseudo-random number in [0,n).
|
||||
// It panics if n <= 0.
|
||||
func Intn(n int) int { return pseudo.Intn(n) }
|
||||
|
||||
// Int63n returns, as an int64, a non-negative pseudo-random number in [0,n).
|
||||
// It panics if n <= 0.
|
||||
func Int63n(n int64) int64 { return pseudo.Int63n(n) }
|
||||
|
||||
// Perm returns, as a slice of n ints, a pseudo-random permutation of the integers [0,n).
|
||||
func Perm(n int) []int { return pseudo.Perm(n) }
|
||||
|
||||
// Seed uses the provided seed value to initialize the default Source to a
|
||||
// deterministic state. If Seed is not called, the generator behaves as if
|
||||
// seeded by Seed(1).
|
||||
func Seed(n int64) { pseudo.Seed(n) }
|
||||
|
||||
var pseudo = rand.New(&source{src: rand.NewSource(1)})
|
||||
|
||||
type source struct {
|
||||
src rand.Source
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func (s *source) Int63() int64 {
|
||||
s.mu.Lock()
|
||||
n := s.src.Int63()
|
||||
s.mu.Unlock()
|
||||
return n
|
||||
}
|
||||
|
||||
func (s *source) Seed(seed int64) {
|
||||
s.mu.Lock()
|
||||
s.src.Seed(seed)
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
// Shuffle pseudo-randomizes the order of elements.
|
||||
// n is the number of elements.
|
||||
// swap swaps the elements with indexes i and j.
|
||||
func Shuffle(n int, swap func(i, j int)) { pseudo.Shuffle(n, swap) }
|
||||
+3
@@ -0,0 +1,3 @@
|
||||
package internal
|
||||
|
||||
const RedisNull = "<nil>"
|
||||
+1000
File diff suppressed because it is too large
Load Diff
+144
@@ -0,0 +1,144 @@
|
||||
package routing
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type RequestPolicy uint8
|
||||
|
||||
const (
|
||||
ReqDefault RequestPolicy = iota
|
||||
|
||||
ReqAllNodes
|
||||
|
||||
ReqAllShards
|
||||
|
||||
ReqMultiShard
|
||||
|
||||
ReqSpecial
|
||||
)
|
||||
|
||||
const (
|
||||
ReadOnlyCMD string = "readonly"
|
||||
)
|
||||
|
||||
func (p RequestPolicy) String() string {
|
||||
switch p {
|
||||
case ReqDefault:
|
||||
return "default"
|
||||
case ReqAllNodes:
|
||||
return "all_nodes"
|
||||
case ReqAllShards:
|
||||
return "all_shards"
|
||||
case ReqMultiShard:
|
||||
return "multi_shard"
|
||||
case ReqSpecial:
|
||||
return "special"
|
||||
default:
|
||||
return fmt.Sprintf("unknown_request_policy(%d)", p)
|
||||
}
|
||||
}
|
||||
|
||||
func ParseRequestPolicy(raw string) (RequestPolicy, error) {
|
||||
switch strings.ToLower(raw) {
|
||||
case "", "default", "none":
|
||||
return ReqDefault, nil
|
||||
case "all_nodes":
|
||||
return ReqAllNodes, nil
|
||||
case "all_shards":
|
||||
return ReqAllShards, nil
|
||||
case "multi_shard":
|
||||
return ReqMultiShard, nil
|
||||
case "special":
|
||||
return ReqSpecial, nil
|
||||
default:
|
||||
return ReqDefault, fmt.Errorf("routing: unknown request_policy %q", raw)
|
||||
}
|
||||
}
|
||||
|
||||
type ResponsePolicy uint8
|
||||
|
||||
const (
|
||||
RespDefaultKeyless ResponsePolicy = iota
|
||||
RespDefaultHashSlot
|
||||
RespAllSucceeded
|
||||
RespOneSucceeded
|
||||
RespAggSum
|
||||
RespAggMin
|
||||
RespAggMax
|
||||
RespAggLogicalAnd
|
||||
RespAggLogicalOr
|
||||
RespSpecial
|
||||
)
|
||||
|
||||
func (p ResponsePolicy) String() string {
|
||||
switch p {
|
||||
case RespDefaultKeyless:
|
||||
return "default(keyless)"
|
||||
case RespDefaultHashSlot:
|
||||
return "default(hashslot)"
|
||||
case RespAllSucceeded:
|
||||
return "all_succeeded"
|
||||
case RespOneSucceeded:
|
||||
return "one_succeeded"
|
||||
case RespAggSum:
|
||||
return "agg_sum"
|
||||
case RespAggMin:
|
||||
return "agg_min"
|
||||
case RespAggMax:
|
||||
return "agg_max"
|
||||
case RespAggLogicalAnd:
|
||||
return "agg_logical_and"
|
||||
case RespAggLogicalOr:
|
||||
return "agg_logical_or"
|
||||
case RespSpecial:
|
||||
return "special"
|
||||
default:
|
||||
return "all_succeeded"
|
||||
}
|
||||
}
|
||||
|
||||
func ParseResponsePolicy(raw string) (ResponsePolicy, error) {
|
||||
switch strings.ToLower(raw) {
|
||||
case "default(keyless)":
|
||||
return RespDefaultKeyless, nil
|
||||
case "default(hashslot)":
|
||||
return RespDefaultHashSlot, nil
|
||||
case "all_succeeded":
|
||||
return RespAllSucceeded, nil
|
||||
case "one_succeeded":
|
||||
return RespOneSucceeded, nil
|
||||
case "agg_sum":
|
||||
return RespAggSum, nil
|
||||
case "agg_min":
|
||||
return RespAggMin, nil
|
||||
case "agg_max":
|
||||
return RespAggMax, nil
|
||||
case "agg_logical_and":
|
||||
return RespAggLogicalAnd, nil
|
||||
case "agg_logical_or":
|
||||
return RespAggLogicalOr, nil
|
||||
case "special":
|
||||
return RespSpecial, nil
|
||||
default:
|
||||
return RespDefaultKeyless, fmt.Errorf("routing: unknown response_policy %q", raw)
|
||||
}
|
||||
}
|
||||
|
||||
type CommandPolicy struct {
|
||||
Request RequestPolicy
|
||||
Response ResponsePolicy
|
||||
// Tips that are not request_policy or response_policy
|
||||
// e.g nondeterministic_output, nondeterministic_output_order.
|
||||
Tips map[string]string
|
||||
}
|
||||
|
||||
func (p *CommandPolicy) CanBeUsedInPipeline() bool {
|
||||
return p.Request != ReqAllNodes && p.Request != ReqAllShards && p.Request != ReqMultiShard
|
||||
}
|
||||
|
||||
func (p *CommandPolicy) IsReadOnly() bool {
|
||||
_, readOnly := p.Tips[ReadOnlyCMD]
|
||||
return readOnly
|
||||
}
|
||||
+57
@@ -0,0 +1,57 @@
|
||||
package routing
|
||||
|
||||
import (
|
||||
"math/rand"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
// ShardPicker chooses “one arbitrary shard” when the request_policy is
|
||||
// ReqDefault and the command has no keys.
|
||||
type ShardPicker interface {
|
||||
Next(total int) int // returns an index in [0,total)
|
||||
}
|
||||
|
||||
// StaticShardPicker always returns the same shard index.
|
||||
type StaticShardPicker struct {
|
||||
index int
|
||||
}
|
||||
|
||||
func NewStaticShardPicker(index int) *StaticShardPicker {
|
||||
return &StaticShardPicker{index: index}
|
||||
}
|
||||
|
||||
func (p *StaticShardPicker) Next(total int) int {
|
||||
if total == 0 || p.index >= total {
|
||||
return 0
|
||||
}
|
||||
return p.index
|
||||
}
|
||||
|
||||
/*───────────────────────────────
|
||||
Round-robin (default)
|
||||
────────────────────────────────*/
|
||||
|
||||
type RoundRobinPicker struct {
|
||||
cnt atomic.Uint32
|
||||
}
|
||||
|
||||
func (p *RoundRobinPicker) Next(total int) int {
|
||||
if total == 0 {
|
||||
return 0
|
||||
}
|
||||
i := p.cnt.Add(1)
|
||||
return int(i-1) % total
|
||||
}
|
||||
|
||||
/*───────────────────────────────
|
||||
Random
|
||||
────────────────────────────────*/
|
||||
|
||||
type RandomPicker struct{}
|
||||
|
||||
func (RandomPicker) Next(total int) int {
|
||||
if total == 0 {
|
||||
return 0
|
||||
}
|
||||
return rand.Intn(total)
|
||||
}
|
||||
+193
@@ -0,0 +1,193 @@
|
||||
package internal
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
var semTimers = sync.Pool{
|
||||
New: func() interface{} {
|
||||
t := time.NewTimer(time.Hour)
|
||||
t.Stop()
|
||||
return t
|
||||
},
|
||||
}
|
||||
|
||||
// FastSemaphore is a channel-based semaphore optimized for performance.
|
||||
// It uses a fast path that avoids timer allocation when tokens are available.
|
||||
// The channel is pre-filled with tokens: Acquire = receive, Release = send.
|
||||
// Closing the semaphore unblocks all waiting goroutines.
|
||||
//
|
||||
// Performance: ~30 ns/op with zero allocations on fast path.
|
||||
// Fairness: Eventual fairness (no starvation) but not strict FIFO.
|
||||
type FastSemaphore struct {
|
||||
tokens chan struct{}
|
||||
max int32
|
||||
}
|
||||
|
||||
// NewFastSemaphore creates a new fast semaphore with the given capacity.
|
||||
func NewFastSemaphore(capacity int32) *FastSemaphore {
|
||||
ch := make(chan struct{}, capacity)
|
||||
// Pre-fill with tokens
|
||||
for i := int32(0); i < capacity; i++ {
|
||||
ch <- struct{}{}
|
||||
}
|
||||
return &FastSemaphore{
|
||||
tokens: ch,
|
||||
max: capacity,
|
||||
}
|
||||
}
|
||||
|
||||
// TryAcquire attempts to acquire a token without blocking.
|
||||
// Returns true if successful, false if no tokens available.
|
||||
func (s *FastSemaphore) TryAcquire() bool {
|
||||
select {
|
||||
case <-s.tokens:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// Acquire acquires a token, blocking if necessary until one is available.
|
||||
// Returns an error if the context is cancelled or the timeout expires.
|
||||
// Uses a fast path to avoid timer allocation when tokens are immediately available.
|
||||
func (s *FastSemaphore) Acquire(ctx context.Context, timeout time.Duration, timeoutErr error) error {
|
||||
// Check context first
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
default:
|
||||
}
|
||||
|
||||
// Try fast path first (no timer needed)
|
||||
select {
|
||||
case <-s.tokens:
|
||||
return nil
|
||||
default:
|
||||
}
|
||||
|
||||
// Slow path: need to wait with timeout
|
||||
timer := semTimers.Get().(*time.Timer)
|
||||
defer semTimers.Put(timer)
|
||||
timer.Reset(timeout)
|
||||
|
||||
select {
|
||||
case <-s.tokens:
|
||||
if !timer.Stop() {
|
||||
<-timer.C
|
||||
}
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
if !timer.Stop() {
|
||||
<-timer.C
|
||||
}
|
||||
return ctx.Err()
|
||||
case <-timer.C:
|
||||
return timeoutErr
|
||||
}
|
||||
}
|
||||
|
||||
// AcquireBlocking acquires a token, blocking indefinitely until one is available.
|
||||
func (s *FastSemaphore) AcquireBlocking() {
|
||||
<-s.tokens
|
||||
}
|
||||
|
||||
// Release releases a token back to the semaphore.
|
||||
func (s *FastSemaphore) Release() {
|
||||
s.tokens <- struct{}{}
|
||||
}
|
||||
|
||||
// Close closes the semaphore, unblocking all waiting goroutines.
|
||||
// After close, all Acquire calls will receive a closed channel signal.
|
||||
func (s *FastSemaphore) Close() {
|
||||
close(s.tokens)
|
||||
}
|
||||
|
||||
// Len returns the current number of acquired tokens.
|
||||
func (s *FastSemaphore) Len() int32 {
|
||||
return s.max - int32(len(s.tokens))
|
||||
}
|
||||
|
||||
// FIFOSemaphore is a channel-based semaphore with strict FIFO ordering.
|
||||
// Unlike FastSemaphore, this guarantees that threads are served in the exact order they call Acquire().
|
||||
// The channel is pre-filled with tokens: Acquire = receive, Release = send.
|
||||
// Closing the semaphore unblocks all waiting goroutines.
|
||||
//
|
||||
// Performance: ~115 ns/op with zero allocations (slower than FastSemaphore due to timer allocation).
|
||||
// Fairness: Strict FIFO ordering guaranteed by Go runtime.
|
||||
type FIFOSemaphore struct {
|
||||
tokens chan struct{}
|
||||
max int32
|
||||
}
|
||||
|
||||
// NewFIFOSemaphore creates a new FIFO semaphore with the given capacity.
|
||||
func NewFIFOSemaphore(capacity int32) *FIFOSemaphore {
|
||||
ch := make(chan struct{}, capacity)
|
||||
// Pre-fill with tokens
|
||||
for i := int32(0); i < capacity; i++ {
|
||||
ch <- struct{}{}
|
||||
}
|
||||
return &FIFOSemaphore{
|
||||
tokens: ch,
|
||||
max: capacity,
|
||||
}
|
||||
}
|
||||
|
||||
// TryAcquire attempts to acquire a token without blocking.
|
||||
// Returns true if successful, false if no tokens available.
|
||||
func (s *FIFOSemaphore) TryAcquire() bool {
|
||||
select {
|
||||
case <-s.tokens:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// Acquire acquires a token, blocking if necessary until one is available.
|
||||
// Returns an error if the context is cancelled or the timeout expires.
|
||||
// Always uses timer to guarantee FIFO ordering (no fast path).
|
||||
func (s *FIFOSemaphore) Acquire(ctx context.Context, timeout time.Duration, timeoutErr error) error {
|
||||
// No fast path - always use timer to guarantee FIFO
|
||||
timer := semTimers.Get().(*time.Timer)
|
||||
defer semTimers.Put(timer)
|
||||
timer.Reset(timeout)
|
||||
|
||||
select {
|
||||
case <-s.tokens:
|
||||
if !timer.Stop() {
|
||||
<-timer.C
|
||||
}
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
if !timer.Stop() {
|
||||
<-timer.C
|
||||
}
|
||||
return ctx.Err()
|
||||
case <-timer.C:
|
||||
return timeoutErr
|
||||
}
|
||||
}
|
||||
|
||||
// AcquireBlocking acquires a token, blocking indefinitely until one is available.
|
||||
func (s *FIFOSemaphore) AcquireBlocking() {
|
||||
<-s.tokens
|
||||
}
|
||||
|
||||
// Release releases a token back to the semaphore.
|
||||
func (s *FIFOSemaphore) Release() {
|
||||
s.tokens <- struct{}{}
|
||||
}
|
||||
|
||||
// Close closes the semaphore, unblocking all waiting goroutines.
|
||||
// After close, all Acquire calls will receive a closed channel signal.
|
||||
func (s *FIFOSemaphore) Close() {
|
||||
close(s.tokens)
|
||||
}
|
||||
|
||||
// Len returns the current number of acquired tokens.
|
||||
func (s *FIFOSemaphore) Len() int32 {
|
||||
return s.max - int32(len(s.tokens))
|
||||
}
|
||||
+113
@@ -0,0 +1,113 @@
|
||||
package internal
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9/internal/util"
|
||||
)
|
||||
|
||||
func Sleep(ctx context.Context, dur time.Duration) error {
|
||||
t := time.NewTimer(dur)
|
||||
defer t.Stop()
|
||||
|
||||
select {
|
||||
case <-t.C:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
func ToLower(s string) string {
|
||||
if isLower(s) {
|
||||
return s
|
||||
}
|
||||
|
||||
b := make([]byte, len(s))
|
||||
for i := range b {
|
||||
c := s[i]
|
||||
if c >= 'A' && c <= 'Z' {
|
||||
c += 'a' - 'A'
|
||||
}
|
||||
b[i] = c
|
||||
}
|
||||
return util.BytesToString(b)
|
||||
}
|
||||
|
||||
func isLower(s string) bool {
|
||||
for i := 0; i < len(s); i++ {
|
||||
c := s[i]
|
||||
if c >= 'A' && c <= 'Z' {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func ReplaceSpaces(s string) string {
|
||||
return strings.ReplaceAll(s, " ", "-")
|
||||
}
|
||||
|
||||
func GetAddr(addr string) string {
|
||||
ind := strings.LastIndexByte(addr, ':')
|
||||
if ind == -1 {
|
||||
return ""
|
||||
}
|
||||
|
||||
if strings.IndexByte(addr, '.') != -1 {
|
||||
return addr
|
||||
}
|
||||
|
||||
if addr[0] == '[' {
|
||||
return addr
|
||||
}
|
||||
return net.JoinHostPort(addr[:ind], addr[ind+1:])
|
||||
}
|
||||
|
||||
func ToInteger(val interface{}) int {
|
||||
switch v := val.(type) {
|
||||
case int:
|
||||
return v
|
||||
case int64:
|
||||
return int(v)
|
||||
case string:
|
||||
i, _ := strconv.Atoi(v)
|
||||
return i
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
func ToFloat(val interface{}) float64 {
|
||||
switch v := val.(type) {
|
||||
case float64:
|
||||
return v
|
||||
case string:
|
||||
f, _ := strconv.ParseFloat(v, 64)
|
||||
return f
|
||||
default:
|
||||
return 0.0
|
||||
}
|
||||
}
|
||||
|
||||
func ToString(val interface{}) string {
|
||||
if str, ok := val.(string); ok {
|
||||
return str
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func ToStringSlice(val interface{}) []string {
|
||||
if arr, ok := val.([]interface{}); ok {
|
||||
result := make([]string, len(arr))
|
||||
for i, v := range arr {
|
||||
result[i] = ToString(v)
|
||||
}
|
||||
return result
|
||||
}
|
||||
return nil
|
||||
}
|
||||
+97
@@ -0,0 +1,97 @@
|
||||
/*
|
||||
© 2023–present Harald Rudell <harald.rudell@gmail.com> (https://haraldrudell.github.io/haraldrudell/)
|
||||
ISC License
|
||||
|
||||
Modified by htemelski-redis
|
||||
Removed the treshold, adapted it to work with float64
|
||||
*/
|
||||
|
||||
package util
|
||||
|
||||
import (
|
||||
"math"
|
||||
|
||||
"go.uber.org/atomic"
|
||||
)
|
||||
|
||||
// AtomicMax is a thread-safe max container
|
||||
// - hasValue indicator true if a value was equal to or greater than threshold
|
||||
// - optional threshold for minimum accepted max value
|
||||
// - if threshold is not used, initialization-free
|
||||
// - —
|
||||
// - wait-free CompareAndSwap mechanic
|
||||
type AtomicMax struct {
|
||||
|
||||
// value is current max
|
||||
value atomic.Float64
|
||||
// whether [AtomicMax.Value] has been invoked
|
||||
// with value equal or greater to threshold
|
||||
hasValue atomic.Bool
|
||||
}
|
||||
|
||||
// NewAtomicMax returns a thread-safe max container
|
||||
// - if threshold is not used, AtomicMax is initialization-free
|
||||
func NewAtomicMax() (atomicMax *AtomicMax) {
|
||||
m := AtomicMax{}
|
||||
m.value.Store((-math.MaxFloat64))
|
||||
return &m
|
||||
}
|
||||
|
||||
// Value updates the container with a possible max value
|
||||
// - isNewMax is true if:
|
||||
// - — value is equal to or greater than any threshold and
|
||||
// - — invocation recorded the first 0 or
|
||||
// - — a new max
|
||||
// - upon return, Max and Max1 are guaranteed to reflect the invocation
|
||||
// - the return order of concurrent Value invocations is not guaranteed
|
||||
// - Thread-safe
|
||||
func (m *AtomicMax) Value(value float64) (isNewMax bool) {
|
||||
// -math.MaxFloat64 as max case
|
||||
var hasValue0 = m.hasValue.Load()
|
||||
if value == (-math.MaxFloat64) {
|
||||
if !hasValue0 {
|
||||
isNewMax = m.hasValue.CompareAndSwap(false, true)
|
||||
}
|
||||
return // -math.MaxFloat64 as max: isNewMax true for first 0 writer
|
||||
}
|
||||
|
||||
// check against present value
|
||||
var current = m.value.Load()
|
||||
if isNewMax = value > current; !isNewMax {
|
||||
return // not a new max return: isNewMax false
|
||||
}
|
||||
|
||||
// store the new max
|
||||
for {
|
||||
|
||||
// try to write value to *max
|
||||
if isNewMax = m.value.CompareAndSwap(current, value); isNewMax {
|
||||
if !hasValue0 {
|
||||
// may be rarely written multiple times
|
||||
// still faster than CompareAndSwap
|
||||
m.hasValue.Store(true)
|
||||
}
|
||||
return // new max written return: isNewMax true
|
||||
}
|
||||
if current = m.value.Load(); current >= value {
|
||||
return // no longer a need to write return: isNewMax false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Max returns current max and value-present flag
|
||||
// - hasValue true indicates that value reflects a Value invocation
|
||||
// - hasValue false: value is zero-value
|
||||
// - Thread-safe
|
||||
func (m *AtomicMax) Max() (value float64, hasValue bool) {
|
||||
if hasValue = m.hasValue.Load(); !hasValue {
|
||||
return
|
||||
}
|
||||
value = m.value.Load()
|
||||
return
|
||||
}
|
||||
|
||||
// Max1 returns current maximum whether zero-value or set by Value
|
||||
// - threshold is ignored
|
||||
// - Thread-safe
|
||||
func (m *AtomicMax) Max1() (value float64) { return m.value.Load() }
|
||||
+96
@@ -0,0 +1,96 @@
|
||||
package util
|
||||
|
||||
/*
|
||||
© 2023–present Harald Rudell <harald.rudell@gmail.com> (https://haraldrudell.github.io/haraldrudell/)
|
||||
ISC License
|
||||
|
||||
Modified by htemelski-redis
|
||||
Adapted from the modified atomic_max, but with inverted logic
|
||||
*/
|
||||
|
||||
import (
|
||||
"math"
|
||||
|
||||
"go.uber.org/atomic"
|
||||
)
|
||||
|
||||
// AtomicMin is a thread-safe Min container
|
||||
// - hasValue indicator true if a value was equal to or greater than threshold
|
||||
// - optional threshold for minimum accepted Min value
|
||||
// - —
|
||||
// - wait-free CompareAndSwap mechanic
|
||||
type AtomicMin struct {
|
||||
|
||||
// value is current Min
|
||||
value atomic.Float64
|
||||
// whether [AtomicMin.Value] has been invoked
|
||||
// with value equal or greater to threshold
|
||||
hasValue atomic.Bool
|
||||
}
|
||||
|
||||
// NewAtomicMin returns a thread-safe Min container
|
||||
// - if threshold is not used, AtomicMin is initialization-free
|
||||
func NewAtomicMin() (atomicMin *AtomicMin) {
|
||||
m := AtomicMin{}
|
||||
m.value.Store(math.MaxFloat64)
|
||||
return &m
|
||||
}
|
||||
|
||||
// Value updates the container with a possible Min value
|
||||
// - isNewMin is true if:
|
||||
// - — value is equal to or greater than any threshold and
|
||||
// - — invocation recorded the first 0 or
|
||||
// - — a new Min
|
||||
// - upon return, Min and Min1 are guaranteed to reflect the invocation
|
||||
// - the return order of concurrent Value invocations is not guaranteed
|
||||
// - Thread-safe
|
||||
func (m *AtomicMin) Value(value float64) (isNewMin bool) {
|
||||
// math.MaxFloat64 as Min case
|
||||
var hasValue0 = m.hasValue.Load()
|
||||
if value == math.MaxFloat64 {
|
||||
if !hasValue0 {
|
||||
isNewMin = m.hasValue.CompareAndSwap(false, true)
|
||||
}
|
||||
return // math.MaxFloat64 as Min: isNewMin true for first 0 writer
|
||||
}
|
||||
|
||||
// check against present value
|
||||
var current = m.value.Load()
|
||||
if isNewMin = value < current; !isNewMin {
|
||||
return // not a new Min return: isNewMin false
|
||||
}
|
||||
|
||||
// store the new Min
|
||||
for {
|
||||
|
||||
// try to write value to *Min
|
||||
if isNewMin = m.value.CompareAndSwap(current, value); isNewMin {
|
||||
if !hasValue0 {
|
||||
// may be rarely written multiple times
|
||||
// still faster than CompareAndSwap
|
||||
m.hasValue.Store(true)
|
||||
}
|
||||
return // new Min written return: isNewMin true
|
||||
}
|
||||
if current = m.value.Load(); current <= value {
|
||||
return // no longer a need to write return: isNewMin false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Min returns current min and value-present flag
|
||||
// - hasValue true indicates that value reflects a Value invocation
|
||||
// - hasValue false: value is zero-value
|
||||
// - Thread-safe
|
||||
func (m *AtomicMin) Min() (value float64, hasValue bool) {
|
||||
if hasValue = m.hasValue.Load(); !hasValue {
|
||||
return
|
||||
}
|
||||
value = m.value.Load()
|
||||
return
|
||||
}
|
||||
|
||||
// Min1 returns current Minimum whether zero-value or set by Value
|
||||
// - threshold is ignored
|
||||
// - Thread-safe
|
||||
func (m *AtomicMin) Min1() (value float64) { return m.value.Load() }
|
||||
+41
@@ -0,0 +1,41 @@
|
||||
package util
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
// ParseFloat parses a Redis RESP3 float reply into a Go float64,
|
||||
// handling "inf", "-inf", "nan" per Redis conventions.
|
||||
func ParseStringToFloat(s string) (float64, error) {
|
||||
switch s {
|
||||
case "inf":
|
||||
return math.Inf(1), nil
|
||||
case "-inf":
|
||||
return math.Inf(-1), nil
|
||||
case "nan", "-nan":
|
||||
return math.NaN(), nil
|
||||
}
|
||||
return strconv.ParseFloat(s, 64)
|
||||
}
|
||||
|
||||
// MustParseFloat is like ParseFloat but panics on parse errors.
|
||||
func MustParseFloat(s string) float64 {
|
||||
f, err := ParseStringToFloat(s)
|
||||
if err != nil {
|
||||
panic(fmt.Sprintf("redis: failed to parse float %q: %v", s, err))
|
||||
}
|
||||
return f
|
||||
}
|
||||
|
||||
// SafeIntToInt32 safely converts an int to int32, returning an error if overflow would occur.
|
||||
func SafeIntToInt32(value int, fieldName string) (int32, error) {
|
||||
if value > math.MaxInt32 {
|
||||
return 0, fmt.Errorf("redis: %s value %d exceeds maximum allowed value %d", fieldName, value, math.MaxInt32)
|
||||
}
|
||||
if value < math.MinInt32 {
|
||||
return 0, fmt.Errorf("redis: %s value %d is below minimum allowed value %d", fieldName, value, math.MinInt32)
|
||||
}
|
||||
return int32(value), nil
|
||||
}
|
||||
+11
@@ -0,0 +1,11 @@
|
||||
//go:build appengine
|
||||
|
||||
package util
|
||||
|
||||
func BytesToString(b []byte) string {
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func StringToBytes(s string) []byte {
|
||||
return []byte(s)
|
||||
}
|
||||
+19
@@ -0,0 +1,19 @@
|
||||
package util
|
||||
|
||||
import "strconv"
|
||||
|
||||
func Atoi(b []byte) (int, error) {
|
||||
return strconv.Atoi(BytesToString(b))
|
||||
}
|
||||
|
||||
func ParseInt(b []byte, base int, bitSize int) (int64, error) {
|
||||
return strconv.ParseInt(BytesToString(b), base, bitSize)
|
||||
}
|
||||
|
||||
func ParseUint(b []byte, base int, bitSize int) (uint64, error) {
|
||||
return strconv.ParseUint(BytesToString(b), base, bitSize)
|
||||
}
|
||||
|
||||
func ParseFloat(b []byte, bitSize int) (float64, error) {
|
||||
return strconv.ParseFloat(BytesToString(b), bitSize)
|
||||
}
|
||||
+5
@@ -0,0 +1,5 @@
|
||||
package util
|
||||
|
||||
func ToPtr[T any](v T) *T {
|
||||
return &v
|
||||
}
|
||||
+17
@@ -0,0 +1,17 @@
|
||||
//go:build !appengine
|
||||
|
||||
package util
|
||||
|
||||
import (
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
// BytesToString converts byte slice to string.
|
||||
func BytesToString(b []byte) string {
|
||||
return unsafe.String(unsafe.SliceData(b), len(b))
|
||||
}
|
||||
|
||||
// StringToBytes converts string to byte slice.
|
||||
func StringToBytes(s string) []byte {
|
||||
return unsafe.Slice(unsafe.StringData(s), len(s))
|
||||
}
|
||||
Reference in New Issue
Block a user