Merge branch 'main' into remote-persistent-safe

Author: MarcusSorealheis

Cargo.lock

@@ -10,7 +10,7 @@ version = "0.7.6"
 nativelink-error = { path = "../nativelink-error" }
 
 byte-unit = { version = "5.1.6", default-features = false, features = ["byte"] }
-humantime = { version = "2.2.0", default-features = false }
+humantime = { version = "2.3.0", default-features = false }
 rand = { version = "0.9.0", default-features = false, features = [
   "thread_rng",
 ] }

nativelink-config/Cargo.toml

@@ -107,6 +107,7 @@ impl AwaitedAction {
             // client_operation_id to all clients.
             client_operation_id: operation_id.clone(),
             action_digest: action_info.unique_qualifier.digest(),
+            last_transition_timestamp: now,
         });
 
         let ctx = Context::current();

nativelink-scheduler/src/awaited_action_db/awaited_action.rs

@@ -14,6 +14,7 @@
 
 use std::collections::HashMap;
 use std::sync::Arc;
+use std::time::SystemTime;
 
 use async_trait::async_trait;
 use nativelink_error::{Code, Error, ResultExt, make_err};
@@ -267,6 +268,7 @@ impl CacheLookupScheduler {
                         client_operation_id: OperationId::default(),
                         stage: ActionStage::CompletedFromCache(action_result),
                         action_digest: action_info.unique_qualifier.digest(),
+                        last_transition_timestamp: SystemTime::now(),
                     };
 
                     let ctx = Context::current();

nativelink-scheduler/src/cache_lookup_scheduler.rs

@@ -12,11 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::{BTreeSet, HashMap};
 use std::sync::Arc;
 use std::time::{Instant, SystemTime};
 
 use async_trait::async_trait;
-use futures::Future;
+use futures::{Future, StreamExt, future};
 use nativelink_config::schedulers::SimpleSpec;
 use nativelink_error::{Code, Error, ResultExt};
 use nativelink_metric::{MetricsComponent, RootMetricsComponent};
@@ -38,8 +39,7 @@ use opentelemetry::context::{Context, FutureExt as OtelFutureExt};
 use opentelemetry_semantic_conventions::attribute::ENDUSER_ID;
 use tokio::sync::{Notify, mpsc};
 use tokio::time::Duration;
-use tokio_stream::StreamExt;
-use tracing::{error, info_span};
+use tracing::{error, info, info_span};
 
 use crate::api_worker_scheduler::ApiWorkerScheduler;
 use crate::awaited_action_db::{AwaitedActionDb, CLIENT_KEEPALIVE_DURATION};
