iop/apps/node/internal/bootstrap/runtime_supervisor.go

package bootstrap

import (
	"context"
	"fmt"
	"sync"
	"time"

	"go.uber.org/fx"
	"go.uber.org/zap"

	"iop/apps/node/internal/transport"
	"iop/packages/go/config"
	"iop/packages/go/observability"
)

// runtimeSupervisor owns the node's single active edge connection and the
// reconnect loop. It is created once per fx lifecycle and drives
// connect/reconnect/shutdown. The mutex guards current across the reconnect
// goroutine and the OnStop hook.
type runtimeSupervisor struct {
	cfg        *config.NodeConfig
	logger     *zap.Logger
	dialer     DialFunc
	sleeper    func(ctx context.Context, d time.Duration)
	shutdowner fx.Shutdowner

	mu               sync.Mutex
	current          *runtimeOwner
	cancelSupervisor context.CancelFunc
}

// start establishes the initial connection, launches the reconnect supervisor
// goroutine, and starts the metrics server. It maps a connect failure to the
// same "bootstrap: %w" error the OnStart hook returned before extraction.
func (s *runtimeSupervisor) start(ctx context.Context) error {
	owner, err := connectRuntime(ctx, s.cfg, s.logger, s.dialer)
	if err != nil {
		return fmt.Errorf("bootstrap: %w", err)
	}

	s.mu.Lock()
	s.current = owner
	s.mu.Unlock()

	supCtx, cancel := context.WithCancel(context.Background())
	s.cancelSupervisor = cancel

	go s.run(supCtx, owner.sess)

	go func() {
		if err := observability.ServeMetrics(s.cfg.Metrics.Port); err != nil {
			s.logger.Warn("metrics server exited", zap.Error(err))
		}
	}()
	return nil
}

// run waits for the active session to end, then drives a reconnect. It exits on
// supervisor cancellation, a local shutdown, or when reconnect returns nil
// (cancelled or exhausted).
func (s *runtimeSupervisor) run(supCtx context.Context, sess *transport.Session) {
	for {
		select {
		case <-supCtx.Done():
			return
		case <-sess.Done():
		}

		// Guard against context cancellation racing with disconnect.
		select {
		case <-supCtx.Done():
			return
		default:
		}

		if sess.IsLocalShutdown() {
			return
		}

		newOwner := s.reconnect(supCtx)
		if newOwner == nil {
			return
		}
		sess = newOwner.sess
	}
}

// reconnect runs the bounded reconnect loop. On success it swaps in the new
// owner and returns it. On supervisor cancellation it returns nil without
// changing state. On exhaustion it closes the current owner, requests node
// shutdown, and returns nil. The attempt/interval defaults, sleep, dial timeout,
// and logging match the prior inline loop.
func (s *runtimeSupervisor) reconnect(supCtx context.Context) *runtimeOwner {
	maxAttempts := s.cfg.Reconnect.MaxAttempts
	if maxAttempts <= 0 {
		maxAttempts = 10
	}
	intervalSec := s.cfg.Reconnect.IntervalSec
	if intervalSec < 0 {
		intervalSec = 10
	}

	for attempt := 1; attempt <= maxAttempts; attempt++ {
		select {
		case <-supCtx.Done():
			return nil
		default:
		}

		s.sleeper(supCtx, time.Duration(intervalSec)*time.Second)
		if supCtx.Err() != nil {
			return nil
		}

		s.logger.Info("reconnecting to edge",
			zap.Int("attempt", attempt),
			zap.Int("max_attempts", maxAttempts),
			zap.Int("interval_sec", intervalSec),
		)

		dialCtx, dialCancel := context.WithTimeout(supCtx, 30*time.Second)
		o, dialErr := connectRuntime(dialCtx, s.cfg, s.logger, s.dialer)
		dialCancel()
		if dialErr == nil {
			s.swapOwner(o)
			return o
		}

		s.logger.Warn("reconnect attempt failed",
			zap.Int("attempt", attempt),
			zap.Int("max_attempts", maxAttempts),
			zap.Error(dialErr),
		)
	}

	s.exhaust(maxAttempts)
	return nil
}

// swapOwner installs owner as the current connection and closes the previous one.
func (s *runtimeSupervisor) swapOwner(owner *runtimeOwner) {
	s.mu.Lock()
	prev := s.current
	s.current = owner
	s.mu.Unlock()
	if prev != nil {
		prev.close()
	}
}

// exhaust clears the current owner, logs the exhaustion, and requests a node
// shutdown with a non-zero exit code.
func (s *runtimeSupervisor) exhaust(maxAttempts int) {
	s.mu.Lock()
	prev := s.current
	s.current = nil
	s.mu.Unlock()
	if prev != nil {
		prev.close()
	}
	s.logger.Error("reconnect exhausted, shutting down node",
		zap.Int("max_attempts", maxAttempts),
	)
	_ = s.shutdowner.Shutdown(fx.ExitCode(1))
}

// stop cancels the supervisor goroutine and closes the current connection.
func (s *runtimeSupervisor) stop() {
	if s.cancelSupervisor != nil {
		s.cancelSupervisor()
	}
	s.mu.Lock()
	owner := s.current
	s.current = nil
	s.mu.Unlock()
	if owner != nil {
		owner.close()
	}
}