bybit_btcusdt_ingest/orderbook_handler.go

package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"log"
	"sort"
	"strconv"
	"sync"
	"time"
)

// OrderbookHandler maintains in-memory L2 orderbook state and periodically flushes snapshots & features.
type OrderbookHandler struct {
	cfg      Config
	storage  *StreamStorage
	mu       sync.Mutex
	bids     map[float64]float64 // price -> size
	asks     map[float64]float64 // price -> size
	lastTS   int64
	hotDB    *sql.DB
	featDB   *sql.DB
	stopChan chan struct{}
	wg       sync.WaitGroup
}

func NewOrderbookHandler(cfg Config, sm *StorageManager) (*OrderbookHandler, error) {
	ss := sm.GetStreamStorage("orderbook")
	if ss == nil {
		return nil, fmt.Errorf("orderbook stream storage not found")
	}

	hotDB, err := OpenDBWithAutoVacuum(ss.DBPath("hot_snapshots.db"))
	if err != nil {
		return nil, fmt.Errorf("open hot_snapshots db: %w", err)
	}

	featDB, err := OpenDBWithAutoVacuum(ss.DBPath("features.db"))
	if err != nil {
		hotDB.Close()
		return nil, fmt.Errorf("open orderbook features db: %w", err)
	}

	ob := &OrderbookHandler{
		cfg:      cfg,
		storage:  ss,
		bids:     make(map[float64]float64),
		asks:     make(map[float64]float64),
		hotDB:    hotDB,
		featDB:   featDB,
		stopChan: make(chan struct{}),
	}

	interval := time.Duration(cfg.Streams.Orderbook.SnapshotIntervalMs) * time.Millisecond
	if interval <= 0 {
		interval = 5 * time.Second
	}

	ob.wg.Add(1)
	go ob.runSnapshotLoop(interval)

	return ob, nil
}

func (ob *OrderbookHandler) Topics() []string {
	depth := ob.cfg.Streams.Orderbook.Depth
	if depth <= 0 {
		depth = 50
	}
	return []string{fmt.Sprintf("orderbook.%d.%s", depth, ob.cfg.Symbol)}
}

func (ob *OrderbookHandler) HandleMessage(data []byte) {
	var msg BybitOrderbookMessage
	if err := json.Unmarshal(data, &msg); err != nil {
		return
	}

	ob.mu.Lock()
	defer ob.mu.Unlock()

	ob.lastTS = msg.TS
	if ob.lastTS == 0 {
		ob.lastTS = time.Now().UnixMilli()
	}

	if msg.Type == "snapshot" {
		ob.bids = make(map[float64]float64)
		ob.asks = make(map[float64]float64)
	}

	for _, b := range msg.Data.B {
		if len(b) < 2 {
			continue
		}
		p, _ := strconv.ParseFloat(b[0], 64)
		s, _ := strconv.ParseFloat(b[1], 64)
		if s == 0 {
			delete(ob.bids, p)
		} else {
			ob.bids[p] = s
		}
	}

	for _, a := range msg.Data.A {
		if len(a) < 2 {
			continue
		}
		p, _ := strconv.ParseFloat(a[0], 64)
		s, _ := strconv.ParseFloat(a[1], 64)
		if s == 0 {
			delete(ob.asks, p)
		} else {
			ob.asks[p] = s
		}
	}
}

func (ob *OrderbookHandler) runSnapshotLoop(interval time.Duration) {
	defer ob.wg.Done()
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-ob.stopChan:
			ob.takeSnapshot()
			return
		case <-ticker.C:
			ob.takeSnapshot()
		}
	}
}

