bybit_btcusdt_ingest/main.go

package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

func main() {
	log.SetFlags(log.LstdFlags | log.Lmicroseconds)

	// Load configuration
	cfg, err := LoadConfig("config.json")
	if err != nil {
		log.Fatalf("Config error: %v", err)
	}

	if len(os.Args) < 2 {
		printUsage()
		return
	}

	command := os.Args[1]

	switch command {
	case "run":
		// Check for daemon flag
		isDaemon := false
		for _, arg := range os.Args {
			if arg == "--daemon" {
				isDaemon = true
				break
			}
		}

		if isDaemon {
			startDaemon(cfg.LogFile)
			return
		}
		runEngine(cfg)

	case "recover":
		log.Println("=== Starting Recovery Mode ===")
		sm, err := NewStorageManager(cfg)
		if err != nil {
			log.Fatal(err)
		}
		if err := sm.StartupRecovery(); err != nil {
			log.Fatal(err)
		}
		log.Println("Recovery completed.")

	case "maintain":
		log.Println("=== Starting Manual Maintenance ===")
		sm, err := NewStorageManager(cfg)
		if err != nil {
			log.Fatal(err)
		}
		sm.RunHourlyMaintenance()
		log.Println("Maintenance completed.")

	case "help":
		printUsage()

	default:
		log.Printf("Unknown command: %s", command)
		printUsage()
	}
}

func printUsage() {
	fmt.Println("Usage: engine [command] [--daemon]")
	fmt.Println("\nCommands:")
	fmt.Println("  run       Start the WebSocket ingestor and processing engine")
	fmt.Println("            Use --daemon to run in background and log to the file defined in config")
	fmt.Println("  recover   Run startup recovery to migrate stale ticks")
	fmt.Println("  maintain  Run a single maintenance cycle (cleanup/retention)")
	fmt.Println("  help      Show this help message")
}

func startDaemon(logPath string) {
	if logPath == "" {
		logPath = "engine.log"
	}

	// Prepare arguments for the child process (remove --daemon)
	args := []string{}
	for _, arg := range os.Args[1:] {
		if arg != "--daemon" {
			args = append(args, arg)
		}
	}

	// Prepare the command to re-run the current binary
	cmd := exec.Command(os.Args[0], args...)
	
	// Open the configured log file for output redirection
	logFile, err := os.OpenFile(logPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
	if err != nil {
		log.Fatalf("Failed to open log file %s: %v", logPath, err)
	}

	cmd.Stdout = logFile
	cmd.Stderr = logFile

	// Start the process in the background
	if err := cmd.Start(); err != nil {
		log.Fatalf("Failed to start daemon: %v", err)
	}

	fmt.Printf("Engine started in background (PID: %d). Logs: %s\n", cmd.Process.Pid, logPath)
	os.Exit(0)
}

func runEngine(cfg Config) {
	log.Println("=== Bybit BTC/USDT Tick Ingest Engine ===")
	log.Printf("Config: symbol=%s, hot_retention=%dh, feature_retention=%dd",
		cfg.Symbol, cfg.HotRetentionHours, cfg.FeatureRetentionDays)

	// Initialize storage (creates dirs, databases, tables)
	sm, err := NewStorageManager(cfg)
	if err != nil {
		log.Fatalf("Storage init error: %v", err)
	}

	// Startup recovery: migrate any stale ticks from previous runs
	if err := sm.StartupRecovery(); err != nil {
		log.Fatalf("Startup recovery error: %v", err)
	}

	// Create the tick channel (buffered to absorb bursts)
	tickCh := make(chan Tick, cfg.TickChannelBuffer)

	// Initialize aggregator (5-second feature bucketing)
	agg, err := NewAggregator(sm)
	if err != nil {
		log.Fatalf("Aggregator init error: %v", err)
	}

	// Initialize batch writer
	writer, err := NewWriter(sm, tickCh, cfg)
	if err != nil {
		log.Fatalf("Writer init error: %v", err)
	}

	// Create ingestor
	ingestor := NewIngestor(cfg, tickCh, agg)

	// Context for graceful shutdown
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Catch OS signals
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)

	var wg sync.WaitGroup

	// Goroutine 1: Batch Writer
	wg.Add(1)
	go func() {
		defer wg.Done()
		writer.Run(ctx)
	}()

	// Goroutine 2: WebSocket Ingestor
	wg.Add(1)
	go func() {
		defer wg.Done()
		ingestor.Run(ctx)
	}()

	// Goroutine 3: Hourly Maintenance
	wg.Add(1)
	go func() {
		defer wg.Done()
		maintInterval := time.Duration(cfg.MaintenanceIntervalMinutes) * time.Minute
		ticker := time.NewTicker(maintInterval)
		defer ticker.Stop()

		log.Printf("[maintenance] Scheduled every %d minutes.", cfg.MaintenanceIntervalMinutes)
		for {
			select {
			case <-ctx.Done():
				log.Println("[maintenance] Maintenance goroutine stopped.")
				return
			case <-ticker.C:
				sm.RunHourlyMaintenance()
			}
		}
	}()

	// Wait for shutdown signal
	sig := <-sigCh
	log.Printf("Received signal %v, initiating graceful shutdown...", sig)
	cancel()

	// Close the tick channel so the writer drains remaining ticks
	close(tickCh)

	// Flush any remaining aggregator bucket
	agg.Close()

	// Wait for all goroutines to finish
	wg.Wait()
	log.Println("=== Engine shut down cleanly. ===")
}