enable the supervisor to handle longer service outages

src/agent/onefuzz-supervisor/src/agent.rs

@@ -25,6 +25,7 @@ pub struct Agent {
     worker_runner: Box<dyn IWorkerRunner>,
     heartbeat: Option<AgentHeartbeatClient>,
     previous_state: NodeState,
+    last_poll_command: Option<PollCommandResult>,
 }
 
 impl Agent {
@@ -39,6 +40,7 @@ impl Agent {
     ) -> Self {
         let scheduler = Some(scheduler);
         let previous_state = NodeState::Init;
+        let last_poll_command = None;
 
         Self {
             coordinator,
@@ -49,6 +51,7 @@ impl Agent {
             worker_runner,
             heartbeat,
             previous_state,
+            last_poll_command,
         }
     }
 
@@ -277,13 +280,30 @@ impl Agent {
     }
 
     async fn execute_pending_commands(&mut self) -> Result<()> {
-        let cmd = self.coordinator.poll_commands().await?;
+        let result = self.coordinator.poll_commands().await?;
 
-        if let Some(cmd) = cmd {
-            info!("agent received node command: {:?}", cmd);
-            self.scheduler()?.execute_command(cmd).await?;
+        match &result {
+            PollCommandResult::None => {}
+            PollCommandResult::Command(cmd) => {
+                info!("agent received node command: {:?}", cmd);
+                self.scheduler()?.execute_command(cmd).await?;
+            }
+            PollCommandResult::RequestFailed(err) => {
+                error!("error polling the service for commands: {:?}", err);
+            }
+            PollCommandResult::ClaimFailed(err) => {
+                if matches!(
+                    self.last_poll_command,
+                    Some(PollCommandResult::ClaimFailed(..))
+                ) {
+                    bail!("repeated command claim attempt failures: {:?}", err);
+                }
+                error!("error polling the service for commands: {:?}", err);
+            }
         }
 
+        self.last_poll_command = Some(result);
+
         Ok(())
     }
 

src/agent/onefuzz-supervisor/src/commands.rs

@@ -159,7 +159,7 @@ pub async fn add_ssh_key(key_info: SshKeyInfo) -> Result<()> {
 }
 
 #[cfg(target_family = "unix")]
-pub async fn add_ssh_key(key_info: SshKeyInfo) -> Result<()> {
+pub async fn add_ssh_key(key_info: &SshKeyInfo) -> Result<()> {
     if get_scaleset_name().await?.is_none() {
         warn!("adding ssh keys only supported on managed nodes");
         return Ok(());

src/agent/onefuzz-supervisor/src/coordinator.rs

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-use anyhow::{Context, Result};
+use anyhow::{Context, Error, Result};
 use downcast_rs::Downcast;
 use onefuzz::{auth::AccessToken, http::ResponseExt, process::Output};
 use reqwest::{Client, RequestBuilder, Response, StatusCode};
@@ -148,7 +148,7 @@ pub struct TaskInfo {
 
 #[async_trait]
 pub trait ICoordinator: Downcast {
-    async fn poll_commands(&mut self) -> Result<Option<NodeCommand>>;
+    async fn poll_commands(&mut self) -> Result<PollCommandResult>;
 
     async fn emit_event(&mut self, event: NodeEvent) -> Result<()>;
 
@@ -159,7 +159,7 @@ impl_downcast!(ICoordinator);
 
 #[async_trait]
 impl ICoordinator for Coordinator {
-    async fn poll_commands(&mut self) -> Result<Option<NodeCommand>> {
+    async fn poll_commands(&mut self) -> Result<PollCommandResult> {
         self.poll_commands().await
     }
 
@@ -172,6 +172,13 @@ impl ICoordinator for Coordinator {
     }
 }
 
+pub enum PollCommandResult {
+    None,
+    Command(NodeCommand),
+    RequestFailed(Error),
+    ClaimFailed(Error),
+}
+
 pub struct Coordinator {
     client: Client,
     registration: Registration,
@@ -194,18 +201,25 @@ impl Coordinator {
     ///
     /// If the request fails due to an expired access token, we will retry once
     /// with a fresh one.
-    pub async fn poll_commands(&mut self) -> Result<Option<NodeCommand>> {
+    pub async fn poll_commands(&mut self) -> Result<PollCommandResult> {
         let request = PollCommandsRequest {
             machine_id: self.registration.machine_id,
         };
 
         let url = self.registration.dynamic_config.commands_url.clone();
         let request = self.client.get(url).json(&request);
 
-        let pending: PendingNodeCommand = self
-            .send_request(request)
-            .await
-            .context("PollCommands")?
+        let response = self.send_request(request).await.context("PollCommands");
+
+        // Treat communication issues with the service as `no commands
+        // available`.  This adds resiliency to the supervisor during longer
+        // outages.  Given poll_commands runs on a 10 second cycle, this should
+        // provide eventual recovery.
+        if let Err(response) = response {
+            return Ok(PollCommandResult::RequestFailed(response));
+        }
+
+        let pending: PendingNodeCommand = response?
             .json()
             .await
             .context("parsing PollCommands response")?;
@@ -219,11 +233,17 @@ impl Coordinator {
             let url = self.registration.dynamic_config.commands_url.clone();
             let request = self.client.delete(url).json(&request);
 
-            self.send_request(request).await.context("ClaimCommand")?;
+            let response = self.send_request(request).await.context("ClaimCommand");
+
+            // similar polling available commands, this treats issues claiming
+            // commands as `no commands available`
+            if let Err(response) = response {
+                return Ok(PollCommandResult::ClaimFailed(response));
+            }
 
-            Ok(Some(envelope.command))
+            Ok(PollCommandResult::Command(envelope.command))
         } else {
-            Ok(None)
+            Ok(PollCommandResult::None)
         }
     }
 

src/agent/onefuzz-supervisor/src/coordinator/double.rs

@@ -11,8 +11,11 @@ pub struct CoordinatorDouble {
 
 #[async_trait]
 impl ICoordinator for CoordinatorDouble {
-    async fn poll_commands(&mut self) -> Result<Option<NodeCommand>> {
-        Ok(self.commands.pop())
+    async fn poll_commands(&mut self) -> Result<PollCommandResult> {
+        match self.commands.pop() {
+            Some(x) => Ok(PollCommandResult::Command(x)),
+            None => Ok(PollCommandResult::None),
+        }
     }
 
     async fn emit_event(&mut self, event: NodeEvent) -> Result<()> {

src/agent/onefuzz-supervisor/src/scheduler.rs

@@ -54,7 +54,7 @@ impl Scheduler {
         Self::default()
     }
 
-    pub async fn execute_command(&mut self, cmd: NodeCommand) -> Result<()> {
+    pub async fn execute_command(&mut self, cmd: &NodeCommand) -> Result<()> {
         match cmd {
             NodeCommand::AddSshKey(ssh_key_info) => {
                 add_ssh_key(ssh_key_info).await?;