kabu_metrics/
influxdb_actor.rs

1use eyre::{Result, eyre};
2use influxdb::{Client, ReadQuery, WriteQuery};
3use kabu_core_components::Component;
4use reth_tasks::TaskExecutor;
5use std::collections::HashMap;
6use std::time::Duration;
7use tokio::sync::broadcast;
8use tokio::time::timeout;
9use tracing::{error, info, warn};
10
11pub async fn start_influxdb_worker(
12    url: String,
13    database: String,
14    tags: HashMap<String, String>,
15    mut event_receiver: broadcast::Receiver<WriteQuery>,
16) {
17    let client = Client::new(url, database.clone());
18    let create_db_stmt = format!("CREATE DATABASE {database}");
19    let result = client.query(ReadQuery::new(create_db_stmt)).await;
20    match result {
21        Ok(_) => info!("Database created with name: {}", database),
22        Err(e) => info!("Database creation failed or already exists: {:?}", e),
23    }
24    // event_receiver is already a Receiver
25    loop {
26        let event_result = event_receiver.recv().await;
27        match event_result {
28            Ok(mut event) => {
29                for (key, value) in tags.iter() {
30                    event = event.add_tag(key, value.clone());
31                }
32                let client_clone = client.clone();
33                tokio::task::spawn(async move {
34                    match timeout(Duration::from_millis(2000), client_clone.query(event)).await {
35                        Ok(inner_result) => {
36                            if let Err(e) = inner_result {
37                                error!("InfluxDB Write failed: {:?}", e);
38                            }
39                        }
40                        Err(elapsed) => {
41                            error!("InfluxDB Query timed out: {}", elapsed);
42                        }
43                    }
44                });
45            }
46            Err(e) => match e {
47                tokio::sync::broadcast::error::RecvError::Closed => {
48                    error!("InfluxDB channel closed");
49                    break;
50                }
51                tokio::sync::broadcast::error::RecvError::Lagged(lagged) => {
52                    warn!("InfluxDB lagged: {:?}", lagged);
53                    continue;
54                }
55            },
56        }
57    }
58}
59
60#[derive(Clone)]
61pub struct InfluxDbWriterComponent {
62    url: String,
63    database: String,
64    tags: HashMap<String, String>,
65    influxdb_write_channel_rx: Option<broadcast::Sender<WriteQuery>>,
66}
67
68impl InfluxDbWriterComponent {
69    pub fn new(url: String, database: String, tags: HashMap<String, String>) -> Self {
70        Self { url, database, tags, influxdb_write_channel_rx: None }
71    }
72
73    pub fn with_channel(self, channel: Option<broadcast::Sender<WriteQuery>>) -> Self {
74        Self { influxdb_write_channel_rx: channel, ..self }
75    }
76}
77
78impl Component for InfluxDbWriterComponent {
79    fn spawn(self, executor: TaskExecutor) -> Result<()> {
80        let name = self.name();
81
82        let influxdb_write_channel_rx = self.influxdb_write_channel_rx.ok_or_else(|| eyre!("InfluxDB write channel is not set"))?;
83
84        let event_receiver = influxdb_write_channel_rx.subscribe();
85
86        executor.spawn_critical(name, start_influxdb_worker(self.url, self.database, self.tags, event_receiver));
87
88        Ok(())
89    }
90
91    fn name(&self) -> &'static str {
92        "InfluxDbWriterComponent"
93    }
94}
kabu_metrics/influxdb_actor.rs

kabu_metrics/
influxdb_actor.rs
