chore: Adds UUID to http events table

ixpantia · Oct 28, 2024 · 65edeb0 · 65edeb0
1 parent ae0acb6
commit 65edeb0
Show file tree

Hide file tree

Showing 7 changed files with 56 additions and 6 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -38,12 +38,13 @@ fxhash = "0.2.1"
 toml = "0.8.19"
 nix = "0.29.0"
 ctrlc = { version = "3.4.5", features = ["termination"] }
-tokio-postgres = { version = "0.7.12", features = ["with-chrono-0_4"] }
+tokio-postgres = { version = "0.7.12", features = ["with-chrono-0_4", "with-uuid-1"] }
 deadpool-postgres = { version = "0.14.0", features = ["rt_tokio_1"] }
 bytes = "1.8.0"
 tokio-postgres-rustls = "0.13.0"
 rustls = "0.23.15"
 chrono = "0.4.38"
+uuid = { version = "1.11.0", features = ["v7"] }
 
 [dev-dependencies]
 rand = "0.8.5"
diff --git a/README.md b/README.md
@@ -206,6 +206,7 @@ the following SQL query:
 
 ```sql
 CREATE TABLE faucet_http_events (
+    request_uuid UUID,
     namespace TEXT,
     target TEXT,
     worker_route TEXT,

diff --git a/postgres.sql b/postgres.sql
@@ -1,4 +1,5 @@
 CREATE TABLE faucet_http_events (
+    request_uuid UUID,
     namespace TEXT,
     target TEXT,
     worker_route TEXT,
@@ -14,3 +15,6 @@ CREATE TABLE faucet_http_events (
     time TIMESTAMPTZ
 );
 
+CREATE INDEX faucet_http_events_request_uuid_idx 
+ON faucet_http_events USING BTREE (request_uuid);
+
diff --git a/src/server/logging.rs b/src/server/logging.rs
@@ -72,6 +72,7 @@ pub mod logger {
 
 #[derive(Clone, Copy)]
 pub struct StateData {
+    pub uuid: uuid::Uuid,
     pub ip: IpAddr,
     pub worker_route: Option<&'static str>,
     pub worker_id: usize,
@@ -85,11 +86,13 @@ trait StateLogData: Send + Sync + 'static {
 impl StateLogData for State {
     #[inline(always)]
     fn get_state_data(&self) -> StateData {
+        let uuid = self.uuid;
         let ip = self.remote_addr;
         let worker_id = self.client.config.worker_id;
         let worker_route = self.client.config.worker_route;
         let target = self.client.config.target;
         StateData {
+            uuid,
             ip,
             worker_id,
             worker_route,
@@ -255,6 +258,7 @@ mod tests {
         impl StateLogData for MockState {
             fn get_state_data(&self) -> StateData {
                 StateData {
+                    uuid: uuid::Uuid::now_v7(),
                     ip: IpAddr::V4([127, 0, 0, 1].into()),
                     target: "test",
                     worker_id: 1,
@@ -352,6 +356,7 @@ mod tests {
 
         let log_data = LogData {
             state_data: StateData {
+                uuid: uuid::Uuid::now_v7(),
                 target: "test",
                 ip: IpAddr::V4([127, 0, 0, 1].into()),
                 worker_route: None,

diff --git a/src/server/service.rs b/src/server/service.rs
@@ -5,22 +5,41 @@ use crate::{
     error::FaucetError,
     server::load_balancing::LoadBalancer,
 };
-use hyper::body::Incoming;
+use hyper::{body::Incoming, header::HeaderValue};
 
 use super::onion::{Layer, Service};
 
 #[derive(Clone)]
 pub(crate) struct State {
+    pub uuid: uuid::Uuid,
     pub remote_addr: IpAddr,
     pub client: Client,
 }
 
+impl State {
+    #[inline(always)]
+    fn new(remote_addr: IpAddr, client: Client) -> State {
+        let uuid = uuid::Uuid::now_v7();
+        State {
+            remote_addr,
+            client,
+            uuid,
+        }
+    }
+}
+
 #[derive(Clone)]
 pub struct AddStateService<S> {
     inner: S,
     load_balancer: LoadBalancer,
 }
 
+fn uuid_to_header_value(uuid: uuid::Uuid) -> HeaderValue {
+    let mut buffer = [0u8; uuid::fmt::Hyphenated::LENGTH];
+    HeaderValue::from_str(uuid.hyphenated().encode_lower(&mut buffer))
+        .expect("Unable to convert from uuid to header value, this is a bug")
+}
+
 impl<S, ReqBody> Service<hyper::Request<ReqBody>> for AddStateService<S>
 where
     ReqBody: hyper::body::Body + Send + Sync + 'static,
@@ -46,11 +65,17 @@ where
                 return Err(e);
             }
         };
+
         let client = self.load_balancer.get_client(remote_addr).await?;
-        req.extensions_mut().insert(State {
-            remote_addr,
-            client,
-        });
+
+        let state = State::new(remote_addr, client);
+
+        // Add the state's UUID to the request. `X-` headers are depracted
+        // https://www.rfc-editor.org/rfc/rfc6648
+        req.headers_mut()
+            .insert("Faucet-Request-Uuid", uuid_to_header_value(state.uuid));
+
+        req.extensions_mut().insert(state);
         self.inner.call(req, Some(remote_addr)).await
     }
 }

diff --git a/src/telemetry/mod.rs b/src/telemetry/mod.rs
@@ -57,6 +57,7 @@ impl TelemetryManager {
             let pool = pool.clone();
             tokio::task::spawn(async move {
                 let types = &[
+                    PgType::UUID,        // UUID
                     PgType::TEXT,        // Namespace
                     PgType::TEXT,        // Target
                     PgType::TEXT,        // Worker Route
@@ -106,6 +107,7 @@ impl TelemetryManager {
                             );
 
                             'write: for (timespamp, log_data) in logs_buffer.drain(..) {
+                                let uuid = &log_data.state_data.uuid;
                                 let target = &log_data.state_data.target;
                                 let worker_id = log_data.state_data.worker_id as i32;
                                 let worker_route = log_data.state_data.worker_route;
@@ -133,6 +135,7 @@ impl TelemetryManager {
                                 let copy_result = copy_in_writer
                                     .as_mut()
                                     .write(&[
+                                        uuid,
                                         &namespace,
                                         target,
                                         &worker_route,