Refactor the HTTP retry code to use a single implementation (#135)

* error: include the underlying error in the error text

While these may not be the optimal error messages in all cases, this
should make them more useful by including the underlying error message
so that you get, for example, "HTTP client error occurred: HTTP status
client error (405 Method Not Allowed) for url (http://127.0.0.1:35195/)"
instead of "HTTP client error ocurred".

* http: provide a generic HTTP client interface that retries

Add a generic `request()` function which implements the retry rules we
have and provide `get()` and `post()` wrappers for the function. Move
the IMDS and Wireserver interfaces to this new interface so we don't
have multiple implementations of the retry rules.

* Retry when the response body can't be deserialized

In the event we receive something that can't be deserialized, retry the
request.

This seems like a rather unlikely scenario, since I would expect a
partial response to fail with a lower-level reqwest error, but it can
happen in theory, so let's handle it gracefully.

* imds: bump API version and set extended=true to the query parameters

This causes the request to be treated as a request for data for
provisioning purposes. Without it, the server may respond with an HTTP
200 and the response may not include all the data we need to provision.

Using this query parameter should turn those requests into failures
which we'll retry.

Fixes #134

Author: jeremycline

libazureinit/Cargo.toml

@@ -22,6 +22,7 @@ strum = { version = "0.26.3", features = ["derive"] }
 fstab = "0.4.0"
 
 [dev-dependencies]
+tracing-test = { version = "0.2", features = ["no-env-filter"] }
 tempfile = "3"
 tokio = { version = "1", features = ["full"] }
 tokio-util = "0.7.11"

libazureinit/src/error.rs

@@ -24,13 +24,13 @@
 /// ```
 #[derive(thiserror::Error, Debug)]
 pub enum Error {
-    #[error("Unable to deserialize or serialize JSON data")]
+    #[error("Unable to deserialize or serialize JSON data: {0}")]
     Json(#[from] serde_json::Error),
-    #[error("Unable to deserialize or serialize XML data")]
+    #[error("Unable to deserialize or serialize XML data: {0}")]
     Xml(#[from] serde_xml_rs::Error),
-    #[error("HTTP client error ocurred")]
+    #[error("HTTP client error occurred: {0}")]
     Http(#[from] reqwest::Error),
-    #[error("An I/O error occurred")]
+    #[error("An I/O error occurred: {0}")]
     Io(#[from] std::io::Error),
     #[error("HTTP request did not succeed (HTTP {status} from {endpoint})")]
     HttpStatus {
@@ -44,7 +44,7 @@ pub enum Error {
     },
     #[error("failed to construct a C-style string")]
     NulError(#[from] std::ffi::NulError),
-    #[error("nix call failed")]
+    #[error("nix call failed: {0}")]
     Nix(#[from] nix::Error),
     #[error("The user {user} does not exist")]
     UserMissing { user: String },
@@ -54,7 +54,7 @@ pub enum Error {
     InstanceMetadataFailure,
     #[error("Provisioning a user with a non-empty password is not supported")]
     NonEmptyPassword,
-    #[error("Unable to get list of block devices")]
+    #[error("Unable to get list of block devices: {0}")]
     BlockUtils(#[from] block_utils::BlockUtilsError),
     #[error(
         "Failed to set the hostname; none of the provided backends succeeded"
@@ -69,5 +69,11 @@ pub enum Error {
     )]
     NoPasswordProvisioner,
     #[error("A timeout error occurred")]
-    Timeout(#[from] tokio::time::error::Elapsed),
+    Timeout,
+}
+
+impl From<tokio::time::error::Elapsed> for Error {
+    fn from(_: tokio::time::error::Elapsed) -> Self {
+        Self::Timeout
+    }
 }

libazureinit/src/goalstate.rs

@@ -3,15 +3,14 @@
 
 use reqwest::header::HeaderMap;
 use reqwest::header::HeaderValue;
-use reqwest::{Client, StatusCode};
+use reqwest::Client;
+use tracing::instrument;
 
 use std::time::Duration;
 
 use serde::Deserialize;
 use serde_xml_rs::from_str;
 
-use tokio::time::timeout;
-
 use crate::error::Error;
 use crate::http;
 
@@ -101,65 +100,52 @@ const DEFAULT_GOALSTATE_URL: &str =
 ///     Some("http://127.0.0.1:8000/"),
 /// );
 /// ```
+#[instrument(err, skip_all)]
 pub async fn get_goalstate(
     client: &Client,
     retry_interval: Duration,
-    total_timeout: Duration,
+    mut total_timeout: Duration,
     url: Option<&str>,
 ) -> Result<Goalstate, Error> {
-    let url = url.unwrap_or(DEFAULT_GOALSTATE_URL);
-
     let mut headers = HeaderMap::new();
     headers.insert("x-ms-agent-name", HeaderValue::from_static("azure-init"));
     headers.insert("x-ms-version", HeaderValue::from_static("2012-11-30"));
-
-    let response = timeout(total_timeout, async {
-        let now = std::time::Instant::now();
-        loop {
-            if let Ok(response) = client
-                .get(url)
-                .headers(headers.clone())
-                .timeout(Duration::from_secs(http::WIRESERVER_HTTP_TIMEOUT_SEC))
-                .send()
-                .await
-            {
-                let statuscode = response.status();
-
-                if statuscode == StatusCode::OK {
-                    tracing::info!(
-                        "HTTP response succeeded with status {}",
-                        statuscode
+    let url = url.unwrap_or(DEFAULT_GOALSTATE_URL);
+    let request_timeout =
+        Duration::from_secs(http::WIRESERVER_HTTP_TIMEOUT_SEC);
+
+    while !total_timeout.is_zero() {
+        let (response, remaining_timeout) = http::get(
+            client,
+            headers.clone(),
+            request_timeout,
+            retry_interval,
+            total_timeout,
+            url,
+        )
+        .await?;
+        match response.text().await {
+            Ok(body) => {
+                let goalstate = from_str(&body).map_err(|error| {
+                    tracing::warn!(
+                        ?error,
+                        "The response body was invalid and could not be deserialized"
                     );
-                    return Ok(response);
-                }
-
-                if !http::RETRY_CODES.contains(&statuscode) {
-                    return response.error_for_status().map_err(|error| {
-                        tracing::error!(
-                            ?error,
-                            "{}",
-                            format!(
-                                "HTTP call failed due to status {}",
-                                statuscode
-                            )
-                        );
-                        error
-                    });
+                    error.into()
+                });
+                if goalstate.is_ok() {
+                    return goalstate;
                 }
             }
-
-            tracing::info!("Retrying to get HTTP response in {} sec, remaining timeout {} sec.", retry_interval.as_secs(), total_timeout.saturating_sub(now.elapsed()).as_secs());
-
-            tokio::time::sleep(retry_interval).await;
+            Err(error) => {
+                tracing::warn!(?error, "Failed to read the full response body")
+            }
         }
-    })
-    .await?;
-
-    let goalstate_body = response?.text().await?;
 
-    let goalstate: Goalstate = from_str(&goalstate_body)?;
+        total_timeout = remaining_timeout;
+    }
 
-    Ok(goalstate)
+    Err(Error::Timeout)
 }
 
 const DEFAULT_HEALTH_URL: &str = "http://168.63.129.16/machine/?comp=health";
@@ -199,56 +185,36 @@ const DEFAULT_HEALTH_URL: &str = "http://168.63.129.16/machine/?comp=health";
 ///     );
 /// }
 /// ```
+#[instrument(err, skip_all)]
 pub async fn report_health(
     client: &Client,
     goalstate: Goalstate,
     retry_interval: Duration,
     total_timeout: Duration,
     url: Option<&str>,
 ) -> Result<(), Error> {
-    let url = url.unwrap_or(DEFAULT_HEALTH_URL);
-
     let mut headers = HeaderMap::new();
     headers.insert("x-ms-agent-name", HeaderValue::from_static("azure-init"));
     headers.insert("x-ms-version", HeaderValue::from_static("2012-11-30"));
     headers.insert(
         "Content-Type",
         HeaderValue::from_static("text/xml;charset=utf-8"),
     );
+    let request_timeout =
+        Duration::from_secs(http::WIRESERVER_HTTP_TIMEOUT_SEC);
+    let url = url.unwrap_or(DEFAULT_HEALTH_URL);
 
     let post_request = build_report_health_file(goalstate);
 
-    _ = timeout(total_timeout, async {
-        let now = std::time::Instant::now();
-        loop {
-            if let Ok(response) = client
-                .post(url)
-                .headers(headers.clone())
-                .body(post_request.clone())
-                .timeout(Duration::from_secs(http::WIRESERVER_HTTP_TIMEOUT_SEC))
-                .send()
-                .await
-            {
-                let statuscode = response.status();
-
-                if statuscode == StatusCode::OK {
-                    tracing::info!("HTTP response succeeded with status {}", statuscode);
-                    return Ok(response);
-                }
-
-                if !http::RETRY_CODES.contains(&statuscode) {
-                    return response.error_for_status().map_err(|error| {
-                        tracing::error!(?error, "{}", format!("HTTP call failed due to status {}", statuscode));
-                        error
-                    });
-                }
-            }
-
-            tracing::info!("Retrying to get HTTP response in {} sec, remaining timeout {} sec.", retry_interval.as_secs(), total_timeout.saturating_sub(now.elapsed()).as_secs());
-
-            tokio::time::sleep(retry_interval).await;
-        }
-    })
+    _ = http::post(
+        client,
+        headers,
+        post_request,
+        request_timeout,
+        retry_interval,
+        total_timeout,
+        url,
+    )
     .await?;
 
     Ok(())
@@ -458,4 +424,54 @@ mod tests {
             assert!(!run_goalstate_retry(rc).await);
         }
     }
+
+    // Assert malformed responses are retried.
+    //
+    // In this case the server doesn't return XML at all.
+    #[tokio::test]
+    #[tracing_test::traced_test]
+    async fn malformed_response() {
+        let body = "You thought this was XML, but you were wrong";
+        let payload = format!(
+            "HTTP/1.1 {} {}\r\nContent-Type: application/xml\r\nContent-Length: {}\r\n\r\n{}",
+             StatusCode::OK.as_u16(),
+             StatusCode::OK.to_string(),
+             body.len(),
+             body
+        );
+
+        let serverlistener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let addr = serverlistener.local_addr().unwrap();
+        let cancel_token = tokio_util::sync::CancellationToken::new();
+        let server = tokio::spawn(unittest::serve_requests(
+            serverlistener,
+            payload,
+            cancel_token.clone(),
+        ));
+
+        let client = Client::builder()
+            .timeout(std::time::Duration::from_secs(5))
+            .build()
+            .unwrap();
+
+        let res = get_goalstate(
+            &client,
+            Duration::from_millis(10),
+            Duration::from_millis(50),
+            Some(format!("http://{:}:{:}/", addr.ip(), addr.port()).as_str()),
+        )
+        .await;
+
+        cancel_token.cancel();
+
+        let requests = server.await.unwrap();
+        assert!(requests >= 2);
+        assert!(logs_contain(
+            "The response body was invalid and could not be deserialized"
+        ));
+        match res {
+            Err(crate::error::Error::Timeout) => {}
+            _ => panic!("Response should have timed out"),
+        };
+    }
 }