Initial commit.
This commit is contained in:
Kalzu Rekku
@@ -0,0 +1,27 @@
|
||||
BINARY_NAME = bybit_btcusdt_ingest
|
||||
|
||||
.PHONY: all build test clean run deps fmt vet
|
||||
|
||||
all: deps build
|
||||
|
||||
deps:
|
||||
go mod tidy
|
||||
|
||||
build: deps
|
||||
go build -o $(BINARY_NAME) .
|
||||
|
||||
run: build
|
||||
./$(BINARY_NAME)
|
||||
|
||||
test:
|
||||
go test -v -race ./...
|
||||
|
||||
fmt:
|
||||
go fmt ./...
|
||||
|
||||
vet: deps
|
||||
go vet ./...
|
||||
|
||||
clean:
|
||||
rm -f $(BINARY_NAME)
|
||||
rm -rf data/
|
||||
+168
@@ -0,0 +1,168 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"log"
|
||||
"math"
|
||||
"sync"
|
||||
)
|
||||
|
||||
const bucketDurationMs = 5000 // 5 seconds
|
||||
|
||||
// Aggregator accumulates ticks into 5-second buckets and writes completed
|
||||
// feature rows to the features database.
|
||||
type Aggregator struct {
|
||||
mu sync.Mutex
|
||||
currentBucket int64 // epoch ms of current bucket start
|
||||
ticks []Tick // ticks within the current bucket
|
||||
featDB *sql.DB // connection to features.db
|
||||
lastPrice float64 // closing price of the previous bucket (for log return)
|
||||
}
|
||||
|
||||
// NewAggregator creates an aggregator with a connection to the features database.
|
||||
func NewAggregator(sm *StorageManager) (*Aggregator, error) {
|
||||
db, err := sm.OpenFeaturesDB()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Attempt to load the last close price from features.db for continuity
|
||||
var lastPrice float64
|
||||
err = db.QueryRow(`
|
||||
SELECT close_price FROM five_second_features
|
||||
ORDER BY timestamp DESC LIMIT 1
|
||||
`).Scan(&lastPrice)
|
||||
if err != nil {
|
||||
lastPrice = 0 // Will be set from first tick
|
||||
}
|
||||
|
||||
return &Aggregator{
|
||||
featDB: db,
|
||||
lastPrice: lastPrice,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ProcessTick adds a tick to the current bucket. If the tick crosses a 5-second
|
||||
// boundary, the previous bucket is aggregated and flushed to the database.
|
||||
func (a *Aggregator) ProcessTick(tick Tick) {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
|
||||
// Determine which 5-second bucket this tick belongs to
|
||||
tickBucket := (tick.Timestamp / bucketDurationMs) * bucketDurationMs
|
||||
|
||||
if a.currentBucket == 0 {
|
||||
// First tick ever — initialize the bucket
|
||||
a.currentBucket = tickBucket
|
||||
}
|
||||
|
||||
if tickBucket > a.currentBucket {
|
||||
// This tick belongs to a new bucket — flush the previous one
|
||||
if len(a.ticks) > 0 {
|
||||
a.flushBucket()
|
||||
}
|
||||
a.currentBucket = tickBucket
|
||||
a.ticks = a.ticks[:0] // reset slice, keep backing array
|
||||
}
|
||||
|
||||
a.ticks = append(a.ticks, tick)
|
||||
}
|
||||
|
||||
// Close shuts down the aggregator, flushing any remaining bucket and closing the DB.
|
||||
func (a *Aggregator) Close() {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
|
||||
if len(a.ticks) > 0 {
|
||||
a.flushBucket()
|
||||
}
|
||||
if a.featDB != nil {
|
||||
a.featDB.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// flushBucket computes features from the accumulated ticks and writes to features.db.
|
||||
// Must be called while holding a.mu.
|
||||
func (a *Aggregator) flushBucket() {
|
||||
if len(a.ticks) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
bucket := a.computeFeatures(a.ticks)
|
||||
|
||||
_, err := a.featDB.Exec(`
|
||||
INSERT OR IGNORE INTO five_second_features
|
||||
(timestamp, log_return, realized_vol, ofi, volume_sum, close_price)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
`, bucket.Timestamp, bucket.LogReturn, bucket.RealizedVol,
|
||||
bucket.OFI, bucket.VolumeSum, bucket.ClosePrice)
|
||||
|
||||
if err != nil {
|
||||
log.Printf("[aggregator] Failed to write feature bucket: %v", err)
|
||||
}
|
||||
|
||||
// Update lastPrice for next bucket's log return calculation
|
||||
a.lastPrice = bucket.ClosePrice
|
||||
}
|
||||
|
||||
// computeFeatures calculates all feature columns from a slice of ticks.
|
||||
func (a *Aggregator) computeFeatures(ticks []Tick) FeatureBucket {
|
||||
n := len(ticks)
|
||||
startPrice := ticks[0].Price
|
||||
closePrice := ticks[n-1].Price
|
||||
|
||||
// Log return: ln(close / start_of_bucket_or_prev_close)
|
||||
refPrice := a.lastPrice
|
||||
if refPrice == 0 {
|
||||
refPrice = startPrice
|
||||
}
|
||||
logReturn := math.Log(closePrice / refPrice)
|
||||
|
||||
// Order Flow Imbalance (OFI) and total volume
|
||||
var buyVol, sellVol, volumeSum float64
|
||||
for _, t := range ticks {
|
||||
volumeSum += t.Volume
|
||||
if t.Side == "Buy" {
|
||||
buyVol += t.Volume
|
||||
} else {
|
||||
sellVol += t.Volume
|
||||
}
|
||||
}
|
||||
ofi := buyVol - sellVol
|
||||
|
||||
// Realized volatility: standard deviation of tick-to-tick log returns
|
||||
realizedVol := 0.0
|
||||
if n > 1 {
|
||||
logReturns := make([]float64, n-1)
|
||||
for i := 1; i < n; i++ {
|
||||
if ticks[i-1].Price > 0 {
|
||||
logReturns[i-1] = math.Log(ticks[i].Price / ticks[i-1].Price)
|
||||
}
|
||||
}
|
||||
|
||||
// Calculate mean
|
||||
var sum float64
|
||||
for _, lr := range logReturns {
|
||||
sum += lr
|
||||
}
|
||||
mean := sum / float64(len(logReturns))
|
||||
|
||||
// Calculate variance
|
||||
var variance float64
|
||||
for _, lr := range logReturns {
|
||||
diff := lr - mean
|
||||
variance += diff * diff
|
||||
}
|
||||
variance /= float64(len(logReturns))
|
||||
realizedVol = math.Sqrt(variance)
|
||||
}
|
||||
|
||||
return FeatureBucket{
|
||||
Timestamp: a.currentBucket,
|
||||
LogReturn: logReturn,
|
||||
RealizedVol: realizedVol,
|
||||
OFI: ofi,
|
||||
VolumeSum: volumeSum,
|
||||
ClosePrice: closePrice,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
)
|
||||
|
||||
// Config holds all tunable parameters for the ingest engine.
|
||||
type Config struct {
|
||||
// WebSocket endpoint for Bybit V5 public linear trades.
|
||||
WebSocketURL string `json:"websocket_url"`
|
||||
|
||||
// Symbol to subscribe to.
|
||||
Symbol string `json:"symbol"`
|
||||
|
||||
// DataDir is the base directory for all database files.
|
||||
DataDir string `json:"data_dir"`
|
||||
|
||||
// HotRetentionHours is how many hours of raw ticks to keep in the hot DB.
|
||||
HotRetentionHours int `json:"hot_retention_hours"`
|
||||
|
||||
// FeatureRetentionDays is how many days of 5s features to keep.
|
||||
FeatureRetentionDays int `json:"feature_retention_days"`
|
||||
|
||||
// WriterFlushIntervalMs is how often (ms) the batch writer flushes to disk.
|
||||
WriterFlushIntervalMs int `json:"writer_flush_interval_ms"`
|
||||
|
||||
// WriterBatchSize is the max number of ticks before a forced flush.
|
||||
WriterBatchSize int `json:"writer_batch_size"`
|
||||
|
||||
// TickChannelBuffer is the capacity of the tick channel between WS reader and DB writer.
|
||||
TickChannelBuffer int `json:"tick_channel_buffer"`
|
||||
|
||||
// MaintenanceIntervalMinutes controls how often the hourly maintenance runs.
|
||||
MaintenanceIntervalMinutes int `json:"maintenance_interval_minutes"`
|
||||
}
|
||||
|
||||
// DefaultConfig returns a Config populated with sensible defaults.
|
||||
func DefaultConfig() Config {
|
||||
return Config{
|
||||
WebSocketURL: "wss://stream.bybit.com/v5/public/linear",
|
||||
Symbol: "BTCUSDT",
|
||||
DataDir: "data",
|
||||
HotRetentionHours: 12,
|
||||
FeatureRetentionDays: 30,
|
||||
WriterFlushIntervalMs: 500,
|
||||
WriterBatchSize: 100,
|
||||
TickChannelBuffer: 10000,
|
||||
MaintenanceIntervalMinutes: 60,
|
||||
}
|
||||
}
|
||||
|
||||
// LoadConfig reads a JSON config file and merges with defaults.
|
||||
// If the file does not exist, defaults are returned and a config file is written.
|
||||
func LoadConfig(path string) (Config, error) {
|
||||
cfg := DefaultConfig()
|
||||
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
// Write default config for user reference
|
||||
if writeErr := writeDefaultConfig(path, cfg); writeErr != nil {
|
||||
return cfg, fmt.Errorf("failed to write default config: %w", writeErr)
|
||||
}
|
||||
fmt.Printf("No config found, wrote defaults to %s\n", path)
|
||||
return cfg, nil
|
||||
}
|
||||
return cfg, fmt.Errorf("failed to read config: %w", err)
|
||||
}
|
||||
|
||||
if err := json.Unmarshal(data, &cfg); err != nil {
|
||||
return cfg, fmt.Errorf("failed to parse config: %w", err)
|
||||
}
|
||||
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
func writeDefaultConfig(path string, cfg Config) error {
|
||||
data, err := json.MarshalIndent(cfg, "", " ")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return os.WriteFile(path, data, 0644)
|
||||
}
|
||||
+11
@@ -0,0 +1,11 @@
|
||||
{
|
||||
"websocket_url": "wss://stream.bybit.com/v5/public/linear",
|
||||
"symbol": "BTCUSDT",
|
||||
"data_dir": "data",
|
||||
"hot_retention_hours": 12,
|
||||
"feature_retention_days": 30,
|
||||
"writer_flush_interval_ms": 500,
|
||||
"writer_batch_size": 100,
|
||||
"tick_channel_buffer": 10000,
|
||||
"maintenance_interval_minutes": 60
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
module bybit_btcusdt_ingest
|
||||
|
||||
go 1.25.0
|
||||
|
||||
require (
|
||||
modernc.org/sqlite v1.53.0
|
||||
nhooyr.io/websocket v1.8.17
|
||||
)
|
||||
|
||||
require (
|
||||
github.com/dustin/go-humanize v1.0.1 // indirect
|
||||
github.com/google/uuid v1.6.0 // indirect
|
||||
github.com/mattn/go-isatty v0.0.20 // indirect
|
||||
github.com/ncruces/go-strftime v1.0.0 // indirect
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
|
||||
golang.org/x/sys v0.44.0 // indirect
|
||||
modernc.org/libc v1.73.4 // indirect
|
||||
modernc.org/mathutil v1.7.1 // indirect
|
||||
modernc.org/memory v1.11.0 // indirect
|
||||
)
|
||||
@@ -0,0 +1,53 @@
|
||||
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
|
||||
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
|
||||
github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs=
|
||||
github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA=
|
||||
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
|
||||
github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
|
||||
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
|
||||
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
|
||||
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
|
||||
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
|
||||
golang.org/x/mod v0.36.0 h1:JJjpVx6myfUsUdAzZuOSTTmRE0PfZeNWzzvKrP7amb4=
|
||||
golang.org/x/mod v0.36.0/go.mod h1:moc6ELqsWcOw5Ef3xVprK5ul/MvtVvkIXLziUOICjUQ=
|
||||
golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
|
||||
golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
|
||||
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.44.0 h1:ildZl3J4uzeKP07r2F++Op7E9B29JRUy+a27EibtBTQ=
|
||||
golang.org/x/sys v0.44.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
|
||||
golang.org/x/tools v0.45.0 h1:18qN3FAooORvApf5XjCXgsuayZOEtXf6JK18I3+ONa8=
|
||||
golang.org/x/tools v0.45.0/go.mod h1:LuUGqqaXcXMEFEruIVJVm5mgDD8vww/z/SR1gQ4uE/0=
|
||||
modernc.org/cc/v4 v4.28.4 h1:Hd/4Es+MBj+/7hSdZaisNyu6bv3V0Dp2MdllyfqaH+c=
|
||||
modernc.org/cc/v4 v4.28.4/go.mod h1:OnovgIhbbMXMu1aISnJ0wvVD1KnW+cAUJkIrAWh+kVI=
|
||||
modernc.org/ccgo/v4 v4.34.4 h1:OVnSOWQjVKOYkFxoHYB+qQmSHK5gqMqARM+K9DpR/Ws=
|
||||
modernc.org/ccgo/v4 v4.34.4/go.mod h1:qdKqE8FNIYyysougB1RX9MxCzp5oJOcQXSobANJ4TuE=
|
||||
modernc.org/fileutil v1.4.0 h1:j6ZzNTftVS054gi281TyLjHPp6CPHr2KCxEXjEbD6SM=
|
||||
modernc.org/fileutil v1.4.0/go.mod h1:EqdKFDxiByqxLk8ozOxObDSfcVOv/54xDs/DUHdvCUU=
|
||||
modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI=
|
||||
modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito=
|
||||
modernc.org/gc/v3 v3.1.3 h1:6QAplYyVO+KdPW3pGnqmJDUxtkec8ooEWvks/hhU3lc=
|
||||
modernc.org/gc/v3 v3.1.3/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY=
|
||||
modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks=
|
||||
modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI=
|
||||
modernc.org/libc v1.73.4 h1:+ra4Ui8ngyt8HDcO1FTDPWlkAh6yOdaO2yAoh8MddQA=
|
||||
modernc.org/libc v1.73.4/go.mod h1:DXZ3eO8qMCNn2SnmTNCiC71nJ9Rcq3PsnpU6Vc4rWK8=
|
||||
modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU=
|
||||
modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg=
|
||||
modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI=
|
||||
modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw=
|
||||
modernc.org/opt v0.2.0 h1:tGyef5ApycA7FSEOMraay9SaTk5zmbx7Tu+cJs4QKZg=
|
||||
modernc.org/opt v0.2.0/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns=
|
||||
modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w=
|
||||
modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE=
|
||||
modernc.org/sqlite v1.53.0 h1:20WG8N9q4ji/dEqGk4uiI0c6OPjSeLTNYGFCc3+7c1M=
|
||||
modernc.org/sqlite v1.53.0/go.mod h1:xoEpOIpGrgT48H5iiyt/YXPCZPEzlfmfFwtk8Lklw8s=
|
||||
modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0=
|
||||
modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A=
|
||||
modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y=
|
||||
modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM=
|
||||
nhooyr.io/websocket v1.8.17 h1:KEVeLJkUywCKVsnLIDlD/5gtayKp8VoCkksHCGGfT9Y=
|
||||
nhooyr.io/websocket v1.8.17/go.mod h1:rN9OFWIUwuxg4fR5tELlYC04bXYowCP9GX47ivo2l+c=
|
||||
@@ -0,0 +1,112 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log"
|
||||
"os"
|
||||
"os/signal"
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
)
|
||||
|
||||
func main() {
|
||||
log.SetFlags(log.LstdFlags | log.Lmicroseconds)
|
||||
log.Println("=== Bybit BTC/USDT Tick Ingest Engine ===")
|
||||
|
||||
// Load configuration
|
||||
cfg, err := LoadConfig("config.json")
|
||||
if err != nil {
|
||||
log.Fatalf("Config error: %v", err)
|
||||
}
|
||||
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. ===")
|
||||
}
|
||||
+353
@@ -0,0 +1,353 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
_ "modernc.org/sqlite"
|
||||
)
|
||||
|
||||
// StorageManager handles all SQLite database operations: initialization,
|
||||
// raw tick writing, feature writing, hourly migration, and pruning.
|
||||
type StorageManager struct {
|
||||
cfg Config
|
||||
hotDBPath string
|
||||
featDBPath string
|
||||
archiveDir string
|
||||
}
|
||||
|
||||
// NewStorageManager creates directories, initializes databases, and returns a ready manager.
|
||||
func NewStorageManager(cfg Config) (*StorageManager, error) {
|
||||
archiveDir := filepath.Join(cfg.DataDir, "archive")
|
||||
if err := os.MkdirAll(archiveDir, 0755); err != nil {
|
||||
return nil, fmt.Errorf("create archive dir: %w", err)
|
||||
}
|
||||
|
||||
sm := &StorageManager{
|
||||
cfg: cfg,
|
||||
hotDBPath: filepath.Join(cfg.DataDir, "hot_ticks.db"),
|
||||
featDBPath: filepath.Join(cfg.DataDir, "features.db"),
|
||||
archiveDir: archiveDir,
|
||||
}
|
||||
|
||||
if err := sm.initHotDB(); err != nil {
|
||||
return nil, fmt.Errorf("init hot db: %w", err)
|
||||
}
|
||||
if err := sm.initFeaturesDB(); err != nil {
|
||||
return nil, fmt.Errorf("init features db: %w", err)
|
||||
}
|
||||
|
||||
return sm, nil
|
||||
}
|
||||
|
||||
// OpenHotDB returns a new connection to the hot ticks database.
|
||||
func (sm *StorageManager) OpenHotDB() (*sql.DB, error) {
|
||||
db, err := sql.Open("sqlite", sm.hotDBPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if _, err := db.Exec("PRAGMA journal_mode=WAL"); err != nil {
|
||||
db.Close()
|
||||
return nil, err
|
||||
}
|
||||
if _, err := db.Exec("PRAGMA synchronous=NORMAL"); err != nil {
|
||||
db.Close()
|
||||
return nil, err
|
||||
}
|
||||
return db, nil
|
||||
}
|
||||
|
||||
// OpenFeaturesDB returns a new connection to the features database.
|
||||
func (sm *StorageManager) OpenFeaturesDB() (*sql.DB, error) {
|
||||
db, err := sql.Open("sqlite", sm.featDBPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if _, err := db.Exec("PRAGMA journal_mode=WAL"); err != nil {
|
||||
db.Close()
|
||||
return nil, err
|
||||
}
|
||||
if _, err := db.Exec("PRAGMA synchronous=NORMAL"); err != nil {
|
||||
db.Close()
|
||||
return nil, err
|
||||
}
|
||||
return db, nil
|
||||
}
|
||||
|
||||
func (sm *StorageManager) initHotDB() error {
|
||||
db, err := sm.OpenHotDB()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
_, err = db.Exec(`
|
||||
PRAGMA auto_vacuum=INCREMENTAL;
|
||||
|
||||
CREATE TABLE IF NOT EXISTS btc_ticks (
|
||||
timestamp INTEGER NOT NULL,
|
||||
price REAL NOT NULL,
|
||||
volume REAL NOT NULL,
|
||||
side TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_timestamp ON btc_ticks(timestamp);
|
||||
`)
|
||||
return err
|
||||
}
|
||||
|
||||
func (sm *StorageManager) initFeaturesDB() error {
|
||||
db, err := sm.OpenFeaturesDB()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
_, err = db.Exec(`
|
||||
PRAGMA auto_vacuum=INCREMENTAL;
|
||||
|
||||
CREATE TABLE IF NOT EXISTS five_second_features (
|
||||
timestamp INTEGER PRIMARY KEY,
|
||||
log_return REAL NOT NULL,
|
||||
realized_vol REAL NOT NULL,
|
||||
ofi REAL NOT NULL,
|
||||
volume_sum REAL NOT NULL,
|
||||
close_price REAL NOT NULL
|
||||
);
|
||||
`)
|
||||
return err
|
||||
}
|
||||
|
||||
// weeklyArchivePath resolves the archive DB file path for a given timestamp.
|
||||
func (sm *StorageManager) weeklyArchivePath(timestampMs int64) string {
|
||||
t := time.UnixMilli(timestampMs)
|
||||
year, week := t.ISOWeek()
|
||||
filename := fmt.Sprintf("btc_ticks_%d_W%02d.db", year, week)
|
||||
return filepath.Join(sm.archiveDir, filename)
|
||||
}
|
||||
|
||||
// initArchiveDB ensures the archive database has the correct schema.
|
||||
func (sm *StorageManager) initArchiveDB(path string) error {
|
||||
db, err := sql.Open("sqlite", path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
_, err = db.Exec(`
|
||||
PRAGMA journal_mode=WAL;
|
||||
PRAGMA synchronous=NORMAL;
|
||||
|
||||
CREATE TABLE IF NOT EXISTS btc_ticks (
|
||||
timestamp INTEGER NOT NULL,
|
||||
price REAL NOT NULL,
|
||||
volume REAL NOT NULL,
|
||||
side TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_timestamp ON btc_ticks(timestamp);
|
||||
`)
|
||||
return err
|
||||
}
|
||||
|
||||
// RunHourlyMaintenance performs tick migration and feature pruning.
|
||||
func (sm *StorageManager) RunHourlyMaintenance() {
|
||||
nowMs := time.Now().UnixMilli()
|
||||
log.Println("[maintenance] Starting hourly database maintenance...")
|
||||
|
||||
// Part A: Migrate raw ticks older than retention window
|
||||
cutoff12h := nowMs - int64(sm.cfg.HotRetentionHours)*60*60*1000
|
||||
if err := sm.migrateRawTicks(cutoff12h); err != nil {
|
||||
log.Printf("[maintenance] Raw tick migration failed: %v", err)
|
||||
}
|
||||
|
||||
// Part B: Prune features older than retention window
|
||||
cutoffFeatures := nowMs - int64(sm.cfg.FeatureRetentionDays)*24*60*60*1000
|
||||
if err := sm.pruneFeatures(cutoffFeatures); err != nil {
|
||||
log.Printf("[maintenance] Feature pruning failed: %v", err)
|
||||
}
|
||||
|
||||
log.Println("[maintenance] Hourly maintenance complete.")
|
||||
}
|
||||
|
||||
// StartupRecovery checks for stale ticks and migrates them before normal operation.
|
||||
func (sm *StorageManager) StartupRecovery() error {
|
||||
log.Println("[startup] Checking for stale ticks in hot database...")
|
||||
|
||||
db, err := sm.OpenHotDB()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
cutoffMs := time.Now().UnixMilli() - int64(sm.cfg.HotRetentionHours)*60*60*1000
|
||||
var count int64
|
||||
err = db.QueryRow("SELECT COUNT(*) FROM btc_ticks WHERE timestamp < ?", cutoffMs).Scan(&count)
|
||||
if err != nil {
|
||||
return fmt.Errorf("count stale ticks: %w", err)
|
||||
}
|
||||
|
||||
if count > 0 {
|
||||
log.Printf("[startup] Found %d stale ticks, triggering immediate migration...", count)
|
||||
if err := sm.migrateRawTicks(cutoffMs); err != nil {
|
||||
return fmt.Errorf("startup migration: %w", err)
|
||||
}
|
||||
} else {
|
||||
log.Println("[startup] No stale ticks found, hot database is clean.")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (sm *StorageManager) migrateRawTicks(cutoffMs int64) error {
|
||||
db, err := sm.OpenHotDB()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
// Find the range of ticks that need migration
|
||||
var minTS, maxTS sql.NullInt64
|
||||
err = db.QueryRow(`
|
||||
SELECT MIN(timestamp), MAX(timestamp)
|
||||
FROM btc_ticks
|
||||
WHERE timestamp < ?
|
||||
`, cutoffMs).Scan(&minTS, &maxTS)
|
||||
if err != nil {
|
||||
return fmt.Errorf("query migration range: %w", err)
|
||||
}
|
||||
|
||||
if !minTS.Valid || !maxTS.Valid {
|
||||
log.Println("[maintenance] No raw ticks older than cutoff to migrate.")
|
||||
return nil
|
||||
}
|
||||
|
||||
// Process ticks week by week to handle cross-week boundaries
|
||||
currentTs := minTS.Int64
|
||||
for currentTs <= maxTS.Int64 {
|
||||
t := time.UnixMilli(currentTs)
|
||||
year, week := t.ISOWeek()
|
||||
archivePath := sm.weeklyArchivePath(currentTs)
|
||||
|
||||
log.Printf("[maintenance] Migrating ticks for %d W%02d -> %s",
|
||||
year, week, filepath.Base(archivePath))
|
||||
|
||||
// Calculate the start of the next ISO week (Monday 00:00:00 UTC)
|
||||
nextWeekStart := sm.nextISOWeekStartMs(currentTs)
|
||||
|
||||
// The upper bound for this chunk: either next week start or cutoff, whichever is smaller
|
||||
chunkEnd := nextWeekStart
|
||||
if cutoffMs < chunkEnd {
|
||||
chunkEnd = cutoffMs
|
||||
}
|
||||
|
||||
// Ensure archive DB exists and has schema
|
||||
if err := sm.initArchiveDB(archivePath); err != nil {
|
||||
return fmt.Errorf("init archive db: %w", err)
|
||||
}
|
||||
|
||||
// Attach and atomically migrate
|
||||
if err := sm.atomicMigrate(db, archivePath, currentTs, chunkEnd); err != nil {
|
||||
return fmt.Errorf("atomic migrate: %w", err)
|
||||
}
|
||||
|
||||
currentTs = nextWeekStart
|
||||
}
|
||||
|
||||
// Reclaim space
|
||||
if _, err := db.Exec("PRAGMA incremental_vacuum(500)"); err != nil {
|
||||
log.Printf("[maintenance] incremental_vacuum warning: %v", err)
|
||||
}
|
||||
|
||||
log.Println("[maintenance] Raw tick migration completed successfully.")
|
||||
return nil
|
||||
}
|
||||
|
||||
// nextISOWeekStartMs calculates the epoch ms of the next Monday 00:00:00 UTC
|
||||
// relative to the given timestamp.
|
||||
func (sm *StorageManager) nextISOWeekStartMs(timestampMs int64) int64 {
|
||||
t := time.UnixMilli(timestampMs).UTC()
|
||||
// Calculate days until next Monday
|
||||
daysUntilMonday := (8 - int(t.Weekday())) % 7
|
||||
if daysUntilMonday == 0 {
|
||||
daysUntilMonday = 7
|
||||
}
|
||||
nextMonday := time.Date(t.Year(), t.Month(), t.Day()+daysUntilMonday, 0, 0, 0, 0, time.UTC)
|
||||
return nextMonday.UnixMilli()
|
||||
}
|
||||
|
||||
func (sm *StorageManager) atomicMigrate(hotDB *sql.DB, archivePath string, fromMs, toMs int64) error {
|
||||
// Attach the archive database
|
||||
_, err := hotDB.Exec("ATTACH DATABASE ? AS archive", archivePath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("attach archive: %w", err)
|
||||
}
|
||||
defer hotDB.Exec("DETACH DATABASE archive")
|
||||
|
||||
tx, err := hotDB.Begin()
|
||||
if err != nil {
|
||||
return fmt.Errorf("begin migration tx: %w", err)
|
||||
}
|
||||
|
||||
// Insert into archive
|
||||
_, err = tx.Exec(`
|
||||
INSERT INTO archive.btc_ticks (timestamp, price, volume, side)
|
||||
SELECT timestamp, price, volume, side
|
||||
FROM main.btc_ticks
|
||||
WHERE timestamp >= ? AND timestamp < ?
|
||||
`, fromMs, toMs)
|
||||
if err != nil {
|
||||
tx.Rollback()
|
||||
return fmt.Errorf("insert into archive: %w", err)
|
||||
}
|
||||
|
||||
// Delete from hot
|
||||
result, err := tx.Exec(`
|
||||
DELETE FROM main.btc_ticks
|
||||
WHERE timestamp >= ? AND timestamp < ?
|
||||
`, fromMs, toMs)
|
||||
if err != nil {
|
||||
tx.Rollback()
|
||||
return fmt.Errorf("delete from hot: %w", err)
|
||||
}
|
||||
|
||||
migrated, _ := result.RowsAffected()
|
||||
if err := tx.Commit(); err != nil {
|
||||
return fmt.Errorf("commit migration: %w", err)
|
||||
}
|
||||
|
||||
log.Printf("[maintenance] Migrated %d ticks to archive.", migrated)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (sm *StorageManager) pruneFeatures(cutoffMs int64) error {
|
||||
db, err := sm.OpenFeaturesDB()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
cutoffTime := time.UnixMilli(cutoffMs)
|
||||
log.Printf("[maintenance] Pruning features older than %s", cutoffTime.Format(time.RFC3339))
|
||||
|
||||
result, err := db.Exec("DELETE FROM five_second_features WHERE timestamp < ?", cutoffMs)
|
||||
if err != nil {
|
||||
return fmt.Errorf("delete old features: %w", err)
|
||||
}
|
||||
|
||||
deleted, _ := result.RowsAffected()
|
||||
if deleted > 0 {
|
||||
log.Printf("[maintenance] Pruned %d old feature rows.", deleted)
|
||||
if _, err := db.Exec("PRAGMA incremental_vacuum(100)"); err != nil {
|
||||
log.Printf("[maintenance] feature vacuum warning: %v", err)
|
||||
}
|
||||
} else {
|
||||
log.Println("[maintenance] No feature rows needed pruning.")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
package main
|
||||
|
||||
// Tick represents a single raw trade event from the Bybit WebSocket.
|
||||
type Tick struct {
|
||||
Timestamp int64 // Epoch millisecond timestamp
|
||||
Price float64 // Transacted trade price
|
||||
Volume float64 // Trade quantity
|
||||
Side string // "Buy" or "Sell"
|
||||
}
|
||||
|
||||
// FeatureBucket holds aggregated 5-second feature data ready for insertion into features.db.
|
||||
type FeatureBucket struct {
|
||||
Timestamp int64 // Epoch millisecond (start of 5s bucket)
|
||||
LogReturn float64 // ln(Price_end / Price_start)
|
||||
RealizedVol float64 // Volatility of ticks inside the bucket
|
||||
OFI float64 // Net volume (Buy volume - Sell volume)
|
||||
VolumeSum float64 // Total volume exchanged
|
||||
ClosePrice float64 // Final transaction price in the bucket
|
||||
}
|
||||
|
||||
// BybitWSMessage represents the top-level WebSocket message from Bybit V5 publicTrade.
|
||||
type BybitWSMessage struct {
|
||||
Topic string `json:"topic"`
|
||||
Type string `json:"type"`
|
||||
TS int64 `json:"ts"`
|
||||
Data []BybitTradeRaw `json:"data"`
|
||||
}
|
||||
|
||||
// BybitTradeRaw represents a single trade object within the Bybit WebSocket data array.
|
||||
// Price and Volume arrive as strings from the API and need conversion.
|
||||
type BybitTradeRaw struct {
|
||||
T int64 `json:"T"` // Timestamp (ms) that the order is filled
|
||||
S string `json:"s"` // Symbol name
|
||||
SD string `json:"S"` // Side of taker: "Buy" or "Sell"
|
||||
V string `json:"v"` // Trade size (string)
|
||||
P string `json:"p"` // Trade price (string)
|
||||
I string `json:"i"` // Trade ID
|
||||
}
|
||||
+143
@@ -0,0 +1,143 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"nhooyr.io/websocket"
|
||||
)
|
||||
|
||||
// Ingestor connects to the Bybit V5 WebSocket, parses trade messages,
|
||||
// sends ticks to the writer channel, and feeds ticks to the aggregator.
|
||||
type Ingestor struct {
|
||||
cfg Config
|
||||
tickCh chan<- Tick
|
||||
aggregator *Aggregator
|
||||
}
|
||||
|
||||
// NewIngestor creates an Ingestor wired to the tick channel and aggregator.
|
||||
func NewIngestor(cfg Config, tickCh chan<- Tick, agg *Aggregator) *Ingestor {
|
||||
return &Ingestor{
|
||||
cfg: cfg,
|
||||
tickCh: tickCh,
|
||||
aggregator: agg,
|
||||
}
|
||||
}
|
||||
|
||||
// Run connects to the WebSocket and processes messages until ctx is cancelled.
|
||||
// It automatically reconnects on connection failures.
|
||||
func (ing *Ingestor) Run(ctx context.Context) {
|
||||
log.Println("[ingestor] Starting WebSocket ingestor...")
|
||||
defer log.Println("[ingestor] Ingestor stopped.")
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
if err := ing.connectAndConsume(ctx); err != nil {
|
||||
if ctx.Err() != nil {
|
||||
return // Context cancelled, clean exit
|
||||
}
|
||||
log.Printf("[ingestor] Connection error: %v. Reconnecting in 5s...", err)
|
||||
select {
|
||||
case <-time.After(5 * time.Second):
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (ing *Ingestor) connectAndConsume(ctx context.Context) error {
|
||||
url := ing.cfg.WebSocketURL
|
||||
log.Printf("[ingestor] Connecting to %s ...", url)
|
||||
|
||||
conn, _, err := websocket.Dial(ctx, url, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("dial: %w", err)
|
||||
}
|
||||
defer conn.CloseNow()
|
||||
|
||||
// Set a generous read limit for large trade batches (up to 1024 trades per message)
|
||||
conn.SetReadLimit(1 << 20) // 1 MB
|
||||
|
||||
// Subscribe to public trades
|
||||
topic := fmt.Sprintf("publicTrade.%s", ing.cfg.Symbol)
|
||||
subMsg := map[string]interface{}{
|
||||
"op": "subscribe",
|
||||
"args": []string{topic},
|
||||
}
|
||||
subJSON, _ := json.Marshal(subMsg)
|
||||
|
||||
if err := conn.Write(ctx, websocket.MessageText, subJSON); err != nil {
|
||||
return fmt.Errorf("subscribe: %w", err)
|
||||
}
|
||||
log.Printf("[ingestor] Subscribed to %s", topic)
|
||||
|
||||
// Read loop
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
conn.Close(websocket.StatusNormalClosure, "shutting down")
|
||||
return nil
|
||||
default:
|
||||
}
|
||||
|
||||
_, data, err := conn.Read(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("read: %w", err)
|
||||
}
|
||||
|
||||
ing.handleMessage(data)
|
||||
}
|
||||
}
|
||||
|
||||
func (ing *Ingestor) handleMessage(data []byte) {
|
||||
var msg BybitWSMessage
|
||||
if err := json.Unmarshal(data, &msg); err != nil {
|
||||
// Could be a subscription confirmation or ping/pong — ignore
|
||||
return
|
||||
}
|
||||
|
||||
// Only process trade data messages
|
||||
if msg.Data == nil || len(msg.Data) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
for _, raw := range msg.Data {
|
||||
price, err := strconv.ParseFloat(raw.P, 64)
|
||||
if err != nil {
|
||||
log.Printf("[ingestor] bad price %q: %v", raw.P, err)
|
||||
continue
|
||||
}
|
||||
volume, err := strconv.ParseFloat(raw.V, 64)
|
||||
if err != nil {
|
||||
log.Printf("[ingestor] bad volume %q: %v", raw.V, err)
|
||||
continue
|
||||
}
|
||||
|
||||
tick := Tick{
|
||||
Timestamp: raw.T,
|
||||
Price: price,
|
||||
Volume: volume,
|
||||
Side: raw.SD,
|
||||
}
|
||||
|
||||
// Feed to aggregator (5s feature bucketing) inline
|
||||
ing.aggregator.ProcessTick(tick)
|
||||
|
||||
// Send to writer channel (non-blocking drop if channel full)
|
||||
select {
|
||||
case ing.tickCh <- tick:
|
||||
default:
|
||||
log.Println("[ingestor] WARNING: tick channel full, dropping tick")
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,104 @@
|
||||
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))
|
||||
}
|
||||
Reference in New Issue
Block a user