From 2d5a574520f45732641725e61d75bf07d8e6b18e Mon Sep 17 00:00:00 2001
From: Blaise Bruer <github.blaise@allada.com>
Date: Mon, 9 Sep 2024 11:55:31 -0500
Subject: [PATCH] ByteStreamServer now responds with no-data-received instead
 of NotFound

If a client tries to upload a file, but the first payload fails, the
client may call query_write_status(), which we now return "0" (no data
received yet) instead of NotFound to increase compatibility with
clients.
---
 nativelink-service/src/bytestream_server.rs      |  9 ++++++++-
 .../tests/bytestream_server_test.rs              | 16 +++++++++-------
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/nativelink-service/src/bytestream_server.rs b/nativelink-service/src/bytestream_server.rs
index 71b7eff7a..bb5ab6c91 100644
--- a/nativelink-service/src/bytestream_server.rs
+++ b/nativelink-service/src/bytestream_server.rs
@@ -556,7 +556,14 @@ impl ByteStreamServer {
 
         let has_fut = store_clone.has(digest);
         let Some(item_size) = has_fut.await.err_tip(|| "Failed to call .has() on store")? else {
-            return Err(make_err!(Code::NotFound, "{}", "not found"));
+            // We lie here and say that the stream needs to start over, even though
+            // it was never started. This can happen when the client disconnects
+            // before sending the first payload, but the client thinks it did send
+            // the payload.
+            return Ok(Response::new(QueryWriteStatusResponse {
+                committed_size: 0,
+                complete: false,
+            }));
         };
         Ok(Response::new(QueryWriteStatusResponse {
             committed_size: item_size as i64,
diff --git a/nativelink-service/tests/bytestream_server_test.rs b/nativelink-service/tests/bytestream_server_test.rs
index 2b4d453ae..9d1124d29 100644
--- a/nativelink-service/tests/bytestream_server_test.rs
+++ b/nativelink-service/tests/bytestream_server_test.rs
@@ -866,19 +866,21 @@ pub async fn test_query_write_status_smoke_test() -> Result<(), Box<dyn std::err
     let resource_name = make_resource_name(raw_data.len());
 
     {
+        // If the write does not exist it should respond with a 0 size.
+        // This is because the client may have tried to write, but the
+        // connection closed before the payload had been received.
         let response = bs_server
             .query_write_status(Request::new(QueryWriteStatusRequest {
                 resource_name: resource_name.clone(),
             }))
             .await;
-        assert_eq!(response.is_err(), true);
-        let expected_err = make_err!(
-            Code::NotFound,
-            "{}{}",
-            "status: NotFound, message: \"not found : Failed on query_write_status() ",
-            "command\", details: [], metadata: MetadataMap { headers: {} }"
+        assert_eq!(
+            response.unwrap().into_inner(),
+            QueryWriteStatusResponse {
+                committed_size: 0,
+                complete: false,
+            }
         );
-        assert_eq!(Into::<Error>::into(response.unwrap_err()), expected_err);
     }
 
     // Setup stream.