Sweep forgotten client operation IDs (#1965)

* Sweep forgotten client operation IDs

* add helpful log

Author: MarcusSorealheis

nativelink-scheduler/src/store_awaited_action_db.rs

@@ -614,6 +614,46 @@ where
         let Some(operation_id) = maybe_operation_id else {
             return Ok(None);
         };
+
+        // Validate that the internal operation actually exists.
+        // If it doesn't, this is an orphaned client operation mapping that should be cleaned up.
+        // This can happen when an operation is deleted (completed/timed out) but the
+        // client_id -> operation_id mapping persists in the store.
+        let maybe_awaited_action = match self
+            .store
+            .get_and_decode(OperationIdToAwaitedAction(Cow::Borrowed(&operation_id)))
+            .await
+        {
+            Ok(maybe_action) => maybe_action,
+            Err(err) if err.code == Code::NotFound => {
+                tracing::warn!(
+                    "Orphaned client operation mapping detected: client_id={} maps to operation_id={}, \
+                    but the operation does not exist in the store (NotFound). This typically happens when \
+                    an operation completes or times out but the client mapping persists.",
+                    client_operation_id,
+                    operation_id
+                );
+                None
+            }
+            Err(err) => {
+                // Some other error occurred
+                return Err(err).err_tip(
+                    || "In RedisAwaitedActionDb::get_awaited_action_by_id::validate_operation",
+                );
+            }
+        };
+
+        if maybe_awaited_action.is_none() {
+            tracing::warn!(
+                "Found orphaned client operation mapping: client_id={} -> operation_id={}, \
+                but operation no longer exists. Returning None to prevent client from polling \
+                a non-existent operation.",
+                client_operation_id,
+                operation_id
+            );
+            return Ok(None);
+        }
+
         Ok(Some(OperationSubscriber::new(
             Some(client_operation_id.clone()),
             OperationIdToAwaitedAction(Cow::Owned(operation_id)),

nativelink-scheduler/tests/redis_store_awaited_action_db_test.rs

@@ -688,6 +688,25 @@ async fn add_action_smoke_test() -> Result<(), Error> {
             ])),
             None,
         )
+        // Validation HMGET: Check if the internal operation exists (orphan detection)
+        .expect(
+            MockCommand {
+                cmd: Str::from_static("HMGET"),
+                subcommand: None,
+                args: vec![
+                    format!("aa_{WORKER_OPERATION_ID}").as_bytes().into(),
+                    "version".as_bytes().into(),
+                    "data".as_bytes().into(),
+                ],
+            },
+            Ok(RedisValue::Array(vec![
+                // Version.
+                "1".into(),
+                // Data.
+                RedisValue::Bytes(Bytes::from(serde_json::to_string(&worker_awaited_action).unwrap())),
+            ])),
+            None,
+        )
         .expect(
             MockCommand {
                 cmd: Str::from_static("HMGET"),
@@ -1029,3 +1048,64 @@ async fn test_outdated_version() -> Result<(), Error> {
 
     Ok(())
 }
+
+/// Test that orphaned client operation ID mappings return None.
+///
+/// This tests the scenario where:
+/// 1. A client operation ID mapping exists (cid_* → operation_id)
+/// 2. The actual operation (aa_*) has been deleted (completed/timed out)
+/// 3. get_awaited_action_by_id should return None instead of a subscriber to a non-existent operation
+#[nativelink_test]
+async fn test_orphaned_client_operation_id_returns_none() -> Result<(), Error> {
+    const CLIENT_OPERATION_ID: &str = "orphaned_client_id";
+    const INTERNAL_OPERATION_ID: &str = "deleted_internal_operation_id";
+    const SUB_CHANNEL: &str = "sub_channel";
+
+    let worker_operation_id = Arc::new(Mutex::new(INTERNAL_OPERATION_ID));
+    let worker_operation_id_clone = worker_operation_id.clone();
+
+    let internal_operation_id = OperationId::from(INTERNAL_OPERATION_ID);
+
+    // Use FakeRedisBackend which handles SUBSCRIBE automatically
+    let mocks = Arc::new(FakeRedisBackend::new());
+    let store = make_redis_store(SUB_CHANNEL, mocks.clone());
+    mocks.set_subscription_manager(store.subscription_manager().unwrap());
+
+    // Manually set up the orphaned state in the fake backend:
+    // 1. Add client_id → operation_id mapping (cid_* key)
+    {
+        let mut table = mocks.table.lock();
+        let mut client_fields = HashMap::new();
+        client_fields.insert(
+            "data".into(),
+            RedisValue::Bytes(Bytes::from(
+                serde_json::to_string(&internal_operation_id).unwrap(),
+            )),
+        );
+        table.insert(format!("cid_{CLIENT_OPERATION_ID}"), client_fields);
+    }
+    // 2. Don't add the actual operation (aa_* key) - this simulates it being deleted/orphaned
+
+    let notifier = Arc::new(Notify::new());
+    let awaited_action_db = StoreAwaitedActionDb::new(
+        store.clone(),
+        notifier.clone(),
+        MockInstantWrapped::default,
+        move || worker_operation_id_clone.lock().clone().into(),
+    )
+    .unwrap();
+
+    // Try to get the awaited action by the client operation ID
+    // This should return None because the internal operation doesn't exist (orphaned mapping)
+    let result = awaited_action_db
+        .get_awaited_action_by_id(&OperationId::from(CLIENT_OPERATION_ID))
+        .await
+        .expect("Should not error when checking orphaned client operation");
+
+    assert!(
+        result.is_none(),
+        "Expected None for orphaned client operation ID, but got a subscription"
+    );
+
+    Ok(())
+}