redis_store_awaited_action_db_test.rs

// Copyright 2024 The NativeLink Authors. All rights reserved.
//
// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    See LICENSE file for details
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use core::time::Duration;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::SystemTime;

use futures::StreamExt;
use mock_instant::global::SystemTime as MockSystemTime;
use nativelink_config::schedulers::SimpleSpec;
use nativelink_config::stores::RedisSpec;
use nativelink_error::{Error, ResultExt};
use nativelink_macro::nativelink_test;
use nativelink_proto::build::bazel::remote::execution::v2::{
    ExecuteRequest, Platform, digest_function,
};
use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{
    ConnectionResult, StartExecute, UpdateForWorker, update_for_worker,
};
use nativelink_redis_tester::FakeRedisBackend;
use nativelink_scheduler::awaited_action_db::{
    AwaitedAction, AwaitedActionDb, AwaitedActionSubscriber,
};
use nativelink_scheduler::simple_scheduler::SimpleScheduler;
use nativelink_scheduler::store_awaited_action_db::StoreAwaitedActionDb;
use nativelink_scheduler::worker::Worker;
use nativelink_scheduler::worker_scheduler::WorkerScheduler;
use nativelink_store::redis_store::{RedisStore, RedisSubscriptionManager};
use nativelink_util::action_messages::{
    ActionInfo, ActionStage, ActionUniqueKey, ActionUniqueQualifier, OperationId, WorkerId,
};
use nativelink_util::common::DigestInfo;
use nativelink_util::digest_hasher::DigestHasherFunc;
use nativelink_util::instant_wrapper::MockInstantWrapped;
use nativelink_util::operation_state_manager::{ClientStateManager, OperationFilter};
use nativelink_util::platform_properties::PlatformProperties;
use nativelink_util::store_trait::SchedulerStore;
use parking_lot::Mutex;
use pretty_assertions::assert_eq;
use redis::Value;
use tokio::sync::mpsc::unbounded_channel;
use tokio::sync::{Notify, mpsc};
use tonic::Code;
use utils::scheduler_utils::update_eq;

mod utils {
    pub(crate) mod scheduler_utils;
}

const INSTANCE_NAME: &str = "instance_name";

async fn verify_initial_connection_message(
    worker_id: WorkerId,
    rx: &mut mpsc::UnboundedReceiver<UpdateForWorker>,
) {
    // Worker should have been sent an execute command.
    let expected_msg_for_worker = UpdateForWorker {
        update: Some(update_for_worker::Update::ConnectionResult(
            ConnectionResult {
                worker_id: worker_id.into(),
            },
        )),
    };
    let msg_for_worker = rx.recv().await.unwrap();
    assert_eq!(msg_for_worker, expected_msg_for_worker);
}

const NOW_TIME: u64 = 10000;

async fn setup_new_worker(
    scheduler: &SimpleScheduler,
    worker_id: WorkerId,
    props: PlatformProperties,
) -> Result<mpsc::UnboundedReceiver<UpdateForWorker>, Error> {
    let (tx, mut rx) = unbounded_channel();
    let worker = Worker::new(worker_id.clone(), props, tx, NOW_TIME, 0);
    scheduler
        .add_worker(worker)
        .await
        .err_tip(|| "Failed to add worker")?;
    tokio::task::yield_now().await; // Allow task<->worker matcher to run.
    verify_initial_connection_message(worker_id, &mut rx).await;
    Ok(rx)
}

fn make_awaited_action(operation_id: &str) -> AwaitedAction {
    AwaitedAction::new(
        operation_id.into(),
        Arc::new(ActionInfo {
            command_digest: DigestInfo::zero_digest(),
            input_root_digest: DigestInfo::zero_digest(),
            timeout: Duration::from_secs(1),
            platform_properties: HashMap::new(),
            priority: 0,
            load_timestamp: SystemTime::UNIX_EPOCH,
            insert_timestamp: SystemTime::UNIX_EPOCH,
            unique_qualifier: ActionUniqueQualifier::Cacheable(ActionUniqueKey {
                instance_name: INSTANCE_NAME.to_string(),
                digest_function: DigestHasherFunc::Sha256,
                digest: DigestInfo::zero_digest(),
            }),
        }),
        MockSystemTime::now().into(),
    )
}

