clickhouse_qol/worker/
insert_worker.rs

1use std::time::Duration;
2
3use clickhouse::{Client, Row};
4use serde::Serialize;
5use tokio::sync::mpsc::{Receiver, Sender};
6
7pub enum WorkerCommand<T>
8where
9    T: Row + Serialize,
10{
11    Insert(T),
12    Ping,
13}
14
15pub struct StreamInsertWorker<T>
16where
17    T: Row + Serialize,
18{
19    rx: Receiver<WorkerCommand<T>>,
20    client: Client,
21    max_size: u64,
22    period_sec: u64,
23    table_name: String,
24}
25
26impl<T> StreamInsertWorker<T>
27where
28    T: Row + Serialize,
29{
30    pub fn new(client: Client, table_name: String, max_size: u64, period_sec: u64) -> (Self, Sender<WorkerCommand<T>>) {
31        let (tx, rx) = tokio::sync::mpsc::channel::<WorkerCommand<T>>(1024);
32        let worker = StreamInsertWorker {
33            rx,
34            client,
35            max_size,
36            period_sec,
37            table_name,
38        };
39        (worker, tx)
40    }
41
42    pub async fn recv(&mut self) -> anyhow::Result<()> {
43        let mut current = 0;
44        let mut inserter = self
45            .client
46            .inserter(&self.table_name)?
47            .with_max_bytes(100_000)
48            .with_max_rows(self.max_size)
49            .with_period(Some(Duration::from_secs(self.period_sec)));
50
51        while let Some(command) = self.rx.recv().await {
52            match command {
53                WorkerCommand::Insert(data) => {
54                    let row = data;
55                    match inserter.write(&row) {
56                        Ok(_) => current += 1,
57                        Err(e) => {
58                            log::error!("failed to write: {:?}", e);
59                        }
60                    }
61
62                    if current >= self.max_size {
63                        match inserter.commit().await {
64                            Ok(res) => {
65                                log::info!("commit result: {:?}", res);
66                                current = 0;
67                            }
68                            Err(e) => {
69                                log::error!("failed to commit: {:?}", e);
70                            }
71                        }
72                    }
73                }
74                WorkerCommand::Ping => {
75                    if let Some(left) = inserter.time_left() {
76                        if left.as_millis() <= 100 {
77                            match inserter.commit().await {
78                                Ok(res) => {
79                                    if res.rows > 0 {
80                                        log::info!("commit result: {:?}", res);
81                                    }
82                                    current = 0;
83                                }
84                                Err(e) => {
85                                    log::error!("failed to commit: {:?}", e);
86                                }
87                            }
88                        }
89                    }
90                }
91            }
92        }
93
94        Ok(())
95    }
96}
clickhouse_qol/worker/insert_worker.rs

clickhouse_qol/worker/
insert_worker.rs
