-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
90dfa51
commit c53d4ca
Showing
12 changed files
with
334 additions
and
58 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,147 @@ | ||
| use std::time::Duration; | ||
|
|
||
| use anyhow::{anyhow, Context, Error}; | ||
| use clickhouse::{inserter::Inserter, Client}; | ||
| use serde::{Deserialize, Serialize}; | ||
| use tokio::sync::mpsc::{channel, Receiver, Sender}; | ||
| use tokio_util::{sync::CancellationToken, task::TaskTracker}; | ||
| use tracing::{debug, error, warn}; | ||
|
|
||
| use crate::cli; | ||
|
|
||
| #[derive(clickhouse::Row, Serialize, Deserialize)] | ||
| pub struct Row { | ||
| pub env: String, | ||
| pub hostname: String, | ||
| #[serde(with = "clickhouse::serde::time::datetime")] | ||
| pub date: time::OffsetDateTime, | ||
| #[serde(with = "clickhouse::serde::uuid")] | ||
| pub request_id: uuid::Uuid, | ||
| pub method: String, | ||
| pub http_version: String, | ||
| pub status: u16, | ||
| pub domain: String, | ||
| pub host: String, | ||
| pub path: String, | ||
| pub canister_id: String, | ||
| pub error_cause: String, | ||
| pub tls_version: String, | ||
| pub tls_cipher: String, | ||
| pub request_size: u64, | ||
| pub response_size: u64, | ||
| pub duration: f64, | ||
| pub duration_full: f64, | ||
| pub duration_conn: f64, | ||
| } | ||
|
|
||
| pub struct Clickhouse { | ||
| token: CancellationToken, | ||
| tracker: TaskTracker, | ||
| tx: Sender<Row>, | ||
| } | ||
|
|
||
| impl Clickhouse { | ||
| pub fn new(cli: &cli::Clickhouse) -> Result<Self, Error> { | ||
| let (tx, rx) = channel(65536); | ||
| let token = CancellationToken::new(); | ||
| let actor = ClickhouseActor::new(cli.clone(), rx)?; | ||
|
|
||
| let child_token = token.child_token(); | ||
| let tracker = TaskTracker::new(); | ||
| tracker.spawn(async move { | ||
| if let Err(e) = actor.run(child_token).await { | ||
| error!("Clickhouse: error during run: {e}"); | ||
| } | ||
| }); | ||
|
|
||
| Ok(Self { tx, tracker, token }) | ||
| } | ||
|
|
||
| pub fn send(&self, r: Row) { | ||
| // If it fails we'll lose the message, but it's better than to block & eat memory. | ||
| let _ = self.tx.try_send(r); | ||
| } | ||
|
|
||
| pub async fn stop(&self) { | ||
| self.token.cancel(); | ||
| self.tracker.close(); | ||
| self.tracker.wait().await; | ||
| } | ||
| } | ||
|
|
||
| pub struct ClickhouseActor { | ||
| inserter: Inserter<Row>, | ||
| rx: Receiver<Row>, | ||
| } | ||
|
|
||
| impl ClickhouseActor { | ||
| pub fn new(c: cli::Clickhouse, rx: Receiver<Row>) -> Result<Self, Error> { | ||
| let mut client = Client::default().with_url( | ||
| c.log_clickhouse_url | ||
| .ok_or_else(|| anyhow!("no URL specified"))?, | ||
| ); | ||
| if let Some(v) = c.log_clickhouse_user { | ||
| client = client.with_user(v); | ||
| } | ||
| if let Some(v) = c.log_clickhouse_pass { | ||
| client = client.with_password(v); | ||
| } | ||
| if let Some(v) = c.log_clickhouse_db { | ||
| client = client.with_database(v); | ||
| } | ||
|
|
||
| let inserter = client | ||
| .inserter( | ||
| &c.log_clickhouse_table | ||
| .ok_or_else(|| anyhow!("no table specified"))?, | ||
| )? | ||
| .with_max_entries(c.log_clickhouse_batch) | ||
| .with_period(Some(c.log_clickhouse_interval)) | ||
| .with_period_bias(0.1); // add 10% random variance to interval | ||
|
|
||
| Ok(Self { inserter, rx }) | ||
| } | ||
|
|
||
| async fn run(mut self, token: CancellationToken) -> Result<(), Error> { | ||
| let mut interval = tokio::time::interval(Duration::from_secs(1)); | ||
|
|
||
| warn!("Clickhouse: started"); | ||
| loop { | ||
| tokio::select! { | ||
| biased; | ||
|
|
||
| () = token.cancelled() => { | ||
| // Close the channel | ||
| self.rx.close(); | ||
|
|
||
| // Drain remaining rows | ||
| while let Some(v) = self.rx.recv().await { | ||
| self.inserter.write(&v).await.context("unable insert row")?; | ||
| } | ||
|
|
||
| // Flush the buffer | ||
| self.inserter.end().await.context("unable to flush buffer")?; | ||
| warn!("Clickhouse: stopped"); | ||
| return Ok(()); | ||
| }, | ||
|
|
||
| // Periodically poke inserter to commit if time has come. | ||
| // If the thresholds are not reached - it doesn't do anything. | ||
| _ = interval.tick() => { | ||
| match self.inserter.commit().await { | ||
| Ok(v) => debug!("Clickhouse: {} rows inserted", v.entries), | ||
| Err(e) => error!("Clickhouse: unable to commit: {e}"), | ||
| } | ||
| } | ||
|
|
||
| row = self.rx.recv() => { | ||
| if let Some(v) = row { | ||
| if let Err(e) = self.inserter.write(&v).await { | ||
| error!("Clickhouse: unable to insert row: {e}"); | ||
| } | ||
| } | ||
| } | ||
| } | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,3 +1,5 @@ | ||
| pub mod clickhouse; | ||
|
|
||
| use std::time::{SystemTime, UNIX_EPOCH}; | ||
|
|
||
| use anyhow::{Context, Error}; | ||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.