// TODO: This test needs to be rewritten to use workers (like test_multiple_clients_subscribe_to_same_action).
#[nativelink_test]
#[ignore = "needs rewrite to use workers (like test_multiple_clients_subscribe_to_same_action)"]
async fn add_action_smoke_test() -> Result<(), Error> {
    const CLIENT_OPERATION_ID: &str = "my_client_operation_id";
    const WORKER_OPERATION_ID: &str = "my_worker_operation_id";
    const SUB_CHANNEL: &str = "sub_channel";

    let worker_awaited_action = make_awaited_action(WORKER_OPERATION_ID);
    let new_awaited_action = {
        let mut new_awaited_action = worker_awaited_action.clone();
        let mut new_state = new_awaited_action.state().as_ref().clone();
        new_state.stage = ActionStage::Executing;
        new_state.last_transition_timestamp = SystemTime::now();
        new_awaited_action.worker_set_state(Arc::new(new_state), MockSystemTime::now().into());
        new_awaited_action
    };

    // Use FakeRedisBackend which handles all Redis commands dynamically
    // This is more maintainable than the standard fake redis which requires exact command sequences
    let fake_redis_backend: FakeRedisBackend<RedisSubscriptionManager> = FakeRedisBackend::new();
    let fake_redis_port = fake_redis_backend.clone().run().await;
    let spec = RedisSpec {
        addresses: vec![format!("redis://127.0.0.1:{fake_redis_port}")],
        experimental_pub_sub_channel: Some(SUB_CHANNEL.to_string()),
        ..Default::default()
    };
    let store = RedisStore::new_standard(spec).await.expect("Working spec");
    fake_redis_backend.set_subscription_manager(store.subscription_manager().await.unwrap());

    let notifier = Arc::new(Notify::new());
    let awaited_action_db = StoreAwaitedActionDb::new(
        store.clone(),
        notifier.clone(),
        MockInstantWrapped::default,
        move || WORKER_OPERATION_ID.into(),
        60,
    )
    .await
    .unwrap();

    let mut subscription = awaited_action_db
        .add_action(
            CLIENT_OPERATION_ID.into(),
            worker_awaited_action.action_info().clone(),
            Duration::from_mins(1),
        )
        .await
        .unwrap();

    {
        // Check initial change state.
        let changed_awaited_action_res = subscription.changed().await;

        assert_eq!(
            changed_awaited_action_res.unwrap().state().stage,
            ActionStage::Queued
        );
    }

    {
        let get_subscription = awaited_action_db
            .get_awaited_action_by_id(&OperationId::from(CLIENT_OPERATION_ID))
            .await
            .unwrap()
            .unwrap();

        let get_res = get_subscription.borrow().await;

        assert_eq!(get_res.unwrap().state().stage, ActionStage::Queued);
    }

    {
        // Update the action and check the new state.
        let (changed_awaited_action_res, update_res) = tokio::join!(
            subscription.changed(),
            awaited_action_db.update_awaited_action(new_awaited_action.clone())
        );
        assert_eq!(update_res, Ok(()));

        assert_eq!(
            changed_awaited_action_res.unwrap().state().stage,
            ActionStage::Executing
        );
    }

    {
        let get_subscription = awaited_action_db
            .get_awaited_action_by_id(&OperationId::from(CLIENT_OPERATION_ID))
            .await
            .unwrap()
            .unwrap();

        let get_res = get_subscription.borrow().await;

        assert_eq!(get_res.unwrap().state().stage, ActionStage::Executing);
    }

    Ok(())
}

