Add safe request timeout for running actions manager

TraceMachina · Mar 10, 2024 · c345f7a · c345f7a
1 parent a651f2c
commit c345f7a
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/nativelink-worker/src/running_actions_manager.rs b/nativelink-worker/src/running_actions_manager.rs
@@ -673,6 +673,12 @@ impl RunningActionImpl {
             .current_dir(format!("{}/{}", self.work_directory, command_proto.working_directory))
             .env_clear();
 
+        let requested_timeout = if self.action_info.timeout.is_zero() {
+            self.running_actions_manager.max_action_timeout
+        } else {
+            self.action_info.timeout
+        };
+
         let mut maybe_side_channel_file: Option<Cow<'_, OsStr>> = None;
         if let Some(additional_environment) = &self
             .running_actions_manager
@@ -688,7 +694,7 @@ impl RunningActionImpl {
                         .get(property)
                         .map_or_else(|| Cow::Borrowed(""), |v| v.as_str()),
                     EnvironmentSource::value(value) => Cow::Borrowed(value.as_str()),
-                    EnvironmentSource::timeout_millis => Cow::Owned(self.timeout.as_millis().to_string()),
+                    EnvironmentSource::timeout_millis => Cow::Owned(requested_timeout.as_millis().to_string()),
                     EnvironmentSource::side_channel_file => {
                         let file_cow = format!(
                             "{}/{}/{}",
@@ -833,7 +839,7 @@ impl RunningActionImpl {
                     };
 
                     let maybe_error_override = if let Some(side_channel_file) = maybe_side_channel_file {
-                        process_side_channel_file(side_channel_file.clone(), &args, self.timeout).await
+                        process_side_channel_file(side_channel_file.clone(), &args, requested_timeout).await
                         .err_tip(|| format!("Error processing side channel file: {side_channel_file:?}"))?
                     } else {
                         None
@@ -1604,7 +1610,7 @@ impl RunningActionsManager for RunningActionsManagerImpl {
                     output_upload_start_timestamp: SystemTime::UNIX_EPOCH,
                     output_upload_completed_timestamp: SystemTime::UNIX_EPOCH,
                 };
-                let timeout = if action_info.timeout == Duration::ZERO || self.timeout_handled_externally {
+                let timeout = if action_info.timeout.is_zero() || self.timeout_handled_externally {
                     self.max_action_timeout
                 } else {
                     action_info.timeout