105 lines
2.2 KiB
Go
105 lines
2.2 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"log"
|
|
"time"
|
|
)
|
|
|
|
// 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
|
|
}
|
|
|
|
// 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 INTO btc_ticks (timestamp, price, volume, side) VALUES (?, ?, ?, ?)")
|
|
if err != nil {
|
|
log.Printf("[writer] prepare failed: %v", err)
|
|
tx.Rollback()
|
|
return
|
|
}
|
|
defer stmt.Close()
|
|
|
|
for _, t := range batch {
|
|
if _, err := stmt.Exec(t.Timestamp, t.Price, t.Volume, t.Side); 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))
|
|
}
|