170 lines
2.9 KiB
Go
170 lines
2.9 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"log"
|
|
"time"
|
|
)
|
|
|
|
const insertTickSQL = `
|
|
INSERT OR IGNORE INTO btc_ticks (
|
|
seq,
|
|
trade_id,
|
|
trade_ts,
|
|
message_ts,
|
|
recv_ts,
|
|
symbol,
|
|
side,
|
|
price,
|
|
volume,
|
|
tick_dir,
|
|
block_trade,
|
|
rpi
|
|
)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
`
|
|
|
|
// Writer drains the tick channel and batch-writes to hot_ticks.db.
|
|
type Writer struct {
|
|
tickCh <-chan Tick
|
|
hotDB *sql.DB
|
|
batchSize int
|
|
flushMs int
|
|
lastSeq int64
|
|
}
|
|
|
|
// NewWriter creates a Writer with its own hot DB connection.
|
|
func NewWriter(sm *StorageManager, tickCh <-chan Tick, cfg Config) (*Writer, error) {
|
|
db, err := sm.OpenHotDB()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
db.SetMaxOpenConns(1)
|
|
|
|
return &Writer{
|
|
tickCh: tickCh,
|
|
hotDB: db,
|
|
batchSize: cfg.WriterBatchSize,
|
|
flushMs: cfg.WriterFlushIntervalMs,
|
|
}, nil
|
|
}
|
|
|
|
// Run starts the writer loop. Blocks until ctx is cancelled.
|
|
func (w *Writer) Run(ctx context.Context) {
|
|
log.Println("[writer] Batch writer started.")
|
|
defer log.Println("[writer] Batch writer stopped.")
|
|
defer w.hotDB.Close()
|
|
|
|
batch := make([]Tick, 0, w.batchSize)
|
|
flushInterval := time.Duration(w.flushMs) * time.Millisecond
|
|
timer := time.NewTimer(flushInterval)
|
|
defer timer.Stop()
|
|
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
if len(batch) > 0 {
|
|
w.flush(batch)
|
|
}
|
|
return
|
|
case tick, ok := <-w.tickCh:
|
|
if !ok {
|
|
if len(batch) > 0 {
|
|
w.flush(batch)
|
|
}
|
|
return
|
|
}
|
|
batch = append(batch, tick)
|
|
if len(batch) >= w.batchSize {
|
|
w.flush(batch)
|
|
batch = batch[:0]
|
|
timer.Reset(flushInterval)
|
|
}
|
|
case <-timer.C:
|
|
if len(batch) > 0 {
|
|
w.flush(batch)
|
|
batch = batch[:0]
|
|
}
|
|
timer.Reset(flushInterval)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (w *Writer) flush(batch []Tick) {
|
|
if len(batch) == 0 {
|
|
return
|
|
}
|
|
|
|
tx, err := w.hotDB.Begin()
|
|
if err != nil {
|
|
log.Printf("[writer] begin tx failed: %v", err)
|
|
return
|
|
}
|
|
|
|
stmt, err := tx.Prepare(`
|
|
INSERT OR IGNORE INTO btc_ticks (
|
|
seq,
|
|
trade_id,
|
|
trade_ts,
|
|
message_ts,
|
|
recv_ts,
|
|
symbol,
|
|
side,
|
|
price,
|
|
volume,
|
|
tick_dir,
|
|
block_trade,
|
|
rpi
|
|
)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
`)
|
|
if err != nil {
|
|
log.Printf("[writer] prepare failed: %v", err)
|
|
tx.Rollback()
|
|
return
|
|
}
|
|
defer stmt.Close()
|
|
|
|
for _, t := range batch {
|
|
|
|
if t.Seq < w.lastSeq {
|
|
log.Printf(
|
|
"[writer] OUT OF ORDER SEQ: prev=%d current=%d",
|
|
w.lastSeq,
|
|
t.Seq,
|
|
)
|
|
}
|
|
|
|
w.lastSeq = t.Seq
|
|
|
|
_, err := stmt.Exec(
|
|
t.Seq,
|
|
t.TradeID,
|
|
t.TradeTS,
|
|
t.MessageTS,
|
|
t.RecvTS,
|
|
t.Symbol,
|
|
t.Side,
|
|
t.Price,
|
|
t.Volume,
|
|
t.TickDir,
|
|
t.BlockTrade,
|
|
t.RPI,
|
|
)
|
|
|
|
if err != nil {
|
|
log.Printf("[writer] insert failed: %v", err)
|
|
tx.Rollback()
|
|
return
|
|
}
|
|
}
|
|
|
|
if err := tx.Commit(); err != nil {
|
|
log.Printf("[writer] commit failed: %v", err)
|
|
return
|
|
}
|
|
|
|
log.Printf("[writer] Flushed %d ticks to hot_ticks.db", len(batch))
|
|
}
|