go-redis/internal/hitless/client_integration.go

package hitless

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

	"github.com/redis/go-redis/v9/internal/pool"
	"github.com/redis/go-redis/v9/push"
)

// ClientIntegrator provides integration between hitless upgrade handlers and Redis clients
type ClientIntegrator struct {
	upgradeHandler *UpgradeHandler
	mu             sync.RWMutex

	// Simple atomic state for pool redirection
	isMoving    int32  // atomic: 0 = not moving, 1 = moving
	newEndpoint string // only written during MOVING, read-only after
}

// NewClientIntegrator creates a new client integrator with client timeout configuration
func NewClientIntegrator(defaultReadTimeout, defaultWriteTimeout time.Duration) *ClientIntegrator {
	return &ClientIntegrator{
		upgradeHandler: NewUpgradeHandler(defaultReadTimeout, defaultWriteTimeout),
	}
}

// GetUpgradeHandler returns the upgrade handler for direct access
func (ci *ClientIntegrator) GetUpgradeHandler() *UpgradeHandler {
	return ci.upgradeHandler
}

// HandlePushNotification is the main entry point for processing upgrade notifications
func (ci *ClientIntegrator) HandlePushNotification(ctx context.Context, handlerCtx push.NotificationHandlerContext, notification []interface{}) error {
	// Handle MOVING notifications for pool redirection
	if len(notification) > 0 {
		if notificationType, ok := notification[0].(string); ok && notificationType == "MOVING" {
			if len(notification) >= 3 {
				if newEndpoint, ok := notification[2].(string); ok {
					// Simple atomic state update - no locks needed
					ci.newEndpoint = newEndpoint
					atomic.StoreInt32(&ci.isMoving, 1)
				}
			}
		}
	}

	return ci.upgradeHandler.HandlePushNotification(ctx, handlerCtx, notification)
}

// Close shuts down the client integrator
func (ci *ClientIntegrator) Close() error {
	ci.mu.Lock()
	defer ci.mu.Unlock()

	// Reset atomic state
	atomic.StoreInt32(&ci.isMoving, 0)
	ci.newEndpoint = ""

	return nil
}

// IsMoving returns true if the pool is currently moving to a new endpoint
// Uses atomic read - no locks needed
func (ci *ClientIntegrator) IsMoving() bool {
	return atomic.LoadInt32(&ci.isMoving) == 1
}

// GetNewEndpoint returns the new endpoint if moving, empty string otherwise
// Safe to read without locks since it's only written during MOVING
func (ci *ClientIntegrator) GetNewEndpoint() string {
	if ci.IsMoving() {
		return ci.newEndpoint
	}
	return ""
}

// PushNotificationHandlerInterface defines the interface for push notification handlers
// This implements the interface expected by the push notification system
type PushNotificationHandlerInterface interface {
	HandlePushNotification(ctx context.Context, handlerCtx push.NotificationHandlerContext, notification []interface{}) error
}

// Ensure ClientIntegrator implements the interface
var _ PushNotificationHandlerInterface = (*ClientIntegrator)(nil)

// PoolRedirector provides pool redirection functionality for hitless upgrades
type PoolRedirector struct {
	poolManager *PoolEndpointManager
	mu          sync.RWMutex
}

// NewPoolRedirector creates a new pool redirector
func NewPoolRedirector() *PoolRedirector {
	return &PoolRedirector{
		poolManager: NewPoolEndpointManager(),
	}
}

// RedirectPool redirects a connection pool to a new endpoint
func (pr *PoolRedirector) RedirectPool(ctx context.Context, pooler pool.Pooler, newEndpoint string, timeout time.Duration) error {
	pr.mu.Lock()
	defer pr.mu.Unlock()

	return pr.poolManager.RedirectPool(ctx, pooler, newEndpoint, timeout)
}

// IsPoolRedirected checks if a pool is currently redirected
func (pr *PoolRedirector) IsPoolRedirected(pooler pool.Pooler) bool {
	pr.mu.RLock()
	defer pr.mu.RUnlock()

	return pr.poolManager.IsPoolRedirected(pooler)
}

// GetRedirection returns redirection information for a pool
func (pr *PoolRedirector) GetRedirection(pooler pool.Pooler) (*EndpointRedirection, bool) {
	pr.mu.RLock()
	defer pr.mu.RUnlock()

	return pr.poolManager.GetRedirection(pooler)
}

// Close shuts down the pool redirector
func (pr *PoolRedirector) Close() error {
	pr.mu.Lock()
	defer pr.mu.Unlock()

	// Clean up all redirections
	ctx := context.Background()
	pr.poolManager.CleanupExpiredRedirections(ctx)

	return nil
}

// ConnectionStateTracker tracks connection states during upgrades
type ConnectionStateTracker struct {
	upgradeHandler *UpgradeHandler
	mu             sync.RWMutex
}

// NewConnectionStateTracker creates a new connection state tracker with timeout configuration
func NewConnectionStateTracker(defaultReadTimeout, defaultWriteTimeout time.Duration) *ConnectionStateTracker {
	return &ConnectionStateTracker{
		upgradeHandler: NewUpgradeHandler(defaultReadTimeout, defaultWriteTimeout),
	}
}

// IsConnectionTransitioning checks if a connection is currently transitioning
func (cst *ConnectionStateTracker) IsConnectionTransitioning(conn *pool.Conn) bool {
	cst.mu.RLock()
	defer cst.mu.RUnlock()

	return cst.upgradeHandler.IsConnectionTransitioning(conn)
}

// GetConnectionState returns the current state of a connection
func (cst *ConnectionStateTracker) GetConnectionState(conn *pool.Conn) (*ConnectionState, bool) {
	cst.mu.RLock()
	defer cst.mu.RUnlock()

	return cst.upgradeHandler.GetConnectionState(conn)
}

// CleanupConnection removes tracking for a connection
func (cst *ConnectionStateTracker) CleanupConnection(conn *pool.Conn) {
	cst.mu.Lock()
	defer cst.mu.Unlock()

	cst.upgradeHandler.CleanupConnection(conn)
}

// Close shuts down the connection state tracker
func (cst *ConnectionStateTracker) Close() error {
	cst.mu.Lock()
	defer cst.mu.Unlock()

	// Clean up all expired states
	cst.upgradeHandler.CleanupExpiredStates()

	return nil
}

// HitlessUpgradeConfig provides configuration for hitless upgrades
type HitlessUpgradeConfig struct {
	Enabled           bool
	TransitionTimeout time.Duration
	CleanupInterval   time.Duration
}

// DefaultHitlessUpgradeConfig returns default configuration for hitless upgrades
func DefaultHitlessUpgradeConfig() *HitlessUpgradeConfig {
	return &HitlessUpgradeConfig{
		Enabled:           true,
		TransitionTimeout: 60 * time.Second, // Longer timeout for transitioning connections
		CleanupInterval:   30 * time.Second, // How often to clean up expired states
	}
}