#[nativelink_test]
async fn test_multiple_clients_subscribe_to_same_action() -> Result<(), Error> {
    const CLIENT_OPERATION_ID_1: &str = "client_operation_id_1";
    const CLIENT_OPERATION_ID_2: &str = "client_operation_id_2";
    const CLIENT_OPERATION_ID_3: &str = "client_operation_id_3";
    const WORKER_OPERATION_ID_1: &str = "worker_operation_id_1";
    const WORKER_OPERATION_ID_2: &str = "worker_operation_id_2";
    const SUB_CHANNEL: &str = "sub_channel";

    let action_info = Arc::new(ActionInfo {
        command_digest: DigestInfo::zero_digest(),
        input_root_digest: DigestInfo::zero_digest(),
        timeout: Duration::from_secs(1),
        platform_properties: HashMap::new(),
        priority: 0,
        load_timestamp: SystemTime::UNIX_EPOCH,
        insert_timestamp: SystemTime::UNIX_EPOCH,
        unique_qualifier: ActionUniqueQualifier::Cacheable(ActionUniqueKey {
            instance_name: INSTANCE_NAME.to_string(),
            digest_function: DigestHasherFunc::Sha256,
            digest: DigestInfo::zero_digest(),
        }),
    });

    // Use FakeRedisBackend which handles all Redis commands dynamically
    // This is more maintainable than the standard fake redis which requires exact command sequences
    let fake_redis_backend: FakeRedisBackend<RedisSubscriptionManager> = FakeRedisBackend::new();
    let fake_redis_port = fake_redis_backend.clone().run().await;
    let spec = RedisSpec {
        addresses: vec![format!("redis://127.0.0.1:{fake_redis_port}")],
        experimental_pub_sub_channel: Some(SUB_CHANNEL.to_string()),
        ..Default::default()
    };
    let store = RedisStore::new_standard(spec).await.expect("Working spec");
    fake_redis_backend.set_subscription_manager(store.subscription_manager().await.unwrap());

    let notifier = Arc::new(Notify::new());
    let worker_operation_id = Arc::new(Mutex::new(WORKER_OPERATION_ID_1));
    let worker_operation_id_clone = worker_operation_id.clone();
    let awaited_action_db = StoreAwaitedActionDb::new(
        store.clone(),
        notifier.clone(),
        MockInstantWrapped::default,
        move || worker_operation_id_clone.lock().clone().into(),
        60,
    )
    .await
    .unwrap();

    let task_change_notify = Arc::new(Notify::new());
    let (scheduler, _worker_scheduler) = SimpleScheduler::new_with_callback(
        &SimpleSpec::default(),
        awaited_action_db,
        || async move {},
        task_change_notify,
        MockInstantWrapped::default,
        None,
    );

    // First client adds the action
    let mut subscription1 = scheduler
        .add_action(CLIENT_OPERATION_ID_1.into(), action_info.clone())
        .await
        .unwrap();

    // Second client tries to add the same action (should subscribe to existing one)
    let _subscription2 = scheduler
        .add_action(CLIENT_OPERATION_ID_2.into(), action_info.clone())
        .await
        .unwrap();

    // Second client should be able to get the action by its client_operation_id
    let get_subscription = scheduler
        .filter_operations(OperationFilter {
            client_operation_id: Some(OperationId::from(CLIENT_OPERATION_ID_2)),
            ..Default::default()
        })
        .await
        .expect("Second client should be able to get action by its client_operation_id")
        .next()
        .await
        .expect("Second client should be able to get action by its client_operation_id");

    let (state, _metadata) = get_subscription
        .as_state()
        .await
        .expect("Unable to get state of operation");
    assert_eq!(state.stage, ActionStage::Queued);

    // Now create a worker and check that it is only allocated the job once.
    let worker_id = WorkerId("worker_id".to_string());
    let mut rx_from_worker =
        setup_new_worker(&scheduler, worker_id.clone(), PlatformProperties::default()).await?;

    // Try to ensure we schedule to the worker.
    scheduler.do_try_match_for_test().await?;

    {
        // Worker should have been sent an execute command.
        let expected_msg_for_worker = UpdateForWorker {
            update: Some(update_for_worker::Update::StartAction(StartExecute {
                execute_request: Some(ExecuteRequest {
                    instance_name: INSTANCE_NAME.to_string(),
                    action_digest: Some(DigestInfo::zero_digest().into()),
                    digest_function: digest_function::Value::Sha256.into(),
                    ..Default::default()
                }),
                operation_id: "Unknown Generated internally".to_string(),
                queued_timestamp: Some(SystemTime::UNIX_EPOCH.into()),
                platform: Some(Platform::default()),
                worker_id: worker_id.clone().into(),
            })),
        };
        let msg_for_worker = rx_from_worker.recv().await.unwrap();
        // Operation ID is random so we ignore it.
        assert!(update_eq(expected_msg_for_worker, msg_for_worker, true));
    }

    let (state, _metadata) = subscription1
        .changed()
        .await
        .expect("No update to subscription");
    assert_eq!(state.stage, ActionStage::Executing);
    let (state, _metadata) = get_subscription
        .as_state()
        .await
        .expect("Unable to get second operation");
    assert_eq!(state.stage, ActionStage::Executing);

    // Immediately try to schedule again to check we don't schedule the job
    // again now that it's executing.
    scheduler.do_try_match_for_test().await?;

    // The worker shouldn't be allocated the job again.
    tokio::select! {
        () = tokio::time::sleep(Duration::from_secs(1)) => {}
        v = rx_from_worker.recv() => {
            panic!("Worker was allocated another job: {v:?}");
        }
    }

    // The worker goes away without completing the task, so the action goes back
    // to queued.
    drop(rx_from_worker);
    scheduler.remove_worker(&worker_id).await?;

    let (state, _metadata) = subscription1
        .changed()
        .await
        .expect("No update to subscription");
    assert_eq!(state.stage, ActionStage::Queued);

    // Create and drop the worker three times to cause the job to complete with
    // a failure.
    for _ in 0..3 {
        let rx_from_worker =
            setup_new_worker(&scheduler, worker_id.clone(), PlatformProperties::default()).await?;
        scheduler.do_try_match_for_test().await?;
        drop(rx_from_worker);
        scheduler.remove_worker(&worker_id).await?;
    }

    // Update the operation ID for the new subscription.
    *worker_operation_id.lock() = WORKER_OPERATION_ID_2;

    // Subscribe with a new operation ID after all that and we should be queued.
    let subscription3 = scheduler
        .add_action(CLIENT_OPERATION_ID_3.into(), action_info.clone())
        .await
        .unwrap();

    let (state, _metadata) = subscription3
        .as_state()
        .await
        .expect("Unable to get state of operation");
    assert_eq!(state.stage, ActionStage::Queued);

    Ok(())
}

