alt/services/worker/cmd/alt-worker/main.go

package main

import (
	"context"
	"fmt"
	"log/slog"
	"os"
	"os/signal"
	"strings"
	"syscall"

	"time"

	"github.com/jackc/pgx/v5/pgxpool"

	workerbacktest "git.toki-labs.com/toki/alt/services/worker/internal/backtest"
	"git.toki-labs.com/toki/alt/services/worker/internal/config"
	"git.toki-labs.com/toki/alt/services/worker/internal/jobs"
	"git.toki-labs.com/toki/alt/services/worker/internal/livetrading"
	"git.toki-labs.com/toki/alt/services/worker/internal/marketdata/aggregation"
	"git.toki-labs.com/toki/alt/services/worker/internal/marketdata/importer"
	"git.toki-labs.com/toki/alt/services/worker/internal/papertrading"
	"git.toki-labs.com/toki/alt/services/worker/internal/providers/kis"
	"git.toki-labs.com/toki/alt/services/worker/internal/scheduler"
	"git.toki-labs.com/toki/alt/services/worker/internal/socket"
	"git.toki-labs.com/toki/alt/services/worker/internal/storage/postgres"
)

func main() {
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	cfg := config.Load()
	runner := jobs.NewRunner()
	jobs.RegisterBuiltins(runner)

	// Store-backed socket surfaces are optional at process startup. If the pool
	// cannot be constructed the worker still serves hello, while market and
	// backtest handlers report unavailable instead of panicking.
	var deps socket.Deps
	var refreshStatusStore *scheduler.RefreshStatusStore
	pool, err := pgxpool.New(ctx, cfg.DatabaseURL)
	if err != nil {
		slog.Error("failed to create database pool; store-backed socket surfaces disabled", "error", err)
	} else {
		defer pool.Close()
		store := postgres.NewStore(pool)
		deps = storeBackedDeps(runner, store)

		// Wire refresh status store so the scheduler tick and socket handler
		// share the same in-memory store and the operator-facing query
		// surface reports actual runtime state.
		refreshStatusStore = scheduler.NewRefreshStatusStore()
		deps.RefreshStatus = refreshStatusStore

		if err := startScheduler(ctx, cfg, runner, refreshStatusStore); err != nil {
			slog.Error("failed to start scheduler", "error", err)
			os.Exit(1)
		}
	}

	server := socket.NewServer(cfg, deps)
	if err := server.Start(ctx); err != nil {
		slog.Error("failed to start worker socket server", "error", err)
		os.Exit(1)
	}

	slog.Info("worker ready",
		"redis_key_prefix", cfg.RedisKeyPrefix,
		"worker_queue", cfg.WorkerQueue,
		"handlers", runner.Len(),
		"backtest_read_surface", deps.Analysis != nil,
		"backtest_start_surface", deps.Starter != nil,
		"market_read_surface", deps.Instruments != nil && deps.Bars != nil,
		"market_import_surface", deps.DailyBarImporter != nil,
		"live_audit_surface", deps.Live != nil,
		"addr", fmt.Sprintf("%s:%d%s", cfg.Host, cfg.Port, cfg.SocketPath),
	)

	<-ctx.Done()
	if err := server.Stop(); err != nil {
		slog.Error("failed to stop worker socket server", "error", err)
		os.Exit(1)
	}
}