@@ -627,8 +627,7 @@ impl SimpleScheduler {
                         tokio::pin!(task_change_fut);
                         tokio::pin!(worker_change_fut);
                         // Wait for either of these futures to be ready.
-                        let state_changed =
-                            futures::future::select(task_change_fut, worker_change_fut);
+                        let state_changed = future::select(task_change_fut, worker_change_fut);
                         if last_match_successful {
                             let _ = state_changed.await;
                         } else {
@@ -638,7 +637,7 @@ impl SimpleScheduler {
                             // hard loop if there's something wrong inside do_try_match.
                             let sleep_fut = tokio::time::sleep(Duration::from_millis(100));
                             tokio::pin!(sleep_fut);
-                            let _ = futures::future::select(state_changed, sleep_fut).await;
+                            let _ = future::select(state_changed, sleep_fut).await;
                         }
 
                         let result = match weak_inner.upgrade() {
@@ -656,6 +655,68 @@ impl SimpleScheduler {
 
                                 let res = scheduler.do_try_match(full_worker_logging).await;
                                 if full_worker_logging {
+                                    let operations_stream = scheduler
+                                        .matching_engine_state_manager
+                                        .filter_operations(OperationFilter::default())
+                                        .await
+                                        .err_tip(|| "In action_scheduler getting filter result");
+
+                                    let mut oldest_actions_in_state: HashMap<
+                                        String,
+                                        BTreeSet<Arc<ActionState>>,
+                                    > = HashMap::new();
+                                    let max_items = 5;
+
+                                    match operations_stream {
+                                        Ok(stream) => {
+                                            let actions = stream
+                                                .filter_map(|item| async move {
+                                                    match item.as_ref().as_state().await {
+                                                        Ok((action_state, _origin_metadata)) => {
+                                                            Some(action_state)
+                                                        }
+                                                        Err(e) => {
+                                                            error!(
+                                                                ?e,
+                                                                "Failed to get action state!"
+                                                            );
+                                                            None
+                                                        }
+                                                    }
+                                                })
+                                                .collect::<Vec<_>>()
+                                                .await;
+                                            for action_state in actions.iter() {
+                                                let name = action_state.stage.name();
+                                                match oldest_actions_in_state.get_mut(&name) {
+                                                    Some(values) => {
+                                                        values.insert(action_state.clone());
+                                                        if values.len() > max_items {
+                                                            values.pop_first();
+                                                        }
+                                                    }
+                                                    None => {
+                                                        let mut values = BTreeSet::new();
+                                                        values.insert(action_state.clone());
+                                                        oldest_actions_in_state
+                                                            .insert(name, values);
+                                                    }
+                                                };
+                                            }
+                                        }
+                                        Err(e) => {
+                                            error!(?e, "Failed to get operations list!");
+                                        }
+                                    }
+
+                                    for value in oldest_actions_in_state.values() {
+                                        let mut items = vec![];
+                                        for item in value {
+                                            items.push(item.to_string());
+                                        }
+                                        info!(?items, "Oldest actions in state");
+                                    }
+
                                     worker_match_logging_last.replace(now);
                                 }
                                 res

nativelink-scheduler/src/simple_scheduler.rs

@@ -320,9 +320,8 @@ where
         // Note: The caller must filter `client_operation_id`.
 
         let mut maybe_reloaded_awaited_action: Option<AwaitedAction> = None;
-        if awaited_action.last_client_keepalive_timestamp() + self.client_action_timeout
-            < (self.now_fn)().now()
-        {
+        let now = (self.now_fn)().now();
+        if awaited_action.last_client_keepalive_timestamp() + self.client_action_timeout < now {
             // This may change if the version is out of date.
             let mut timed_out = true;
             if !awaited_action.state().stage.is_finished() {
@@ -335,6 +334,7 @@ where
                     )),
                     ..ActionResult::default()
                 });
+                state.last_transition_timestamp = now;
                 let state = Arc::new(state);
                 // We may be competing with an client timestamp update, so try
                 // this a few times.
@@ -655,6 +655,7 @@ where
                     // correct client id.
                     client_operation_id: operation_id.clone(),
                     action_digest: awaited_action.action_info().digest(),
+                    last_transition_timestamp: now,
                 }),
                 now,
             );

nativelink-scheduler/src/simple_scheduler_state_manager.rs

@@ -43,6 +43,7 @@ async fn action_state_any_url_test() -> Result<(), Error> {
         // Result is only populated if has_action_result.
         stage: ActionStage::Completed(ActionResult::default()),
         action_digest,
+        last_transition_timestamp: SystemTime::now(),
     };
     let operation: Operation = action_state.as_operation(client_id);
 

