wip.

internal/hitless/client_integration.go

@@ -0,0 +1,200 @@
+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
+	}
+}