// storeBackedDeps wires every store-backed socket dependency around a single
// store, including the backtest execution rail. It is factored out of main so a
// unit test can assert the backtest start/read surfaces are wired without a live
// database. Backtest execution reads imported daily bars from storage and runs
// the bundled strategy resolver; the starter records the pending run and hands
// execution to the same runner the import jobs use, keeping a single async job
// rail rather than a second execution engine.
func storeBackedDeps(runner *jobs.Runner, store *postgres.Store) socket.Deps {
	kisProvider := kis.NewLiveProvider(kis.NewClient(kis.ConfigFromEnv()))
	// One importer instance backs both the async job handler and the socket
	// import command so the operator-facing command and scheduled jobs share the
	// same provider/store wiring.
	kisImporter := importer.New(kisProvider, store)
	jobs.RegisterDailyBarImportHandler(runner, kis.Capability(), kisImporter)

	barSource := workerbacktest.NewStorageBarSource(store, store)
	strategyPort := workerbacktest.NewBuiltInStrategyPort()
	engine := workerbacktest.NewEngine(barSource, strategyPort, store)
	jobs.RegisterRunBacktestHandler(runner, store, engine, time.Now)
	starter := jobs.NewBacktestStarter(runner, store)

	// Paper trading reuses the same imported-bar source and bundled strategy
	// resolver as backtest, but runs synchronously into a process-local
	// in-memory state keyed by account id (no durable store for readiness).
	paperEngine := papertrading.NewEngine(barSource, strategyPort)
	paperService := papertrading.NewService(paperEngine, time.Now)

	// Live trading service: nil broker means live order operations are unavailable
	// until a broker adapter is wired. The audit store is always injected so the
	// audit query surface is available as soon as the database pool is up.
	liveSvc := livetrading.NewService(nil)
	liveSvc.SetAuditStore(store)

	return socket.Deps{
		Starter:           starter,
		Analysis:          store,
		Results:           store,
		Instruments:       store,
		Bars:              store,
		DailyBarImporter:  kisImporter,
		MonthlyAggregator: aggregation.NewAggregator(store, store),
		Paper:             paperService,
		Live:              liveSvc,
	}
}

func startScheduler(ctx context.Context, cfg config.Config, runner *jobs.Runner, statusStore *scheduler.RefreshStatusStore) error {
	if !cfg.SchedulerEnabled {
		return nil
	}

	f, err := os.Open(cfg.ScheduleConfigPath)
	if err != nil {
		return fmt.Errorf("failed to open schedule config: %w", err)
	}
	defer f.Close()

	schedCfg, err := scheduler.LoadConfig(f)
	if err != nil {
		return fmt.Errorf("failed to load schedule config: %w", err)
	}

	// Store the config so the refresh status store can compute next_run values
	// for the socket handler surface.
	if statusStore != nil {
		statusStore.SetScheduleConfig(schedCfg)
	}

	capabilities := scheduler.CapabilityMap(kis.Capability())
	validation := scheduler.ValidateConfig(schedCfg, scheduler.ValidationOptions{
		Now:          time.Now(),
		Capabilities: capabilities,
	})
	if validation.HasRejected() {
		var reasons []string
		for _, item := range validation.Items {
			if item.Status == "rejected" {
				reasons = append(reasons, fmt.Sprintf("%s: %s", item.Schedule, item.Reason))
			}
		}
		return fmt.Errorf("invalid schedule config: %s", strings.Join(reasons, "; "))
	}

	slog.Info("starting scheduler loop",
		"scheduler_enabled", true,
		"schedule_count", len(schedCfg.Schedules),
	)

	go func() {
		ticker := time.NewTicker(1 * time.Minute)
		defer ticker.Stop()

		opts := scheduler.TickOptions{
			Capabilities: capabilities,
			Executor:     runner,
		}

		runSchedulerLoop(ctx, schedCfg, opts, ticker.C, statusStore)
	}()

	return nil
}

func runSchedulerLoop(ctx context.Context, schedCfg scheduler.Config, opts scheduler.TickOptions, ticks <-chan time.Time, statusStore *scheduler.RefreshStatusStore) {
	for {
		select {
		case <-ctx.Done():
			slog.Info("stopping scheduler loop")
			return
		case <-ticks:
			opts.Now = time.Now()
			res, err := scheduler.RunTick(ctx, schedCfg, opts)
			if err != nil {
				slog.Error("scheduler tick execution error", "error", err)
				continue
			}
			for _, item := range res.Items {
				if item.Status == "failed" {
					slog.Error("scheduler job failed", "schedule", item.Schedule, "symbol", item.Symbol, "error", item.Error)
				} else if item.Status == "skipped" {
					slog.Info("scheduler job skipped", "schedule", item.Schedule, "symbol", item.Symbol, "reason", item.Error)
				} else {
					slog.Info("scheduler job succeeded", "schedule", item.Schedule, "symbol", item.Symbol, "readiness", item.Readiness, "freshness", item.Freshness)
				}
			}
			// Record tick results in the refresh status store so the socket
			// handler surface reports actual runtime state.
			if statusStore != nil {
				now := opts.Now
				for _, sched := range schedCfg.Schedules {
					statusStore.RecordTick(sched.Name, res.Items, now)
				}
			}
		}
	}
}