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

// Package bootstrap wires the IOP Node application using go.uber.org/fx.
package bootstrap

import (
	"context"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sync"
	"time"

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

	"iop/apps/node/internal/adapters"
	"iop/apps/node/internal/node"
	"iop/apps/node/internal/router"
	"iop/apps/node/internal/store"
	"iop/apps/node/internal/transport"
	"iop/packages/go/config"
	"iop/packages/go/events"
	"iop/packages/go/observability"
	iop "iop/proto/gen/iop"
)

// DialFunc is the signature used to connect to Edge. Can be replaced in tests.
type DialFunc func(ctx context.Context, addr, token string, logger *zap.Logger) (*transport.RegisterResult, error)

// Option configures Module behaviour; intended for testing only.
type Option func(*moduleOpts)

type moduleOpts struct {
	dialer  DialFunc
	sleeper func(ctx context.Context, d time.Duration)
}

// WithDialer replaces the transport dial function.
func WithDialer(fn DialFunc) Option {
	return func(o *moduleOpts) { o.dialer = fn }
}

// WithSleeper replaces the inter-retry sleep function.
func WithSleeper(fn func(ctx context.Context, d time.Duration)) Option {
	return func(o *moduleOpts) { o.sleeper = fn }
}

// runtimeOwner holds one connection's resources and closes them idempotently.
type runtimeOwner struct {
	reg  *adapters.Registry
	sess *transport.Session
	st   *store.Store
	once sync.Once
}

func (r *runtimeOwner) close() {
	r.once.Do(func() {
		if r.reg != nil {
			_ = r.reg.Stop(context.Background())
		}
		if r.sess != nil {
			_ = r.sess.Close()
		}
		if r.st != nil {
			_ = r.st.Close()
		}
	})
}

// connectRuntime dials edge and wires up adapters, store, router, and node handler.
// On any failure after partial allocation the allocated resources are closed.
func connectRuntime(ctx context.Context, cfg *config.NodeConfig, logger *zap.Logger, dialer DialFunc) (*runtimeOwner, error) {
	result, err := dialer(ctx, cfg.Transport.EdgeAddr, cfg.Transport.Token, logger)
	if err != nil {
		return nil, fmt.Errorf("dial edge: %w", err)
	}
	owner := &runtimeOwner{sess: result.Session}

	reg, err := adapters.BuildFromPayload(result.Config, logger)
	if err != nil {
		owner.close()
		return nil, fmt.Errorf("build adapters: %w", err)
	}
	owner.reg = reg

	dsn, err := storeDSN(result.Config.GetRuntime().GetWorkspaceRoot())
	if err != nil {
		owner.close()
		return nil, err
	}
	st, err := store.New(dsn, logger)
	if err != nil {
		owner.close()
		return nil, fmt.Errorf("store: %w", err)
	}
	owner.st = st

	if err := reg.Start(ctx); err != nil {
		owner.close()
		return nil, fmt.Errorf("start adapters: %w", err)
	}

	rtr := router.New(reg, logger)
	globalConcurrency := int(result.Config.GetRuntime().GetConcurrency())
	n := node.New(result.NodeID, rtr, st, globalConcurrency, os.Stdout, logger)
	result.Session.SetEventHandler(func(event *iop.EdgeNodeEvent) {
		printEdgeEvent(os.Stdout, event)
	})
	result.Session.SetHandler(n)

	logger.Info("connected to edge",
		zap.String("node_id", result.NodeID),
		zap.String("alias", result.Alias),
	)
	return owner, nil
}