func (ob *OrderbookHandler) takeSnapshot() {
	ob.mu.Lock()
	if len(ob.bids) == 0 || len(ob.asks) == 0 {
		ob.mu.Unlock()
		return
	}

	ts := (ob.lastTS / 5000) * 5000

	// Sort bids descending, asks ascending
	sortedBids := make([]OrderbookLevel, 0, len(ob.bids))
	for p, s := range ob.bids {
		sortedBids = append(sortedBids, OrderbookLevel{Price: p, Size: s})
	}
	sort.Slice(sortedBids, func(i, j int) bool {
		return sortedBids[i].Price > sortedBids[j].Price
	})

	sortedAsks := make([]OrderbookLevel, 0, len(ob.asks))
	for p, s := range ob.asks {
		sortedAsks = append(sortedAsks, OrderbookLevel{Price: p, Size: s})
	}
	sort.Slice(sortedAsks, func(i, j int) bool {
		return sortedAsks[i].Price < sortedAsks[j].Price
	})

	ob.mu.Unlock()

	// 1. Write top 50 snapshot levels to hot DB
	tx, err := ob.hotDB.Begin()
	if err == nil {
		stmt, err := tx.Prepare(`
			INSERT OR REPLACE INTO ob_snapshots (timestamp, level, bid_price, bid_size, ask_price, ask_size)
			VALUES (?, ?, ?, ?, ?, ?)
		`)
		if err == nil {
			maxLevels := 50
			for i := 0; i < maxLevels; i++ {
				var bp, bs, ap, as sql.NullFloat64
				if i < len(sortedBids) {
					bp = sql.NullFloat64{Float64: sortedBids[i].Price, Valid: true}
					bs = sql.NullFloat64{Float64: sortedBids[i].Size, Valid: true}
				}
				if i < len(sortedAsks) {
					ap = sql.NullFloat64{Float64: sortedAsks[i].Price, Valid: true}
					as = sql.NullFloat64{Float64: sortedAsks[i].Size, Valid: true}
				}
				if bp.Valid || ap.Valid {
					stmt.Exec(ts, i, bp, bs, ap, as)
				}
			}
			stmt.Close()
			tx.Commit()
		} else {
			tx.Rollback()
		}
	}

	// 2. Compute orderbook features
	bestBid := sortedBids[0].Price
	bestAsk := sortedAsks[0].Price
	spread := bestAsk - bestBid
	midPrice := (bestBid + bestAsk) / 2.0

	var bidDepth5, askDepth5, bidDepth20, askDepth20 float64
	for i := 0; i < len(sortedBids); i++ {
		if i < 5 {
			bidDepth5 += sortedBids[i].Size
		}
		if i < 20 {
			bidDepth20 += sortedBids[i].Size
		}
	}
	for i := 0; i < len(sortedAsks); i++ {
		if i < 5 {
			askDepth5 += sortedAsks[i].Size
		}
		if i < 20 {
			askDepth20 += sortedAsks[i].Size
		}
	}

	depthImb5 := 0.0
	if bidDepth5+askDepth5 > 0 {
		depthImb5 = (bidDepth5 - askDepth5) / (bidDepth5 + askDepth5)
	}

	depthImb20 := 0.0
	if bidDepth20+askDepth20 > 0 {
		depthImb20 = (bidDepth20 - askDepth20) / (bidDepth20 + askDepth20)
	}

	bid1Size := sortedBids[0].Size
	ask1Size := sortedAsks[0].Size
	weightedMid := midPrice
	if bid1Size+ask1Size > 0 {
		weightedMid = (bestBid*ask1Size + bestAsk*bid1Size) / (bid1Size + ask1Size)
	}

	// Top 10 VWAP (combined across bids and asks)
	var sumPV, sumV float64
	for i := 0; i < 10; i++ {
		if i < len(sortedBids) {
			sumPV += sortedBids[i].Price * sortedBids[i].Size
			sumV += sortedBids[i].Size
		}
		if i < len(sortedAsks) {
			sumPV += sortedAsks[i].Price * sortedAsks[i].Size
			sumV += sortedAsks[i].Size
		}
	}
	vwap10 := midPrice
	if sumV > 0 {
		vwap10 = sumPV / sumV
	}

	_, err = ob.featDB.Exec(`
		INSERT OR IGNORE INTO ob_features (
			timestamp, spread, mid_price, bid_depth_5, ask_depth_5, bid_depth_20, ask_depth_20,
			depth_imbalance_5, depth_imbalance_20, weighted_mid, vwap_10
		) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
	`, ts, spread, midPrice, bidDepth5, askDepth5, bidDepth20, askDepth20, depthImb5, depthImb20, weightedMid, vwap10)

	if err != nil {
		log.Printf("[orderbook_handler] features db insert error: %v", err)
	}
}

func (ob *OrderbookHandler) Close() {
	close(ob.stopChan)
	ob.wg.Wait()
	if ob.hotDB != nil {
		ob.hotDB.Close()
	}
	if ob.featDB != nil {
		ob.featDB.Close()
	}
}