use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use crate::engine::traits::{BarSink, SinkError};
use crate::live_engine::CompletedBar;
use super::config::ClickHouseWriterConfig;
use super::flush_thread::{FlushCommand, FlushThreadMetrics, spawn_flush_thread};
use super::row::ClickHouseBarRow;
pub struct ClickHouseWriterSink {
tx: std::sync::mpsc::SyncSender<FlushCommand>,
flush_thread: Option<std::thread::JoinHandle<()>>,
metrics: Arc<FlushThreadMetrics>,
bars_sent: AtomicU64,
bars_dropped: AtomicU64,
}
impl ClickHouseWriterSink {
pub fn new(config: ClickHouseWriterConfig) -> Self {
let (tx, rx) = std::sync::mpsc::sync_channel(config.channel_capacity);
let metrics = Arc::new(FlushThreadMetrics::default());
let flush_thread = spawn_flush_thread(rx, config, Arc::clone(&metrics));
Self {
tx,
flush_thread: Some(flush_thread),
metrics,
bars_sent: AtomicU64::new(0),
bars_dropped: AtomicU64::new(0),
}
}
pub fn metrics(&self) -> &FlushThreadMetrics {
&self.metrics
}
pub fn shared_metrics(&self) -> Arc<FlushThreadMetrics> {
Arc::clone(&self.metrics)
}
pub fn bars_sent(&self) -> u64 {
self.bars_sent.load(Ordering::Relaxed)
}
pub fn bars_dropped(&self) -> u64 {
self.bars_dropped.load(Ordering::Relaxed)
}
}
impl BarSink for ClickHouseWriterSink {
fn on_bar(&mut self, bar: &CompletedBar) -> Result<(), SinkError> {
let row = ClickHouseBarRow::from_completed_bar(bar);
match self.tx.try_send(FlushCommand::Bar(Box::new(row))) {
Ok(()) => {
self.bars_sent.fetch_add(1, Ordering::Relaxed);
Ok(())
}
Err(std::sync::mpsc::TrySendError::Full(cmd)) => {
self.bars_dropped.fetch_add(1, Ordering::Relaxed);
if let FlushCommand::Bar(boxed_row) = cmd {
match super::dead_letter::write_dead_letter_single(*boxed_row) {
Ok(path) => {
tracing::warn!(
path = %path.display(),
"backpressure - bar dead-lettered to Parquet"
);
}
Err(dl_err) => {
tracing::error!(
error = %dl_err,
"CRITICAL: backpressure AND dead-letter write failed -- bar lost"
);
}
}
} else {
tracing::warn!("clickhouse flush thread backpressure - non-bar command dropped");
}
Err(SinkError::Recoverable(
"clickhouse flush thread backpressure".into(),
))
}
Err(std::sync::mpsc::TrySendError::Disconnected(_)) => {
Err(SinkError::Unrecoverable(
"clickhouse flush thread died".into(),
))
}
}
}
fn flush(&mut self) -> Result<(), SinkError> {
self.tx
.send(FlushCommand::Flush)
.map_err(|_| SinkError::Unrecoverable("clickhouse flush thread died".into()))
}
fn name(&self) -> &str {
"clickhouse"
}
}
impl Drop for ClickHouseWriterSink {
fn drop(&mut self) {
let _ = self.tx.send(FlushCommand::Shutdown);
if let Some(handle) = self.flush_thread.take() {
let _ = handle.join();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use opendeviationbar_core::{FixedPoint, OpenDeviationBar};
use std::sync::Arc;
fn make_bar() -> CompletedBar {
let trade = opendeviationbar_core::Tick {
ref_id: 1,
price: FixedPoint::from_str("50000.0").unwrap(),
volume: FixedPoint::from_str("1.0").unwrap(),
first_sub_id: 1,
last_sub_id: 1,
timestamp: opendeviationbar_core::normalize_timestamp(1_700_000_000_000),
is_buyer_maker: false,
is_best_match: None,
best_bid: None,
best_ask: None,
};
CompletedBar {
symbol: Arc::from("BTCUSDT"),
threshold_decimal_bps: 250,
bar: OpenDeviationBar::new(&trade),
}
}
#[test]
fn test_sink_name() {
let config = ClickHouseWriterConfig {
url: "http://127.0.0.1:1".to_string(),
channel_capacity: 10,
max_retries: 0,
flush_period_ms: 60_000,
..Default::default()
};
let mut sink = ClickHouseWriterSink::new(config);
assert_eq!(sink.name(), "clickhouse");
let _ = sink.flush();
}
#[test]
fn test_sink_on_bar_sends_to_channel() {
let config = ClickHouseWriterConfig {
url: "http://127.0.0.1:1".to_string(),
channel_capacity: 100,
max_retries: 0,
flush_period_ms: 60_000,
..Default::default()
};
let mut sink = ClickHouseWriterSink::new(config);
let bar = make_bar();
assert!(sink.on_bar(&bar).is_ok());
assert_eq!(sink.bars_sent(), 1);
assert_eq!(sink.bars_dropped(), 0);
}
#[test]
fn test_sink_backpressure() {
let config = ClickHouseWriterConfig {
url: "http://127.0.0.1:1".to_string(),
channel_capacity: 1, max_retries: 0,
flush_period_ms: 60_000,
max_rows: 10_000, ..Default::default()
};
let mut sink = ClickHouseWriterSink::new(config);
let bar = make_bar();
assert!(sink.on_bar(&bar).is_ok());
let result = sink.on_bar(&bar);
assert!(matches!(result, Err(SinkError::Recoverable(_))));
assert_eq!(sink.bars_dropped(), 1);
}
#[test]
fn test_sink_graceful_shutdown() {
let config = ClickHouseWriterConfig {
url: "http://127.0.0.1:1".to_string(),
channel_capacity: 100,
max_retries: 0,
flush_period_ms: 60_000,
..Default::default()
};
let mut sink = ClickHouseWriterSink::new(config);
let bar = make_bar();
sink.on_bar(&bar).unwrap();
sink.on_bar(&bar).unwrap();
assert_eq!(sink.bars_sent(), 2);
drop(sink);
}
}