Explicitly flush the file

Author: danielsn

libdd-profiling/src/exporter/file_exporter.rs

@@ -8,6 +8,7 @@
 //!
 //! This is primarily used for testing to validate the exact bytes sent over the wire.
 
+use anyhow::Context as _;
 use std::path::PathBuf;
 
 /// HTTP 200 OK response with no body
@@ -124,7 +125,9 @@ async fn run_dump_server_unix(
 ) -> anyhow::Result<()> {
     loop {
         let (stream, _) = listener.accept().await?;
-        handle_connection_async(stream, output_path.clone()).await;
+        if let Err(e) = handle_connection_async(stream, output_path.clone()).await {
+            eprintln!("[dump-server] Error handling connection: {:#}", e);
+        }
     }
 }
 
@@ -139,7 +142,9 @@ async fn run_dump_server_windows(
 
     // Handle first connection
     first_server.connect().await?;
-    handle_connection_async(first_server, output_path.clone()).await;
+    if let Err(e) = handle_connection_async(first_server, output_path.clone()).await {
+        eprintln!("[dump-server] Error handling connection: {:#}", e);
+    }
 
     // Handle subsequent connections
     loop {
@@ -152,7 +157,9 @@ async fn run_dump_server_windows(
         server.connect().await?;
 
         // Handle connection sequentially (this is just a debugging API)
-        handle_connection_async(server, output_path.clone()).await;
+        if let Err(e) = handle_connection_async(server, output_path.clone()).await {
+            eprintln!("[dump-server] Error handling connection: {:#}", e);
+        }
     }
 }
 
@@ -203,39 +210,41 @@ fn is_request_complete(
 }
 
 /// Read complete HTTP request from an async stream
-async fn read_http_request_async<R: tokio::io::AsyncReadExt + Unpin>(stream: &mut R) -> Vec<u8> {
+async fn read_http_request_async<R: tokio::io::AsyncReadExt + Unpin>(
+    stream: &mut R,
+) -> anyhow::Result<Vec<u8>> {
     let mut request_data = Vec::new();
     let mut buffer = [0u8; 8192];
     let mut content_length: Option<usize> = None;
     let mut headers_end_pos: Option<usize> = None;
 
     loop {
-        match stream.read(&mut buffer).await {
-            Ok(0) => break, // Connection closed
-            Ok(n) => {
-                request_data.extend_from_slice(&buffer[..n]);
-
-                // Look for end of headers if we haven't found it yet
-                if headers_end_pos.is_none() {
-                    if let Some(pos) = find_subsequence(&request_data, b"\r\n\r\n") {
-                        headers_end_pos = Some(pos + 4);
-                        content_length = parse_content_length(&request_data[..pos]);
-                    }
-                }
+        let n = stream
+            .read(&mut buffer)
+            .await
+            .context("Failed to read from connection")?;
 
-                // Check if we have the complete request
-                if is_request_complete(&request_data, headers_end_pos, content_length) {
-                    break;
-                }
-            }
-            Err(e) => {
-                eprintln!("[dump-server] Failed to read from connection: {}", e);
-                break;
+        if n == 0 {
+            break; // Connection closed
+        }
+
+        request_data.extend_from_slice(&buffer[..n]);
+
+        // Look for end of headers if we haven't found it yet
+        if headers_end_pos.is_none() {
+            if let Some(pos) = find_subsequence(&request_data, b"\r\n\r\n") {
+                headers_end_pos = Some(pos + 4);
+                content_length = parse_content_length(&request_data[..pos]);
             }
         }
+
+        // Check if we have the complete request
+        if is_request_complete(&request_data, headers_end_pos, content_length) {
+            break;
+        }
     }
 
-    request_data
+    Ok(request_data)
 }
 
 /// Decode chunked transfer encoding
@@ -245,49 +254,51 @@ fn decode_chunked_body(chunked_data: &[u8]) -> Vec<u8> {
 
     while pos < chunked_data.len() {
         // Find the end of the chunk size line (\r\n)
-        if let Some(line_end) = find_subsequence(&chunked_data[pos..], b"\r\n") {
-            // Parse the chunk size (hex)
-            if let Ok(size_str) = std::str::from_utf8(&chunked_data[pos..pos + line_end]) {
-                if let Ok(chunk_size) = usize::from_str_radix(size_str.trim(), 16) {
-                    if chunk_size == 0 {
-                        // End of chunks
-                        break;
-                    }
+        let Some(line_end) = find_subsequence(&chunked_data[pos..], b"\r\n") else {
+            break;
+        };
 
-                    // Move past the size line and \r\n
-                    pos += line_end + 2;
+        // Parse the chunk size (hex)
+        let Ok(size_str) = std::str::from_utf8(&chunked_data[pos..pos + line_end]) else {
+            break;
+        };
 
-                    // Read the chunk data
-                    if pos + chunk_size <= chunked_data.len() {
-                        result.extend_from_slice(&chunked_data[pos..pos + chunk_size]);
-                        pos += chunk_size;
+        let Ok(chunk_size) = usize::from_str_radix(size_str.trim(), 16) else {
+            break;
+        };
 
-                        // Skip the trailing \r\n after the chunk
-                        if pos + 2 <= chunked_data.len() && &chunked_data[pos..pos + 2] == b"\r\n" {
-                            pos += 2;
-                        }
-                    } else {
-                        break;
-                    }
-                } else {
-                    break;
-                }
-            } else {
-                break;
-            }
-        } else {
+        if chunk_size == 0 {
+            break; // End of chunks
+        }
+
+        // Move past the size line and \r\n
+        pos += line_end + 2;
+
+        // Read the chunk data
+        if pos + chunk_size > chunked_data.len() {
             break;
         }
+
+        result.extend_from_slice(&chunked_data[pos..pos + chunk_size]);
+        pos += chunk_size;
+
+        // Skip the trailing \r\n after the chunk
+        if pos + 2 <= chunked_data.len() && &chunked_data[pos..pos + 2] == b"\r\n" {
+            pos += 2;
+        }
     }
 
     result
 }
 
 /// Write request data to file if non-empty (async version)
 /// Decodes chunked transfer encoding if present
-async fn write_request_to_file_async(output_path: &PathBuf, request_data: &[u8]) {
+async fn write_request_to_file_async(
+    output_path: &PathBuf,
+    request_data: &[u8],
+) -> anyhow::Result<()> {
     if request_data.is_empty() {
-        return;
+        return Ok(());
     }
 
     // Check if this is a chunked request and decode it
@@ -296,13 +307,11 @@ async fn write_request_to_file_async(output_path: &PathBuf, request_data: &[u8])
         let body = &request_data[headers_end + 4..];
 
         // Check for transfer-encoding: chunked
-        let is_chunked = if let Ok(headers_str) = std::str::from_utf8(headers) {
+        let is_chunked = std::str::from_utf8(headers).is_ok_and(|headers_str| {
             headers_str
                 .to_lowercase()
                 .contains("transfer-encoding: chunked")
-        } else {
-            false
-        };
+        });
 
         if is_chunked {
             // Decode the chunked body and reconstruct the request with Content-Length
@@ -335,23 +344,37 @@ async fn write_request_to_file_async(output_path: &PathBuf, request_data: &[u8])
         request_data.to_vec()
     };
 
-    if let Err(e) = tokio::fs::write(output_path, data_to_write).await {
-        eprintln!(
-            "[dump-server] Failed to write request dump to {:?}: {}",
-            output_path, e
-        );
-    }
+    // Write to file and explicitly sync to ensure data is on disk before responding
+    use tokio::io::AsyncWriteExt;
+
+    let mut file = tokio::fs::File::create(output_path)
+        .await
+        .context("Failed to create dump file")?;
+
+    file.write_all(&data_to_write)
+        .await
+        .context("Failed to write request dump")?;
+
+    // Sync to ensure data is persisted to disk before sending HTTP response
+    file.sync_all()
+        .await
+        .context("Failed to sync request dump to disk")?;
+
+    Ok(())
 }
 
 /// Handle a connection: read HTTP request, write to file, send response
-async fn handle_connection_async<S>(mut stream: S, output_path: PathBuf)
+async fn handle_connection_async<S>(mut stream: S, output_path: PathBuf) -> anyhow::Result<()>
 where
     S: tokio::io::AsyncReadExt + tokio::io::AsyncWriteExt + Unpin,
 {
-    let request_data = read_http_request_async(&mut stream).await;
-    write_request_to_file_async(&output_path, &request_data).await;
+    let request_data = read_http_request_async(&mut stream).await?;
+    write_request_to_file_async(&output_path, &request_data).await?;
 
-    if let Err(e) = stream.write_all(HTTP_200_RESPONSE).await {
-        eprintln!("[dump-server] Failed to send HTTP response: {}", e);
-    }
+    stream
+        .write_all(HTTP_200_RESPONSE)
+        .await
+        .context("Failed to send HTTP response")?;
+
+    Ok(())
 }

libdd-profiling/tests/file_exporter_test.rs

@@ -18,10 +18,12 @@ fn create_file_exporter(
 
     // Create a unique temp file path
     let temp_dir = std::env::temp_dir();
+    let random_id: u64 = rand::random();
     let file_path = temp_dir.join(format!(
-        "libdd_test_{}_{}.http",
+        "libdd_test_{}_{}_{:x}.http",
         std::process::id(),
-        chrono::Utc::now().timestamp_nanos_opt().unwrap_or(0)
+        chrono::Utc::now().timestamp_nanos_opt().unwrap_or(0),
+        random_id
     ));
 
     let mut endpoint = config::file(file_path.to_string_lossy().as_ref())?;
@@ -73,6 +75,8 @@ mod tests {
         exporter.send(request, None).expect("send to succeed");
 
         // Read the dump file (wait a moment for it to be written)
+        // The file is synced before the 200 response, but we still need a small delay
+        // to ensure the background thread's runtime has fully completed the async operation
         std::thread::sleep(std::time::Duration::from_millis(200));
         let request_bytes = std::fs::read(&file_path).expect("read dump file");
 
@@ -193,6 +197,8 @@ mod tests {
         exporter.send(request, None).expect("send to succeed");
 
         // Read the dump file (wait a moment for it to be written)
+        // The file is synced before the 200 response, but we still need a small delay
+        // to ensure the background thread's runtime has fully completed the async operation
         std::thread::sleep(std::time::Duration::from_millis(200));
         let request_bytes = std::fs::read(&file_path).expect("read dump file");
 
@@ -249,6 +255,8 @@ mod tests {
         exporter.send(request, None).expect("send to succeed");
 
         // Read the dump file (wait a moment for it to be written)
+        // The file is synced before the 200 response, but we still need a small delay
+        // to ensure the background thread's runtime has fully completed the async operation
         std::thread::sleep(std::time::Duration::from_millis(200));
         let request_bytes = std::fs::read(&file_path).expect("read dump file");
 
@@ -312,6 +320,8 @@ mod tests {
         exporter.send(request, None).expect("send to succeed");
 
         // Read the dump file (wait a moment for it to be written)
+        // The file is synced before the 200 response, but we still need a small delay
+        // to ensure the background thread's runtime has fully completed the async operation
         std::thread::sleep(std::time::Duration::from_millis(200));
         let request_bytes = std::fs::read(&file_path).expect("read dump file");
 
@@ -359,6 +369,8 @@ mod tests {
         exporter.send(request, None).expect("send to succeed");
 
         // Read the dump file (wait a moment for it to be written)
+        // The file is synced before the 200 response, but we still need a small delay
+        // to ensure the background thread's runtime has fully completed the async operation
         std::thread::sleep(std::time::Duration::from_millis(200));
         let request_bytes = std::fs::read(&file_path).expect("read dump file");