proxy-pool/internal/gateway/controlplane/session_supervisor.go

package controlplane

import (
	"context"
	"errors"
	"math/rand"
	"time"

	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

var ErrInvalidReconnectOptions = errors.New("invalid gateway reconnect options")

// SessionLifecycle is the long-running worker control-plane session boundary.
// SessionRunner satisfies this interface.
type SessionLifecycle interface {
	Run(context.Context) error
}

// ReconnectOptions controls bounded exponential reconnect delays. Jitter is a
// percentage applied symmetrically to each base delay.
type ReconnectOptions struct {
	InitialDelay time.Duration
	MaxDelay     time.Duration
	Jitter       int
}

// SessionSupervisorRuntime makes reconnect timing deterministic in tests.
// Nil fields use the production timer and random source.
type SessionSupervisorRuntime struct {
	Sleeper SessionSleeper
	Random  SessionRandom
}

type SessionSleeper interface {
	Sleep(context.Context, time.Duration) error
}

type SessionRandom interface {
	Float64() float64
}

// SessionSupervisor owns retry policy around repeated worker sessions without
// coupling it to SnapshotWatcher or RuntimeReporter.
type SessionSupervisor struct {
	lifecycle SessionLifecycle
	options   ReconnectOptions
	runtime   SessionSupervisorRuntime
}

func NewSessionSupervisor(lifecycle SessionLifecycle, options ReconnectOptions, runtimes ...SessionSupervisorRuntime) (*SessionSupervisor, error) {
	if lifecycle == nil || options.InitialDelay <= 0 || options.MaxDelay < options.InitialDelay || options.Jitter < 0 || options.Jitter > 100 || len(runtimes) > 1 {
		return nil, ErrInvalidReconnectOptions
	}
	runtime := SessionSupervisorRuntime{Sleeper: timerSessionSleeper{}, Random: globalSessionRandom{}}
	if len(runtimes) == 1 {
		if runtimes[0].Sleeper != nil {
			runtime.Sleeper = runtimes[0].Sleeper
		}
		if runtimes[0].Random != nil {
			runtime.Random = runtimes[0].Random
		}
	}
	return &SessionSupervisor{lifecycle: lifecycle, options: options, runtime: runtime}, nil
}

// Run restarts recoverable sessions until the context is canceled. A session
// that ends without an error is treated as a closed snapshot stream.
func (supervisor *SessionSupervisor) Run(ctx context.Context) error {
	if supervisor == nil || ctx == nil {
		return ErrInvalidReconnectOptions
	}
	for failedAttempts := uint(0); ; failedAttempts++ {
		err := supervisor.lifecycle.Run(ctx)
		if ctx.Err() != nil {
			return ctx.Err()
		}
		if err == nil {
			err = ErrSnapshotStreamClosed
		}
		if !retryableSessionError(err) {
			return err
		}
		if err := supervisor.runtime.Sleeper.Sleep(ctx, supervisor.retryDelay(failedAttempts)); err != nil {
			if ctx.Err() != nil {
				return ctx.Err()
			}
			return err
		}
	}
}

func (supervisor *SessionSupervisor) retryDelay(failedAttempts uint) time.Duration {
	delay := supervisor.options.InitialDelay
	for attempt := uint(0); attempt < failedAttempts && delay < supervisor.options.MaxDelay; attempt++ {
		if delay >= supervisor.options.MaxDelay/2 {
			delay = supervisor.options.MaxDelay
			break
		}
		delay *= 2
	}
	if supervisor.options.Jitter == 0 {
		return delay
	}
	random := supervisor.runtime.Random.Float64()
	if random < 0 {
		random = 0
	} else if random > 1 {
		random = 1
	}
	spread := float64(supervisor.options.Jitter) / 100
	delay = time.Duration(float64(delay) * (1 + (2*random-1)*spread))
	if delay < time.Millisecond {
		delay = time.Millisecond
	}
	if delay > supervisor.options.MaxDelay {
		return supervisor.options.MaxDelay
	}
	return delay
}

func retryableSessionError(err error) bool {
	if errors.Is(err, ErrInvalidOptions) {
		return false
	}
	switch status.Code(err) {
	case codes.InvalidArgument, codes.PermissionDenied, codes.Unauthenticated, codes.Unimplemented:
		return false
	default:
		return true
	}
}

type timerSessionSleeper struct{}

func (timerSessionSleeper) Sleep(ctx context.Context, delay time.Duration) error {
	timer := time.NewTimer(delay)
	defer timer.Stop()
	select {
	case <-timer.C:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

type globalSessionRandom struct{}

func (globalSessionRandom) Float64() float64 {
	return rand.Float64()
}

var _ SessionLifecycle = (*SessionRunner)(nil)