#[nativelink_test]
async fn test_outdated_version() -> Result<(), Error> {
    const CLIENT_OPERATION_ID: &str = "outdated_operation_id";
    let worker_operation_id = Arc::new(Mutex::new(CLIENT_OPERATION_ID));
    let worker_operation_id_clone = worker_operation_id.clone();

    let fake_redis_backend: FakeRedisBackend<RedisSubscriptionManager> = FakeRedisBackend::new();
    let fake_redis_port = fake_redis_backend.clone().run().await;
    let spec = RedisSpec {
        addresses: vec![format!("redis://127.0.0.1:{fake_redis_port}")],
        experimental_pub_sub_channel: Some("sub_channel".into()),
        ..Default::default()
    };
    let store = RedisStore::new_standard(spec).await.expect("Working spec");
    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(),
        60,
    )
    .await
    .unwrap();

    let worker_awaited_action = make_awaited_action("WORKER_OPERATION_ID");

    let update_res = awaited_action_db
        .update_awaited_action(worker_awaited_action.clone())
        .await;
    assert_eq!(update_res, Ok(()));

    let update_res2 = awaited_action_db
        .update_awaited_action(worker_awaited_action.clone())
        .await;
    assert!(update_res2.is_err());
    assert_eq!(
        update_res2.unwrap_err(),
        Error::new(Code::Aborted, "Could not update AwaitedAction because the version did not match for WORKER_OPERATION_ID".into())
    );

    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 fake_redis_backend: FakeRedisBackend<RedisSubscriptionManager> = FakeRedisBackend::new();
    let fake_redis_port = fake_redis_backend.clone().run().await;
    let spec = RedisSpec {
        addresses: vec![format!("redis://127.0.0.1:{fake_redis_port}")],
        experimental_pub_sub_channel: Some(SUB_CHANNEL.into()),
        ..Default::default()
    };
    let store = RedisStore::new_standard(spec).await.expect("Working spec");
    fake_redis_backend.set_subscription_manager(store.subscription_manager().await.unwrap());

    // Manually set up the orphaned state in the fake backend:
    // 1. Add client_id → operation_id mapping (cid_* key)
    {
        let mut table = fake_redis_backend.table.lock().unwrap();
        let mut client_fields = HashMap::new();
        client_fields.insert(
            "data".into(),
            Value::BulkString(
                serde_json::to_string(&internal_operation_id)
                    .unwrap()
                    .into_bytes(),
            ),
        );
        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(),
        60,
    )
    .await
    .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(())
}

