fix(export): send HTTP headers before serializing large exports

aw-server/src/endpoints/util.rs

@@ -1,4 +1,6 @@
-use std::io::{Cursor, Seek, SeekFrom};
+use std::fs::File;
+use std::io::{copy, pipe, Cursor, PipeReader, PipeWriter, Seek, SeekFrom};
+use std::thread;
 
 use rocket::http::ContentType;
 use rocket::http::Header;
@@ -38,43 +40,121 @@ impl<'r> Responder<'r, 'static> for HttpErrorJson {
 }
 
 pub struct BucketsExportRocket {
-    file: std::fs::File,
+    datastore: aw_datastore::Datastore,
+    bucket_id: Option<String>,
     filename: String,
 }
 
+fn export_filename(
+    datastore: &aw_datastore::Datastore,
+    bucket_id: Option<&str>,
+) -> Result<String, HttpErrorJson> {
+    let name = match bucket_id {
+        Some(id) => {
+            datastore.get_bucket(id)?;
+            Some(id.to_owned())
+        }
+        None => {
+            let buckets = datastore.get_buckets()?;
+            (buckets.len() == 1).then(|| buckets.into_keys().next().unwrap())
+        }
+    };
+    Ok(match name {
+        Some(id) => format!("attachment; filename=aw-bucket-export_{id}.json"),
+        None => "attachment; filename=aw-buckets-export.json".into(),
+    })
+}
+
+#[cfg(not(any(unix, windows)))]
+compile_error!("export streaming requires unix or windows anonymous pipes");
+
+fn pipe_reader_to_file(reader: PipeReader) -> File {
+    #[cfg(unix)]
+    {
+        File::from(std::os::fd::OwnedFd::from(reader))
+    }
+    #[cfg(windows)]
+    {
+        File::from(std::os::windows::io::OwnedHandle::from(reader))
+    }
+}
+
+fn pipe_writer_to_file(writer: PipeWriter) -> File {
+    #[cfg(unix)]
+    {
+        File::from(std::os::fd::OwnedFd::from(writer))
+    }
+    #[cfg(windows)]
+    {
+        File::from(std::os::windows::io::OwnedHandle::from(writer))
+    }
+}
+
+/// Serialize on the datastore worker into a private tempfile, then copy to
+/// the client pipe from this thread. The worker stays disk-paced; a slow
+/// or dropped download must not stall heartbeats (see `ServerState`).
+fn spawn_export_stream(
+    datastore: aw_datastore::Datastore,
+    bucket_id: Option<String>,
+    writer: PipeWriter,
+) {
+    thread::spawn(move || {
+        let staging = match tempfile::tempfile() {
+            Ok(file) => file,
+            Err(err) => {
+                error!("Failed to create export staging file: {err}");
+                return;
+            }
+        };
+        let mut staging = match datastore.export_to_file(bucket_id.as_deref(), staging) {
+            Ok((file, _)) => file,
+            Err(err) => {
+                error!("Export stream failed: {err:?}");
+                return;
+            }
+        };
+        if let Err(err) = staging.seek(SeekFrom::Start(0)) {
+            error!("Failed to rewind export staging file: {err}");
+            return;
+        }
+        let mut writer = pipe_writer_to_file(writer);
+        if let Err(err) = copy(&mut staging, &mut writer) {
+            error!("Export stream copy failed: {err}");
+        }
+    });
+}
+
 impl BucketsExportRocket {
     pub fn new(
         datastore: &aw_datastore::Datastore,
         bucket_id: Option<&str>,
     ) -> Result<Self, HttpErrorJson> {
-        let io_error = |err: std::io::Error| {
-            error!("Failed to prepare export file: {err}");
-            HttpErrorJson::new(
-                Status::InternalServerError,
-                "Failed to prepare export file".into(),
-            )
-        };
-        // tempfile creates a private file and removes it when the response is
-        // dropped. Spooling preserves HTTP errors even if serialization or disk
-        // writes fail, while keeping event buffering bounded.
-        let file = tempfile::tempfile().map_err(io_error)?;
-        let (mut file, name) = datastore.export_to_file(bucket_id, file)?;
-        file.seek(SeekFrom::Start(0)).map_err(io_error)?;
-        let filename = match name {
-            Some(id) => format!("attachment; filename=aw-bucket-export_{id}.json"),
-            None => "attachment; filename=aw-buckets-export.json".into(),
-        };
-        Ok(Self { file, filename })
+        // Resolve the download name and 404 missing buckets before the
+        // response is built. Serialization itself runs after headers so a
+        // slow export does not look like a hung connection.
+        let filename = export_filename(datastore, bucket_id)?;
+        Ok(Self {
+            datastore: datastore.clone(),
+            bucket_id: bucket_id.map(str::to_owned),
+            filename,
+        })
     }
 }
 
 impl<'r> Responder<'r, 'static> for BucketsExportRocket {
     fn respond_to(self, _: &Request) -> response::Result<'static> {
