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"
	"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}

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

	dsn, err := storeDSN()
	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 := set.Registry.Start(ctx); err != nil {
		owner.close()
		return nil, fmt.Errorf("start adapters: %w", err)
	}

	rtr := router.New(set.Registry, logger)
	globalConcurrency := int(result.Config.GetRuntime().GetConcurrency())
	n := node.New(result.NodeID, rtr, st, globalConcurrency, os.Stdout, logger, set)
	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.NewLoggerWithFile(cfg.Logging.Level, cfg.Logging.Pretty, cfg.Logging.Path)
			},
		),

		fx.Invoke(func(lc fx.Lifecycle, cfg *config.NodeConfig, logger *zap.Logger, shutdowner fx.Shutdowner) {
			sup := &runtimeSupervisor{
				cfg:        cfg,
				logger:     logger,
				dialer:     mo.dialer,
				sleeper:    mo.sleeper,
				shutdowner: shutdowner,
			}
			lc.Append(fx.Hook{
				OnStart: func(ctx context.Context) error {
					return sup.start(ctx)
				},
				OnStop: func(_ context.Context) error {
					sup.stop()
					return nil
				},
			})
		}),
	)
}

func storeDSN() (string, error) {
	return "file: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
}