Added klines, liquidations, tickers and trades to be recorded. Bundled the stats as cli argument.
This commit is contained in:
Kalzu Rekku
@@ -10,6 +10,8 @@ import (
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"bybit_btcusdt_ingest/utils/stats"
|
||||
)
|
||||
|
||||
func main() {
|
||||
@@ -30,7 +32,6 @@ func main() {
|
||||
|
||||
switch command {
|
||||
case "run":
|
||||
// Check for daemon flag
|
||||
isDaemon := false
|
||||
for _, arg := range os.Args {
|
||||
if arg == "--daemon" {
|
||||
@@ -65,6 +66,22 @@ func main() {
|
||||
sm.RunHourlyMaintenance()
|
||||
log.Println("Maintenance completed.")
|
||||
|
||||
case "stats":
|
||||
dataDir := cfg.DataDir
|
||||
noColor := false
|
||||
|
||||
for i := 2; i < len(os.Args); i++ {
|
||||
arg := os.Args[i]
|
||||
if arg == "--no-color" {
|
||||
noColor = true
|
||||
} else if (arg == "--data" || arg == "-data") && i+1 < len(os.Args) {
|
||||
dataDir = os.Args[i+1]
|
||||
i++
|
||||
}
|
||||
}
|
||||
|
||||
stats.PrintDashboard(dataDir, noColor)
|
||||
|
||||
case "help":
|
||||
printUsage()
|
||||
|
||||
@@ -77,10 +94,13 @@ func main() {
|
||||
func printUsage() {
|
||||
fmt.Println("Usage: engine [command] [--daemon]")
|
||||
fmt.Println("\nCommands:")
|
||||
fmt.Println(" run Start the WebSocket ingestor and processing engine")
|
||||
fmt.Println(" run Start the Multi-Stream 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(" recover Run startup recovery to migrate stale trade ticks")
|
||||
fmt.Println(" maintain Run a single maintenance cycle (cleanup/retention across all streams)")
|
||||
fmt.Println(" stats Display database statistics across all data streams")
|
||||
fmt.Println(" Use --data <path> to specify custom data directory")
|
||||
fmt.Println(" Use --no-color to disable ANSI color formatting")
|
||||
fmt.Println(" help Show this help message")
|
||||
}
|
||||
|
||||
@@ -89,7 +109,6 @@ func startDaemon(logPath string) {
|
||||
logPath = "engine.log"
|
||||
}
|
||||
|
||||
// Prepare arguments for the child process (remove --daemon)
|
||||
args := []string{}
|
||||
for _, arg := range os.Args[1:] {
|
||||
if arg != "--daemon" {
|
||||
@@ -97,10 +116,8 @@ func startDaemon(logPath string) {
|
||||
}
|
||||
}
|
||||
|
||||
// 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)
|
||||
@@ -109,7 +126,6 @@ func startDaemon(logPath string) {
|
||||
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)
|
||||
}
|
||||
@@ -119,55 +135,88 @@ func startDaemon(logPath string) {
|
||||
}
|
||||
|
||||
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)
|
||||
log.Println("=== Bybit Multi-Stream Ingest Engine ===")
|
||||
log.Printf("Symbol: %s | Data Dir: %s", cfg.Symbol, cfg.DataDir)
|
||||
|
||||
// Initialize storage (creates dirs, databases, tables)
|
||||
// Initialize storage (creates per-stream directories, databases, tables, and migrates old layout if needed)
|
||||
sm, err := NewStorageManager(cfg)
|
||||
if err != nil {
|
||||
log.Fatalf("Storage init error: %v", err)
|
||||
}
|
||||
|
||||
// Startup recovery: migrate any stale ticks from previous runs
|
||||
// Startup recovery for trades
|
||||
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)
|
||||
handlers := make(map[string]MessageHandler)
|
||||
|
||||
// Initialize aggregator (5-second feature bucketing)
|
||||
agg, err := NewAggregator(sm)
|
||||
if err != nil {
|
||||
log.Fatalf("Aggregator init error: %v", err)
|
||||
// Initialize enabled handlers
|
||||
if cfg.Streams.Trades.Enabled {
|
||||
th, err := NewTradeHandler(cfg, sm)
|
||||
if err != nil {
|
||||
log.Fatalf("TradeHandler init error: %v", err)
|
||||
}
|
||||
handlers["publicTrade"] = th
|
||||
log.Println("[init] Trades stream handler enabled.")
|
||||
}
|
||||
|
||||
// Initialize batch writer
|
||||
writer, err := NewWriter(sm, tickCh, cfg)
|
||||
if err != nil {
|
||||
log.Fatalf("Writer init error: %v", err)
|
||||
if cfg.Streams.Ticker.Enabled {
|
||||
tickerH, err := NewTickerHandler(cfg, sm)
|
||||
if err != nil {
|
||||
log.Fatalf("TickerHandler init error: %v", err)
|
||||
}
|
||||
handlers["tickers"] = tickerH
|
||||
log.Println("[init] Tickers stream handler enabled.")
|
||||
}
|
||||
|
||||
// Create ingestor
|
||||
ingestor := NewIngestor(cfg, tickCh, agg)
|
||||
if cfg.Streams.Klines.Enabled {
|
||||
klineH, err := NewKlineHandler(cfg, sm)
|
||||
if err != nil {
|
||||
log.Fatalf("KlineHandler init error: %v", err)
|
||||
}
|
||||
handlers["kline"] = klineH
|
||||
log.Printf("[init] Klines stream handler enabled (intervals: %v).", cfg.Streams.Klines.Intervals)
|
||||
}
|
||||
|
||||
if cfg.Streams.Orderbook.Enabled {
|
||||
obH, err := NewOrderbookHandler(cfg, sm)
|
||||
if err != nil {
|
||||
log.Fatalf("OrderbookHandler init error: %v", err)
|
||||
}
|
||||
handlers["orderbook"] = obH
|
||||
log.Printf("[init] Orderbook stream handler enabled (depth: %d).", cfg.Streams.Orderbook.Depth)
|
||||
}
|
||||
|
||||
if cfg.Streams.Liquidations.Enabled {
|
||||
liqH, err := NewLiquidationHandler(cfg, sm)
|
||||
if err != nil {
|
||||
log.Fatalf("LiquidationHandler init error: %v", err)
|
||||
}
|
||||
handlers["allLiquidation"] = liqH
|
||||
log.Println("[init] Liquidations stream handler enabled.")
|
||||
}
|
||||
|
||||
// Create multi-stream WebSocket Ingestor
|
||||
ingestor := NewIngestor(cfg, handlers)
|
||||
|
||||
// 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 1: Trade Handler background writer (if trades enabled)
|
||||
if th, ok := handlers["publicTrade"].(*TradeHandler); ok {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
th.writer.Run(ctx)
|
||||
}()
|
||||
}
|
||||
|
||||
// Goroutine 2: WebSocket Ingestor
|
||||
wg.Add(1)
|
||||
@@ -176,7 +225,7 @@ func runEngine(cfg Config) {
|
||||
ingestor.Run(ctx)
|
||||
}()
|
||||
|
||||
// Goroutine 3: Hourly Maintenance
|
||||
// Goroutine 3: Periodic Maintenance
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
@@ -201,13 +250,12 @@ func runEngine(cfg Config) {
|
||||
log.Printf("Received signal %v, initiating graceful shutdown...", sig)
|
||||
cancel()
|
||||
|
||||
// Close the tick channel so the writer drains remaining ticks
|
||||
close(tickCh)
|
||||
// Close all handlers
|
||||
for name, h := range handlers {
|
||||
log.Printf("[shutdown] Closing handler %s...", name)
|
||||
h.Close()
|
||||
}
|
||||
|
||||
// Flush any remaining aggregator bucket
|
||||
agg.Close()
|
||||
|
||||
// Wait for all goroutines to finish
|
||||
wg.Wait()
|
||||
log.Println("=== Engine shut down cleanly. ===")
|
||||
}
|
||||
Reference in New Issue
Block a user