go-kratos/internal/websocket/reconnect.go

package websocket

import (
	"log"
	"math/rand"
	"sync"
	"time"
)

// ReconnectConfig holds configuration for reconnection attempts
type ReconnectConfig struct {
	MaxAttempts   int
	InitialDelay  time.Duration
	MaxDelay      time.Duration
	Multiplier    float64
	JitterFactor  float64
	TimeoutPerTry time.Duration
}

// DefaultReconnectConfig returns default reconnection configuration
func DefaultReconnectConfig() *ReconnectConfig {
	return &ReconnectConfig{
		MaxAttempts:   5,
		InitialDelay:  time.Second,
		MaxDelay:      time.Minute,
		Multiplier:    2.0,
		JitterFactor:  0.1,
		TimeoutPerTry: time.Second * 5,
	}
}

// ReconnectManager handles reconnection logic
type ReconnectManager struct {
	config    *ReconnectConfig
	attempts  int
	lastDelay time.Duration
	mu        sync.RWMutex
	active    bool
	stopChan  chan struct{}
	stats     *ReconnectStats
}

// ReconnectStats tracks reconnection statistics
type ReconnectStats struct {
	TotalAttempts     int
	SuccessfulRetries int
	FailedRetries     int
	LastAttemptTime   time.Time
	TotalDowntime     time.Duration
	mu                sync.RWMutex
}

func NewReconnectStats() *ReconnectStats {
	return &ReconnectStats{
		LastAttemptTime: time.Now(),
	}
}

func (rs *ReconnectStats) RecordAttempt(success bool) {
	rs.mu.Lock()
	defer rs.mu.Unlock()

	rs.TotalAttempts++
	if success {
		rs.SuccessfulRetries++
	} else {
		rs.FailedRetries++
	}
	rs.LastAttemptTime = time.Now()
}

func (rs *ReconnectStats) UpdateDowntime(duration time.Duration) {
	rs.mu.Lock()
	defer rs.mu.Unlock()
	rs.TotalDowntime += duration
}

func (rs *ReconnectStats) GetStats() map[string]interface{} {
	rs.mu.RLock()
	defer rs.mu.RUnlock()

	return map[string]interface{}{
		"total_attempts":     rs.TotalAttempts,
		"successful_retries": rs.SuccessfulRetries,
		"failed_retries":     rs.FailedRetries,
		"last_attempt_time":  rs.LastAttemptTime,
		"total_downtime":     rs.TotalDowntime.String(),
		"success_rate":       float64(rs.SuccessfulRetries) / float64(rs.TotalAttempts),
	}
}

func NewReconnectManager(config *ReconnectConfig) *ReconnectManager {
	if config == nil {
		config = DefaultReconnectConfig()
	}
	return &ReconnectManager{
		config:    config,
		stopChan:  make(chan struct{}),
		stats:     NewReconnectStats(),
		lastDelay: config.InitialDelay,
	}
}

func (rm *ReconnectManager) Start(connectFunc func() error) {
	rm.mu.Lock()
	if rm.active {
		rm.mu.Unlock()
		return
	}
	rm.active = true
	rm.mu.Unlock()

	go rm.reconnectLoop(connectFunc)
}

func (rm *ReconnectManager) Stop() {
	rm.mu.Lock()
	defer rm.mu.Unlock()

	if rm.active {
		close(rm.stopChan)
		rm.active = false
	}
}

func (rm *ReconnectManager) reconnectLoop(connectFunc func() error) {
	for {
		select {
		case <-rm.stopChan:
			return
		default:
			if rm.attempts >= rm.config.MaxAttempts {
				log.Printf("Max reconnection attempts (%d) reached", rm.config.MaxAttempts)
				rm.stats.RecordAttempt(false)
				return
			}

			rm.attempts++
			startTime := time.Now()

			err := connectFunc()
			if err == nil {
				log.Printf("Successfully reconnected after %d attempts", rm.attempts)
				rm.stats.RecordAttempt(true)
				rm.stats.UpdateDowntime(time.Since(startTime))
				rm.reset()
				return
			}

			rm.stats.RecordAttempt(false)
			rm.stats.UpdateDowntime(time.Since(startTime))
			log.Printf("Reconnection attempt %d failed: %v", rm.attempts, err)

			delay := rm.calculateDelay()
			time.Sleep(delay)
		}
	}
}

func (rm *ReconnectManager) calculateDelay() time.Duration {
	rm.mu.Lock()
	defer rm.mu.Unlock()

	delay := rm.lastDelay
	rm.lastDelay = time.Duration(float64(rm.lastDelay) * rm.config.Multiplier)
	if rm.lastDelay > rm.config.MaxDelay {
		rm.lastDelay = rm.config.MaxDelay
	}

	// Add jitter
	jitter := time.Duration(float64(rm.lastDelay) * rm.config.JitterFactor)
	delay += time.Duration(float64(jitter) * (2*rand.Float64() - 1))

	return delay
}

func (rm *ReconnectManager) reset() {
	rm.mu.Lock()
	defer rm.mu.Unlock()

	rm.attempts = 0
	rm.lastDelay = rm.config.InitialDelay
}

func (rm *ReconnectManager) GetStats() map[string]interface{} {
	return rm.stats.GetStats()
}