// Module returns the fx options that wire the node application.
func Module(cfg *config.NodeConfig, opts ...Option) fx.Option {
	mo := &moduleOpts{
		dialer: transport.DialEdge,
		sleeper: func(ctx context.Context, d time.Duration) {
			if d <= 0 {
				return
			}
			timer := time.NewTimer(d)
			select {
			case <-ctx.Done():
				timer.Stop()
			case <-timer.C:
			}
		},
	}
	for _, opt := range opts {
		opt(mo)
	}

	return fx.Options(
		fx.Provide(
			func() *config.NodeConfig { return cfg },

			func(cfg *config.NodeConfig) (*zap.Logger, error) {
				return observability.NewLogger(cfg.Logging.Level, cfg.Logging.Pretty)
			},
		),

		fx.Invoke(func(lc fx.Lifecycle, cfg *config.NodeConfig, logger *zap.Logger, shutdowner fx.Shutdowner) {
			var mu sync.Mutex
			var current *runtimeOwner
			var cancelSupervisor context.CancelFunc

			lc.Append(fx.Hook{
				OnStart: func(ctx context.Context) error {
					owner, err := connectRuntime(ctx, cfg, logger, mo.dialer)
					if err != nil {
						return fmt.Errorf("bootstrap: %w", err)
					}

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

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

					go func() {
						sess := owner.sess
						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
							}

							maxAttempts := cfg.Reconnect.MaxAttempts
							if maxAttempts <= 0 {
								maxAttempts = 10
							}
							intervalSec := cfg.Reconnect.IntervalSec
							if intervalSec < 0 {
								intervalSec = 10
							}

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

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

								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, cfg, logger, mo.dialer)
								dialCancel()
								if dialErr == nil {
									newOwner = o
									break
								}

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

							if newOwner == nil {
								mu.Lock()
								prev := current
								current = nil
								mu.Unlock()
								if prev != nil {
									prev.close()
								}
								logger.Error("reconnect exhausted, shutting down node",
									zap.Int("max_attempts", maxAttempts),
								)
								_ = shutdowner.Shutdown(fx.ExitCode(1))
								return
							}

							mu.Lock()
							prev := current
							current = newOwner
							mu.Unlock()
							if prev != nil {
								prev.close()
							}

							sess = newOwner.sess
						}
					}()

					go func() {
						if err := observability.ServeMetrics(cfg.Metrics.Port); err != nil {
							logger.Warn("metrics server exited", zap.Error(err))
						}
					}()
					return nil
				},
				OnStop: func(_ context.Context) error {
					if cancelSupervisor != nil {
						cancelSupervisor()
					}
					mu.Lock()
					owner := current
					current = nil
					mu.Unlock()
					if owner != nil {
						owner.close()
					}
					return nil
				},
			})
		}),
	)
}

func storeDSN(workspaceRoot string) (string, error) {
	if workspaceRoot == "" {
		return "file:iop.db?cache=shared&mode=rwc", nil
	}
	if err := os.MkdirAll(workspaceRoot, 0o755); err != nil {
		return "", fmt.Errorf("bootstrap: workspace: %w", err)
	}
	return "file:" + filepath.Join(workspaceRoot, "iop.db") + "?cache=shared&mode=rwc", nil
}

func printEdgeEvent(out io.Writer, event *iop.EdgeNodeEvent) {
	if out == nil {
		return
	}
	switch event.GetType() {
	case events.TypeEdgeDisconnected:
		fmt.Fprintf(out, "[edge-event] disconnected reason=%q%s\n", event.GetReason(), transportCloseDetail(event.GetMetadata()))
	default:
		fmt.Fprintf(out, "[edge-event] %s reason=%q%s\n", event.GetType(), event.GetReason(), transportCloseDetail(event.GetMetadata()))
	}
}

func transportCloseDetail(metadata map[string]string) string {
	if metadata == nil {
		return ""
	}
	detail := ""
	if reason := metadata[events.MetadataTransportCloseReason]; reason != "" {
		detail += fmt.Sprintf(" transport_close_reason=%q", reason)
	}
	if err := metadata[events.MetadataTransportCloseError]; err != "" {
		detail += fmt.Sprintf(" transport_close_error=%q", err)
	}
	return detail
}