Resilient connection (#153)

microsoft · Oct 28, 2020 · db85341 · db85341
1 parent 1d2fb99
commit db85341
Show file tree

Hide file tree

Showing 19 changed files with 532 additions and 38 deletions.
diff --git a/src/agent/Cargo.lock b/src/agent/Cargo.lock
diff --git a/src/agent/Cargo.toml b/src/agent/Cargo.toml
@@ -6,6 +6,7 @@ members = [
     "onefuzz",
     "onefuzz-agent",
     "onefuzz-supervisor",
+    "reqwest-retry",
     "storage-queue",
     "win-util",
 ]

diff --git a/src/agent/onefuzz-agent/Cargo.toml b/src/agent/onefuzz-agent/Cargo.toml
@@ -30,6 +30,7 @@ url = { version = "2.1", features = ["serde"] }
 uuid = { version = "0.8", features = ["serde", "v4"] }
 onefuzz = { path = "../onefuzz" }
 storage-queue = { path = "../storage-queue" }
+reqwest-retry = { path = "../reqwest-retry" }
 
 [dev-dependencies]
 tempfile = "3.1"
diff --git a/src/agent/onefuzz-agent/src/tasks/merge/generic.rs b/src/agent/onefuzz-agent/src/tasks/merge/generic.rs
@@ -8,6 +8,7 @@ use onefuzz::{
     syncdir::SyncedDir,
 };
 use reqwest::Url;
+use reqwest_retry::SendRetry;
 use serde::Deserialize;
 use std::{
     collections::HashMap,
@@ -117,7 +118,7 @@ async fn try_delete_blob(input_url: Url) -> Result<()> {
     let http_client = reqwest::Client::new();
     match http_client
         .delete(input_url)
-        .send()
+        .send_retry_default()
         .await?
         .error_for_status_with_body()
         .await

diff --git a/src/agent/onefuzz-agent/src/tasks/merge/libfuzzer_merge.rs b/src/agent/onefuzz-agent/src/tasks/merge/libfuzzer_merge.rs
@@ -10,6 +10,7 @@ use onefuzz::{
     syncdir::SyncedDir,
 };
 use reqwest::Url;
+use reqwest_retry::SendRetry;
 use serde::Deserialize;
 use std::{
     collections::HashMap,
@@ -116,7 +117,7 @@ async fn try_delete_blob(input_url: Url) -> Result<()> {
     let http_client = reqwest::Client::new();
     match http_client
         .delete(input_url)
-        .send()
+        .send_retry_default()
         .await?
         .error_for_status_with_body()
         .await

diff --git a/src/agent/onefuzz-agent/src/tasks/report/crash_report.rs b/src/agent/onefuzz-agent/src/tasks/report/crash_report.rs
@@ -8,7 +8,9 @@ use onefuzz::{
     syncdir::SyncedDir,
     telemetry::Event::{new_report, new_unable_to_reproduce, new_unique_report},
 };
+
 use reqwest::StatusCode;
+use reqwest_retry::SendRetry;
 use serde::{Deserialize, Serialize};
 use std::path::PathBuf;
 use uuid::Uuid;
@@ -65,7 +67,7 @@ async fn upload_deduped(report: &CrashReport, container: &BlobContainerUrl) -> R
         .json(report)
         // Conditional PUT, only if-not-exists.
         .header("If-None-Match", "*")
-        .send()
+        .send_retry_default()
         .await?;
     if result.status() != StatusCode::NOT_MODIFIED {
         event!(new_unique_report;);
@@ -77,15 +79,15 @@ async fn upload_report(report: &CrashReport, container: &BlobContainerUrl) -> Re
     event!(new_report;);
     let blob = BlobClient::new();
     let url = container.blob(report.blob_name()).url();
-    blob.put(url).json(report).send().await?;
+    blob.put(url).json(report).send_retry_default().await?;
     Ok(())
 }
 
 async fn upload_no_repro(report: &NoCrash, container: &BlobContainerUrl) -> Result<()> {
     event!(new_unable_to_reproduce;);
     let blob = BlobClient::new();
     let url = container.blob(report.blob_name()).url();
-    blob.put(url).json(report).send().await?;
+    blob.put(url).json(report).send_retry_default().await?;
     Ok(())
 }
 

diff --git a/src/agent/onefuzz-supervisor/Cargo.toml b/src/agent/onefuzz-supervisor/Cargo.toml
@@ -24,3 +24,4 @@ tokio = { version = "0.2.13", features = ["full"] }
 url = { version = "2.1.1", features = ["serde"] }
 uuid = { version = "0.8.1", features = ["serde", "v4"] }
 clap = "2.33"
+reqwest-retry = { path = "../reqwest-retry" }
diff --git a/src/agent/onefuzz-supervisor/src/auth.rs b/src/agent/onefuzz-supervisor/src/auth.rs
@@ -5,6 +5,7 @@ use std::fmt;
 
 use anyhow::Result;
 use onefuzz::http::ResponseExt;
+use reqwest_retry::SendRetry;
 use url::Url;
 use uuid::Uuid;
 
@@ -112,7 +113,7 @@ impl ClientCredentials {
             .post(url)
             .header("Content-Length", "0")
             .form(&self.form_data())
-            .send()
+            .send_retry_default()
             .await?
             .error_for_status_with_body()
             .await?;
@@ -183,7 +184,7 @@ impl ManagedIdentityCredentials {
         let response = reqwest::Client::new()
             .get(self.url())
             .header("Metadata", "true")
-            .send()
+            .send_retry_default()
             .await?
             .error_for_status_with_body()
             .await?;

diff --git a/src/agent/onefuzz-supervisor/src/config.rs b/src/agent/onefuzz-supervisor/src/config.rs
@@ -4,6 +4,7 @@
 use anyhow::Result;
 use onefuzz::{http::ResponseExt, jitter::delay_with_jitter};
 use reqwest::StatusCode;
+use reqwest_retry::SendRetry;
 use std::{
     path::{Path, PathBuf},
     time::{Duration, Instant},
@@ -159,7 +160,7 @@ impl Registration {
                 .header("Content-Length", "0")
                 .bearer_auth(token.secret().expose_ref())
                 .body("")
-                .send()
+                .send_retry_default()
                 .await?
                 .error_for_status();
 
@@ -219,7 +220,7 @@ impl Registration {
         let response = reqwest::Client::new()
             .get(url)
             .bearer_auth(token.secret().expose_ref())
-            .send()
+            .send_retry_default()
             .await?
             .error_for_status_with_body()
             .await?;

diff --git a/src/agent/onefuzz/Cargo.toml b/src/agent/onefuzz/Cargo.toml
@@ -37,6 +37,7 @@ strum = "0.19"
 strum_macros = "0.19"
 tempfile = "3.1"
 process_control = "2.0"
+reqwest-retry = { path = "../reqwest-retry"}
 
 [target.'cfg(target_os = "windows")'.dependencies]
 winreg = "0.7"

diff --git a/src/agent/onefuzz/src/blob/client.rs b/src/agent/onefuzz/src/blob/client.rs
@@ -6,6 +6,7 @@ use std::path::{Path, PathBuf};
 use anyhow::Result;
 use futures::stream::TryStreamExt;
 use reqwest::{Body, RequestBuilder, Response, Url};
+use reqwest_retry::SendRetry;
 use serde::Serialize;
 use tokio::{fs, io};
 use tokio_util::codec;
@@ -31,7 +32,12 @@ impl BlobClient {
     pub async fn get(&self, url: &Url) -> Result<Response> {
         let url = url.clone();
 
-        let r = self.client.get(url).send().await?.error_for_status()?;
+        let r = self
+            .client
+            .get(url)
+            .send_retry_default()
+            .await?
+            .error_for_status()?;
 
         Ok(r)
     }
@@ -62,7 +68,7 @@ impl BlobClient {
             .put(url)
             .header("x-ms-blob-type", "BlockBlob")
             .body(data)
-            .send()
+            .send_retry_default()
             .await?;
 
         Ok(r)
@@ -77,7 +83,7 @@ impl BlobClient {
             .put(url)
             .header("x-ms-blob-type", "BlockBlob")
             .json(&item)
-            .send()
+            .send_retry_default()
             .await?;
 
         Ok(r)
@@ -103,7 +109,7 @@ impl BlobClient {
             .header("Content-Length", &content_length)
             .header("x-ms-blob-type", "BlockBlob")
             .body(body)
-            .send()
+            .send_retry_default()
             .await?;
 
         Ok(resp)

diff --git a/src/agent/onefuzz/src/machine_id.rs b/src/agent/onefuzz/src/machine_id.rs
@@ -3,6 +3,7 @@
 
 use crate::fs::{onefuzz_etc, write_file};
 use anyhow::Result;
+use reqwest_retry::SendRetry;
 use std::time::Duration;
 use tokio::fs;
 use uuid::Uuid;
@@ -27,7 +28,7 @@ pub async fn get_ims_id() -> Result<Uuid> {
                 .get(IMS_ID_URL)
                 .timeout(Duration::from_millis(500))
                 .header("Metadata", "true")
-                .send()
+                .send_retry_default()
                 .await?;
             let body = resp.text().await?;
             write_file(path, &body).await?;
@@ -48,7 +49,7 @@ pub async fn get_machine_name() -> Result<String> {
                 .get(VM_NAME_URL)
                 .timeout(Duration::from_millis(500))
                 .header("Metadata", "true")
-                .send()
+                .send_retry_default()
                 .await?;
             let body = resp.text().await?;
             write_file(path, &body).await?;
@@ -68,7 +69,7 @@ pub async fn get_scaleset_name() -> Result<String> {
                 .get(VM_SCALESET_NAME)
                 .timeout(Duration::from_millis(500))
                 .header("Metadata", "true")
-                .send()
+                .send_retry_default()
                 .await?;
             let body = resp.text().await?;
             write_file(path, &body).await?;

diff --git a/src/agent/onefuzz/src/uploader.rs b/src/agent/onefuzz/src/uploader.rs
@@ -6,6 +6,7 @@ use std::path::Path;
 use anyhow::Result;
 use futures::stream::TryStreamExt;
 use reqwest as r;
+use reqwest_retry::{send_retry_reqwest_default, SendRetry};
 use serde::Serialize;
 use tokio::{fs, io};
 use tokio_util::codec;
@@ -31,13 +32,6 @@ impl BlobUploader {
         let metadata = fs::metadata(file_path).await?;
         let file_len = metadata.len();
 
-        let file = fs::File::open(file_path).await?;
-        let reader = io::BufReader::new(file);
-        let codec = codec::BytesCodec::new();
-        let file_stream = codec::FramedRead::new(reader, codec).map_ok(bytes::BytesMut::freeze);
-
-        let body = r::Body::wrap_stream(file_stream);
-
         let url = {
             let url_path = self.url.path();
             let blob_path = format!("{}/{}", url_path, file_name);
@@ -47,21 +41,29 @@ impl BlobUploader {
         };
 
         // Check if the file already exists before uploading
-        let head = self.client.head(url.clone()).send().await?;
+        let head = self.client.head(url.clone()).send_retry_default().await?;
         if head.status() == reqwest::StatusCode::OK {
             return Ok(head);
         }
 
         let content_length = format!("{}", file_len);
 
-        let resp = self
-            .client
-            .put(url)
-            .header("Content-Length", &content_length)
-            .header("x-ms-blob-type", "BlockBlob")
-            .body(body)
-            .send()
-            .await?;
+        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);
+
+            let request_builder = self
+                .client
+                .put(url.clone())
+                .header("Content-Length", &content_length)
+                .header("x-ms-blob-type", "BlockBlob")
+                .body(r::Body::wrap_stream(file_stream));
+
+            Ok(request_builder)
+        })
+        .await?;
 
         Ok(resp)
     }
@@ -84,7 +86,7 @@ impl BlobUploader {
             .put(url)
             .header("x-ms-blob-type", "BlockBlob")
             .json(&data)
-            .send()
+            .send_retry_default()
             .await?;
 
         Ok(resp)

diff --git a/src/agent/reqwest-retry/Cargo.toml b/src/agent/reqwest-retry/Cargo.toml
@@ -0,0 +1,13 @@
+[package]
+name = "reqwest-retry"
+version = "0.1.0"
+authors = ["fuzzing@microsoft.com"]
+edition = "2018"
+license = "MIT"
+
+[dependencies]
+anyhow = "1.0"
+async-trait = "0.1.36"
+reqwest = { version = "0.10", features = ["json", "stream"] }
+backoff = { version = "0.2.1", features = ["async-std"] }
+log = "0.4"