Tidy DelWorkload handling in a pre-snapshot context

src/inpod/statemanager.rs

@@ -134,6 +134,7 @@ impl WorkloadProxyManagerState {
                     ));
                 };
                 if !self.snapshot_received {
+                    debug!("got workload add before snapshot");
                     self.snapshot_names.insert(poddata.workload_uid.clone());
                 }
                 let netns =
@@ -172,12 +173,20 @@ impl WorkloadProxyManagerState {
                     "pod delete request, shutting down proxy"
                 );
                 if !self.snapshot_received {
-                    // TODO: consider if this is an error. if not, do this instead:
-                    // self.snapshot_names.remove(&workload_uid)
-                    // self.pending_workloads.remove(&workload_uid)
-                    return Err(Error::ProtocolError(
-                        "pod delete received before snapshot".into(),
-                    ));
+                    debug!("got workload delete before snapshot");
+                    // Since we insert here on AddWorkload before we get a snapshot,
+                    // make sure we also opportunistically remove here before we
+                    // get a snapshot
+                    //
+                    // Note that even though AddWorkload starts the workload, we do *not* need
+                    // to stop it here, as it should be auto-dropped subsequently during snapshot
+                    // reconcile(), when we actually get the `SnapshotSent` notification.
+                    self.snapshot_names.remove(&workload_uid);
+                    // `reconcile()` will drop this workload later, but if the workload never successfully
+                    // starts it will stay in the pending queue (which `reconcile()` can't remove it from),
+                    // so clear the pending queue here.
+                    self.pending_workloads.remove(&workload_uid);
+                    return Ok(());
                 }
                 self.del_workload(&workload_uid);
                 Ok(())
@@ -536,7 +545,7 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn idemepotency_add_workload_fails_and_then_deleted() {
+    async fn idempotency_add_workload_fails_and_then_deleted() {
         let fixture = fixture!();
         let mut state = fixture.state;
 
@@ -567,6 +576,45 @@ mod tests {
         state.drain().await;
     }
 
+    #[tokio::test]
+    async fn del_workload_before_snapshot_removes_from_snapshot_and_pending() {
+        let fixture = fixture!();
+        let mut state = fixture.state;
+
+        let ns = new_netns();
+
+        // to make the proxy fail, bind to its ports in its netns
+        let _sock = create_proxy_conflict(&ns);
+
+        let data = WorkloadData {
+            netns: ns,
+            workload_uid: uid(0),
+            workload_info: workload_info(),
+        };
+
+        let ret = state.process_msg(WorkloadMessage::AddWorkload(data)).await;
+
+        assert!(state.snapshot_names.len() == 1);
+        assert!(ret.is_err());
+        assert!(state.have_pending());
+
+        state
+            .process_msg(WorkloadMessage::DelWorkload(uid(0)))
+            .await
+            .unwrap();
+
+        assert!(state.snapshot_names.is_empty());
+
+        state
+            .process_msg(WorkloadMessage::WorkloadSnapshotSent)
+            .await
+            .unwrap();
+
+        assert!(state.snapshot_names.is_empty());
+        assert!(!state.have_pending());
+        state.drain().await;
+    }
+
     #[tokio::test]
     async fn add_delete_add_workload_starts_only_one_proxy() {
         let fixture = fixture!();