Minimise network traffic with ETags

DiamondLightSource · Dec 13, 2023 · 865a3f1 · 865a3f1
1 parent dfa6cea
commit 865a3f1
Show file tree

Hide file tree

Showing 4 changed files with 88 additions and 6 deletions.
diff --git a/bundler/Cargo.lock b/bundler/Cargo.lock
diff --git a/bundler/Cargo.toml b/bundler/Cargo.toml
@@ -6,9 +6,11 @@ edition = "2021"
 [dependencies]
 anyhow = "1.0.75"
 axum = { version = "0.7.2" }
+axum-extra = { version = "0.9.0", features = ["typed-header"] }
 clap = { version = "4.4.11", features = ["derive", "env"] }
 dotenvy = { version = "0.15.7" }
 flate2 = { version = "1.0.28" }
+headers = { version = "0.4.0" }
 humantime = { version = "2.1.0" }
 serde = { version = "1.0.193", features = ["derive"] }
 serde_json = { version = "1.0.108" }

diff --git a/bundler/src/bundle.rs b/bundler/src/bundle.rs
@@ -79,6 +79,10 @@ where
         Ok(Self::new(metadata, proposals, sessions))
     }
 
+    pub fn revision(&self) -> &str {
+        &self.manifest.revision
+    }
+
     pub fn to_tar_gz(&self) -> Result<Vec<u8>, anyhow::Error> {
         let mut bundle_builder = tar::Builder::new(GzEncoder::new(Vec::new(), Compression::best()));
 

diff --git a/bundler/src/main.rs b/bundler/src/main.rs
@@ -3,15 +3,24 @@ mod bundle;
 mod permissionables;
 
 use crate::bundle::{Bundle, NoMetadata};
-use axum::{body::Bytes, extract::State, routing::get, serve, Router};
+use axum::{
+    body::Bytes,
+    extract::State,
+    http::{HeaderMap, StatusCode},
+    response::IntoResponse,
+    routing::get,
+    serve, Router,
+};
+use axum_extra::TypedHeader;
 use clap::Parser;
+use headers::{ETag, HeaderMapExt, IfNoneMatch};
 use serde::Serialize;
 use sqlx::{mysql::MySqlPoolOptions, MySqlPool};
 use std::{
     hash::Hash,
-    marker::PhantomData,
     net::{Ipv4Addr, SocketAddr, SocketAddrV4},
     ops::Add,
+    str::FromStr,
     sync::Arc,
     time::Duration,
 };
@@ -28,8 +37,8 @@ struct BundleFile<Metadata>
 where
     Metadata: Serialize,
 {
+    bundle: Bundle<Metadata>,
     file: Bytes,
-    _metadata: PhantomData<Metadata>,
 }
 
 impl<Metadata> TryFrom<Bundle<Metadata>> for BundleFile<Metadata>
@@ -41,7 +50,7 @@ where
     fn try_from(bundle: Bundle<Metadata>) -> Result<Self, Self::Error> {
         Ok(Self {
             file: bundle.to_tar_gz()?.into(),
-            _metadata: PhantomData,
+            bundle,
         })
     }
 }
@@ -115,6 +124,25 @@ async fn update_bundle(
     }
 }
 
-async fn bundle_endpoint(State(current_bundle): State<CurrentBundle>) -> Bytes {
-    current_bundle.as_ref().read().await.file.clone()
+async fn bundle_endpoint(
+    State(current_bundle): State<CurrentBundle>,
+    if_none_match: Option<TypedHeader<IfNoneMatch>>,
+) -> impl IntoResponse {
+    let etag = ETag::from_str(&format!(
+        r#""{}""#,
+        current_bundle.as_ref().read().await.bundle.revision()
+    ))
+    .unwrap();
+    let mut headers = HeaderMap::new();
+    headers.typed_insert(etag.clone());
+    match if_none_match {
+        Some(TypedHeader(if_none_match)) if !if_none_match.precondition_passes(&etag) => {
+            (StatusCode::NOT_MODIFIED, headers, Bytes::new())
+        }
+        _ => (
+            StatusCode::OK,
+            headers,
+            current_bundle.as_ref().read().await.file.clone(),
+        ),
+    }
 }