+        let (reader, writer) = pipe().map_err(|err| {
+            error!("Failed to open export pipe: {err}");
+            Status::InternalServerError
+        })?;
+        spawn_export_stream(self.datastore, self.bucket_id, writer);
         Response::build()
             .status(Status::Ok)
             .header(Header::new("Content-Disposition", self.filename))
             .header(ContentType::JSON)
-            .streamed_body(rocket::tokio::fs::File::from_std(self.file))
+            .streamed_body(rocket::tokio::fs::File::from_std(pipe_reader_to_file(
+                reader,
+            )))
             .ok()
     }
 }

aw-server/tests/api.rs

@@ -71,6 +71,72 @@ mod api_tests {
         assert!(body["message"].as_str().unwrap().contains("does not exist"));
     }
 
+    #[test]
+    fn export_all_empty_uses_plural_filename_and_opens_before_body() {
+        let server = setup_testserver();
+        let client = Client::untracked(server).unwrap();
+        let response = client
+            .get("/api/0/export")
+            .header(Header::new("Host", "127.0.0.1:5600"))
+            .dispatch();
+        assert_eq!(response.status(), Status::Ok);
+        assert_eq!(response.content_type(), Some(ContentType::JSON));
+        assert_eq!(
+            response.headers().get_one("Content-Disposition"),
+            Some("attachment; filename=aw-buckets-export.json")
+        );
+        let body: Value = serde_json::from_str(&response.into_string().unwrap()).unwrap();
+        assert_eq!(body["buckets"], json!({}));
+    }
+
+    #[test]
+    fn unread_export_body_does_not_block_datastore() {
+        let server = setup_testserver();
+        let datastore = server
+            .state::<endpoints::ServerState>()
+            .unwrap()
+            .datastore
+            .clone();
+        let bucket: Bucket = serde_json::from_value(json!({
+            "id": "big", "type": "test", "client": "test", "hostname": "test"
+        }))
+        .unwrap();
+        datastore.create_bucket(&bucket).unwrap();
+        let mut event = aw_models::Event::default();
+        event
+            .data
+            .insert("blob".into(), json!("x".repeat(128 * 1024)));
+        datastore.insert_events("big", &[event]).unwrap();
+
+        let client = Client::untracked(server).unwrap();
+        let response = client
+            .get("/api/0/buckets/big/export")
+            .header(Header::new("Host", "127.0.0.1:5600"))
+            .dispatch();
+        assert_eq!(response.status(), Status::Ok);
+
+        // Let Command::Export start. A client-paced pipe would fill here and
+        // stall the worker; staging to a tempfile must not.
+        std::thread::sleep(std::time::Duration::from_millis(100));
+        let (tx, rx) = std::sync::mpsc::channel();
+        let ds = datastore.clone();
+        std::thread::spawn(move || {
+            let _ = tx.send(ds.get_buckets());
+        });
+        rx.recv_timeout(std::time::Duration::from_secs(3))
+            .expect("datastore worker blocked by unread export body")
+            .unwrap();
+
+        let body: Value = serde_json::from_str(&response.into_string().unwrap()).unwrap();
+        assert_eq!(
+            body["buckets"]["big"]["events"][0]["data"]["blob"]
+                .as_str()
+                .unwrap()
+                .len(),
+            128 * 1024
+        );
+    }
+
     #[test]
     fn test_bucket() {
         let server = setup_testserver();