Update the logic for checking if a blob exist before uploading (#2503)

* Update the logic for checking if a blob exist before uploading * remove unused ref * update the status code check
microsoft · Oct 12, 2022 · 9e5006f · 9e5006f
1 parent 0fab487
commit 9e5006f
Showing 1 changed file with 32 additions and 37 deletions.
diff --git a/src/agent/onefuzz/src/uploader.rs b/src/agent/onefuzz/src/uploader.rs
@@ -7,7 +7,7 @@ use anyhow::{Context, Result};
 use futures::stream::TryStreamExt;
 use reqwest::{Body, Client, Response, StatusCode, Url};
 use reqwest_retry::{
-    send_retry_reqwest_default, RetryCheck, SendRetry, DEFAULT_RETRY_PERIOD, MAX_RETRY_ATTEMPTS,
+    send_retry_reqwest, RetryCheck, SendRetry, DEFAULT_RETRY_PERIOD, MAX_RETRY_ATTEMPTS,
 };
 use serde::Serialize;
 use tokio::{fs, io};
@@ -42,44 +42,39 @@ impl BlobUploader {
             url
         };
 
-        // Check if the file already exists before uploading
-        if let Ok(head) = self
-            .client
-            .head(url.clone())
-            .send_retry(
-                |code| match code {
-                    StatusCode::NOT_FOUND => RetryCheck::Fail,
-                    _ => RetryCheck::Retry,
-                },
-                DEFAULT_RETRY_PERIOD,
-                MAX_RETRY_ATTEMPTS,
-            )
-            .await
-        {
-            if head.status() == StatusCode::OK {
-                return Ok(head);
-            }
-        }
-
         let content_length = format!("{}", file_len);
 
-        let resp = send_retry_reqwest_default(|| {
-            let file = fs::File::from_std(std::fs::File::open(file_path)?);
-            let reader = io::BufReader::new(file);
-            let codec = codec::BytesCodec::new();
-            let file_stream = codec::FramedRead::new(reader, codec)
-                .map_ok(bytes::BytesMut::freeze)
-                .into_stream();
-
-            let request_builder = self
-                .client
-                .put(url.clone())
-                .header("Content-Length", &content_length)
-                .header("x-ms-blob-type", "BlockBlob")
-                .body(Body::wrap_stream(file_stream));
-
-            Ok(request_builder)
-        })
+        let resp = send_retry_reqwest(
+            || {
+                let file = fs::File::from_std(std::fs::File::open(file_path)?);
+                let reader = io::BufReader::new(file);
+                let codec = codec::BytesCodec::new();
+                let file_stream = codec::FramedRead::new(reader, codec)
+                    .map_ok(bytes::BytesMut::freeze)
+                    .into_stream();
+
+                let request_builder = self
+                    .client
+                    .put(url.clone())
+                    // https://learn.microsoft.com/en-us/rest/api/storageservices/put-blob-from-url#request-headers
+                    .header("Content-Length", &content_length)
+                    .header("x-ms-blob-type", "BlockBlob")
+                    // upload only if the the destination blob does not exist
+                    .header("If-None-Match", "*")
+                    .body(Body::wrap_stream(file_stream));
+
+                Ok(request_builder)
+            },
+            |status| {
+                if status == StatusCode::PRECONDITION_FAILED || status == StatusCode::CONFLICT {
+                    RetryCheck::Succeed
+                } else {
+                    RetryCheck::Retry
+                }
+            },
+            DEFAULT_RETRY_PERIOD,
+            MAX_RETRY_ATTEMPTS,
+        )
         .await
         .context("BlobUploader.upload")?;