nativelink-scheduler/tests/action_messages_test.rs

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 use std::sync::Arc;
-use std::time::UNIX_EPOCH;
+use std::time::{SystemTime, UNIX_EPOCH};
 
 mod utils {
     pub(crate) mod scheduler_utils;
@@ -71,6 +71,7 @@ async fn add_action_handles_skip_cache() -> Result<(), Error> {
             client_operation_id: OperationId::default(),
             stage: ActionStage::Queued,
             action_digest: action_info.unique_qualifier.digest(),
+            last_transition_timestamp: SystemTime::now(),
         }));
     let ActionUniqueQualifier::Cacheable(action_key) = action_info.unique_qualifier.clone() else {
         panic!("This test should be testing when item was cached first");

nativelink-scheduler/tests/cache_lookup_scheduler_test.rs

@@ -14,7 +14,7 @@
 
 use std::collections::HashMap;
 use std::sync::Arc;
-use std::time::UNIX_EPOCH;
+use std::time::{SystemTime, UNIX_EPOCH};
 
 mod utils {
     pub(crate) mod scheduler_utils;
@@ -70,6 +70,7 @@ async fn add_action_adds_property() -> Result<(), Error> {
             client_operation_id: OperationId::default(),
             stage: ActionStage::Queued,
             action_digest: action_info.unique_qualifier.digest(),
+            last_transition_timestamp: SystemTime::now(),
         }));
     let client_operation_id = OperationId::default();
     let (_, (passed_client_operation_id, action_info)) = join!(
@@ -114,6 +115,7 @@ async fn add_action_overwrites_property() -> Result<(), Error> {
             client_operation_id: OperationId::default(),
             stage: ActionStage::Queued,
             action_digest: action_info.unique_qualifier.digest(),
+            last_transition_timestamp: SystemTime::now(),
         }));
     let client_operation_id = OperationId::default();
     let (_, (passed_client_operation_id, action_info)) = join!(
@@ -153,6 +155,7 @@ async fn add_action_property_added_after_remove() -> Result<(), Error> {
             client_operation_id: OperationId::default(),
             stage: ActionStage::Queued,
             action_digest: action_info.unique_qualifier.digest(),
+            last_transition_timestamp: SystemTime::now(),
         }));
     let client_operation_id = OperationId::default();
     let (_, (passed_client_operation_id, action_info)) = join!(
@@ -192,6 +195,7 @@ async fn add_action_property_remove_after_add() -> Result<(), Error> {
             client_operation_id: OperationId::default(),
             stage: ActionStage::Queued,
             action_digest: action_info.unique_qualifier.digest(),
+            last_transition_timestamp: SystemTime::now(),
         }));
     let client_operation_id = OperationId::default();
     let (_, (passed_client_operation_id, action_info)) = join!(
@@ -233,6 +237,7 @@ async fn add_action_property_replace() -> Result<(), Error> {
             client_operation_id: OperationId::default(),
             stage: ActionStage::Queued,
             action_digest: action_info.unique_qualifier.digest(),
+            last_transition_timestamp: SystemTime::now(),
         }));
     let client_operation_id = OperationId::default();
     let (_, (passed_client_operation_id, action_info)) = join!(
@@ -277,6 +282,7 @@ async fn add_action_property_replace_match_value() -> Result<(), Error> {
             client_operation_id: OperationId::default(),
             stage: ActionStage::Queued,
             action_digest: action_info.unique_qualifier.digest(),
+            last_transition_timestamp: SystemTime::now(),
         }));
     let client_operation_id = OperationId::default();
     let (_, (passed_client_operation_id, action_info)) = join!(
@@ -322,6 +328,7 @@ async fn add_action_property_replace_value() -> Result<(), Error> {
             client_operation_id: OperationId::default(),
             stage: ActionStage::Queued,
             action_digest: action_info.unique_qualifier.digest(),
+            last_transition_timestamp: SystemTime::now(),
         }));
     let client_operation_id = OperationId::default();
     let (_, (passed_client_operation_id, action_info)) = join!(
@@ -359,6 +366,7 @@ async fn add_action_property_remove() -> Result<(), Error> {
             client_operation_id: OperationId::default(),
             stage: ActionStage::Queued,
             action_digest: action_info.unique_qualifier.digest(),
+            last_transition_timestamp: SystemTime::now(),
         }));
     // let platform_property_manager = Arc::new(PlatformPropertyManager::new(HashMap::new()));
     let client_operation_id = OperationId::default();

nativelink-scheduler/tests/property_modifier_scheduler_test.rs

@@ -489,6 +489,7 @@ async fn add_action_smoke_test() -> Result<(), Error> {
         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
     };