feat: add prometheus and protobuf messages

Signed-off-by: kjuulh <contact@kjuulh.io>

crates/nodata/Cargo.toml

@@ -35,6 +35,7 @@ tokio-stream = "0.1.15"
 dagger-sdk = "0.11.10"
 rand = "0.8.5"
 notmad = "0.4.0"
+prometheus = "0.13.4"
 
 [dev-dependencies]
 tracing-test = { version = "0.2.5" } #, features = ["no-env-filter"] }

crates/nodata/src/gen/nodata.v1.rs

@@ -52,5 +52,11 @@ pub struct PingRequest {
 #[derive(Clone, Copy, PartialEq, ::prost::Message)]
 pub struct PingResponse {
 }
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct Log {
+    #[prost(bytes="vec", repeated, tag="1")]
+    pub messages: ::prost::alloc::vec::Vec<::prost::alloc::vec::Vec<u8>>,
+}
 include!("nodata.v1.tonic.rs");
 // @@protoc_insertion_point(module)

crates/nodata/src/grpc.rs

@@ -17,13 +17,26 @@ include!("gen/nodata.v1.rs");
 pub struct GrpcServer {
     state: SharedState,
     host: SocketAddr,
+    counter: prometheus::Counter,
 }
 
 impl GrpcServer {
     pub fn new(state: &SharedState, host: SocketAddr) -> Self {
+        let counter = prometheus::Counter::with_opts(prometheus::Opts::new(
+            "messages_pr_second",
+            "how many messages are ingested pr second",
+        ))
+        .unwrap();
+
+        state
+            .metrics_registry
+            .register(Box::new(counter.clone()))
+            .expect("to be able to register metrics");
+
         Self {
             state: state.clone(),
             host,
+            counter,
         }
     }
 }
@@ -45,6 +58,8 @@ impl no_data_service_server::NoDataService for GrpcServer {
             "handling event"
         );
 
+        self.counter.inc();
+
         self.state.ingest().publish(req).await.map_err(|e| {
             tracing::warn!(error = e.to_string(), "failed to handle ingest of data");
             tonic::Status::internal(e.to_string())

crates/nodata/src/http.rs

@@ -3,11 +3,14 @@ use std::net::SocketAddr;
 use anyhow::Context;
 use axum::async_trait;
 use axum::extract::MatchedPath;
+use axum::extract::State;
 use axum::http::Request;
 use axum::routing::get;
 use axum::Router;
 use notmad::Component;
 use notmad::MadError;
+use prometheus::Encoder;
+use prometheus::TextEncoder;
 use tokio_util::sync::CancellationToken;
 use tower_http::trace::TraceLayer;
 
@@ -32,6 +35,7 @@ impl Component for HttpServer {
     async fn run(&self, _cancellation_token: CancellationToken) -> Result<(), notmad::MadError> {
         let app = Router::new()
             .route("/", get(root))
+            .route("/metrics", get(metrics))
             .with_state(self.state.clone())
             .layer(
                 TraceLayer::new_for_http().make_span_with(|request: &Request<_>| {
@@ -63,3 +67,15 @@ impl Component for HttpServer {
 async fn root() -> &'static str {
     "Hello, nodata!"
 }
+
+async fn metrics(State(state): State<SharedState>) -> String {
+    let encoder = TextEncoder::new();
+    let metrics = state.metrics_registry.gather();
+
+    let mut buffer = Vec::new();
+    encoder
+        .encode(&metrics, &mut buffer)
+        .expect("to be able to encode metrics");
+
+    String::from_utf8(buffer).expect("to be able to encode from utf8")
+}

crates/nodata/src/main.rs

@@ -28,9 +28,17 @@ struct Command {
 #[derive(Subcommand)]
 enum Commands {
     Serve {
-        #[arg(env = "SERVICE_HOST", long, default_value = "127.0.0.1:3000")]
+        #[arg(
+            env = "SERVICE_HOST",
+            long = "service-host",
+            default_value = "127.0.0.1:3000"
+        )]
         host: SocketAddr,
-        #[arg(env = "GRPC_SERVICE_HOST", long, default_value = "127.0.0.1:7900")]
+        #[arg(
+            env = "GRPC_SERVICE_HOST",
+            long = "service-grpc-host",
+            default_value = "127.0.0.1:7900"
+        )]
         grpc_host: SocketAddr,
     },
 
@@ -39,7 +47,7 @@ enum Commands {
         //host: String,
         #[arg(
             env = "GRPC_SERVICE_HOST",
-            long,
+            long = "service-grpc-host",
             default_value = "http://127.0.0.1:7900"
         )]
         grpc_host: String,
@@ -57,6 +65,14 @@ enum ClientCommands {
         #[arg(long)]
         value: String,
     },
+    PublishEvents {
+        #[arg(long)]
+        topic: String,
+        #[arg(long)]
+        size: usize,
+        #[arg(long)]
+        threads: usize,
+    },
     GetTopics {},
     SubscribeTopic {
         #[arg(long)]
@@ -99,6 +115,39 @@ async fn main() -> anyhow::Result<()> {
                     })
                     .await?;
             }
+            ClientCommands::PublishEvents {
+                topic,
+                size,
+                threads,
+            } => {
+                let mut handles = Vec::new();
+                for _ in 0..threads {
+                    let topic = topic.clone();
+                    let grpc_host = grpc_host.clone();
+                    let handle = tokio::spawn(async move {
+                        let mut client = create_client(grpc_host).await?;
+
+                        loop {
+                            let _ = client
+                                .publish_event(PublishEventRequest {
+                                    topic: topic.clone(),
+                                    value: vec![0; size],
+                                })
+                                .await?;
+                            //tokio::time::sleep(std::time::Duration::from_millis(5)).await;
+                        }
+
+                        #[allow(unreachable_code)]
+                        Ok::<(), anyhow::Error>(())
+                    });
+
+                    handles.push(handle);
+                }
+
+                for handle in handles {
+                    handle.await??;
+                }
+            }
             ClientCommands::GetTopics {} => {
                 let mut client = create_client(grpc_host).await?;
 

crates/nodata/src/services/staging.rs

@@ -1,4 +1,4 @@
-use std::{collections::BTreeMap, env::temp_dir, sync::Arc};
+use std::{collections::BTreeMap, sync::Arc};
 
 use tokio::sync::RwLock;
 

crates/nodata/src/state.rs

@@ -1,6 +1,7 @@
 use std::{ops::Deref, sync::Arc};
 
 use anyhow::Context;
+use prometheus::Registry;
 use sqlx::{Pool, Postgres};
 
 use crate::services::{consumers::Consumers, handler::Handler, staging::Staging};
@@ -27,6 +28,7 @@ pub struct State {
     pub staging: Staging,
     pub consumers: Consumers,
     pub handler: Handler,
+    pub metrics_registry: Registry,
 }
 
 impl State {
@@ -51,6 +53,7 @@ impl State {
             consumers: Consumers::new(),
             staging,
             handler,
+            metrics_registry: Registry::new(),
         })
     }
 }