/// Regression test for the cid_* orphan accumulation bug.
///
/// Pre-fix, `add_action` wrote the `cid_<client_operation_id>` →
/// `operation_id` pointer with `expiry=None`. PR #2315 later started
/// attaching `retain_completed_for_s` TTL onto the matching `aa_*`
/// key on completion. The TTL mismatch meant: aa_* expired after
/// `retain_completed_for_s`, cid_* lingered forever. A subsequent
/// `WaitExecution` resolving the stale cid_* hit the orphan path,
/// returned `NotFound`, and the client (Bazel) restarted Execute,
/// creating *another* unbounded cid_*. In production this produced
/// ~3.8M cid_* keys with ~4.5% already orphaned.
///
/// This test asserts that `add_action` now writes the cid_* with a
/// bounded TTL. We don't pin the exact TTL value here (that's an
/// implementation detail of the fix), only that *some* TTL was
/// attached — exercising the path that previously passed `None`.
#[nativelink_test]
async fn add_action_attaches_ttl_to_cid_mapping() -> Result<(), Error> {
    const CLIENT_OPERATION_ID: &str = "cid_ttl_test_client";
    const WORKER_OPERATION_ID: &str = "cid_ttl_test_worker";
    const SUB_CHANNEL: &str = "sub_channel";

    let action_info = Arc::new(ActionInfo {
        command_digest: DigestInfo::zero_digest(),
        input_root_digest: DigestInfo::zero_digest(),
        timeout: Duration::from_secs(1),
        platform_properties: HashMap::new(),
        priority: 0,
        load_timestamp: SystemTime::UNIX_EPOCH,
        insert_timestamp: SystemTime::UNIX_EPOCH,
        unique_qualifier: ActionUniqueQualifier::Cacheable(ActionUniqueKey {
            instance_name: INSTANCE_NAME.to_string(),
            digest_function: DigestHasherFunc::Sha256,
            digest: DigestInfo::zero_digest(),
        }),
    });

    let fake_redis_backend: FakeRedisBackend<RedisSubscriptionManager> = FakeRedisBackend::new();
    let fake_redis_port = fake_redis_backend.clone().run().await;
    let spec = RedisSpec {
        addresses: vec![format!("redis://127.0.0.1:{fake_redis_port}")],
        experimental_pub_sub_channel: Some(SUB_CHANNEL.to_string()),
        ..Default::default()
    };
    let store = RedisStore::new_standard(spec).await.expect("Working spec");
    fake_redis_backend.set_subscription_manager(store.subscription_manager().await.unwrap());

    let notifier = Arc::new(Notify::new());
    let awaited_action_db = StoreAwaitedActionDb::new(
        store.clone(),
        notifier.clone(),
        MockInstantWrapped::default,
        move || WORKER_OPERATION_ID.into(),
        60,
    )
    .await
    .unwrap();

    let _subscription = awaited_action_db
        .add_action(
            CLIENT_OPERATION_ID.into(),
            action_info.clone(),
            Duration::from_mins(1),
        )
        .await
        .unwrap();

    // The cid_* key must have an EXPIRE attached. Pre-fix this map
    // was empty for the cid_ key — the assertion below would fail.
    let expiries = fake_redis_backend.expiries.lock().unwrap().clone();
    let cid_key = format!("cid_{CLIENT_OPERATION_ID}");
    let ttl = expiries.get(&cid_key).copied().unwrap_or_else(|| {
        panic!(
            "expected an EXPIRE on {cid_key} after add_action, but none was set. \
             All recorded expiries: {expiries:?}"
        )
    });
    assert!(
        ttl > 0,
        "cid_* TTL must be positive (got {ttl}); a 0 or negative TTL would \
         immediately evict the mapping and break in-flight WaitExecution calls"
    );

    Ok(())
}