Skip to content

Commit

Permalink
Reduce references to self.state_manager.inner (#1060)
Browse files Browse the repository at this point in the history
Removed unused test method and relocated notify function in `add_action`.
  • Loading branch information
adam-singer authored Jun 28, 2024
1 parent 5104778 commit 2eefa75
Show file tree
Hide file tree
Showing 2 changed files with 3 additions and 23 deletions.
4 changes: 3 additions & 1 deletion nativelink-scheduler/src/scheduler_state/state_manager.rs
Original file line number Diff line number Diff line change
Expand Up @@ -466,6 +466,7 @@ impl ClientStateManager for StateManager {
// Check to see if the action is running, if it is and cacheable, merge the actions.
if let Some(running_action) = self.inner.active_actions.get_mut(&action_info) {
self.inner.metrics.add_action_joined_running_action.inc();
self.inner.tasks_or_workers_change_notify.notify_one();
return Ok(Arc::new(ClientActionStateResult::new(
running_action.notify_channel.subscribe(),
)));
Expand Down Expand Up @@ -497,6 +498,7 @@ impl ClientStateManager for StateManager {
.queued_actions
.insert(arc_action_info.clone(), queued_action);
self.inner.queued_actions_set.insert(arc_action_info);
self.inner.tasks_or_workers_change_notify.notify_one();
return Ok(result);
}

Expand Down Expand Up @@ -525,7 +527,7 @@ impl ClientStateManager for StateManager {
worker_id: None,
},
);

self.inner.tasks_or_workers_change_notify.notify_one();
return Ok(Arc::new(ClientActionStateResult::new(rx)));
}

Expand Down
22 changes: 0 additions & 22 deletions nativelink-scheduler/src/simple_scheduler.rs
Original file line number Diff line number Diff line change
Expand Up @@ -88,10 +88,6 @@ impl SimpleSchedulerImpl {
action_info: ActionInfo,
) -> Result<watch::Receiver<Arc<ActionState>>, Error> {
let add_action_result = self.state_manager.add_action(action_info).await?;
self.state_manager
.inner
.tasks_or_workers_change_notify
.notify_one();
add_action_result.as_receiver().await.cloned()
}

Expand Down Expand Up @@ -484,24 +480,6 @@ impl SimpleScheduler {
.contains(worker_id)
}

/// Checks to see if the worker can accept work. Should only be used in unit tests.
pub async fn can_worker_accept_work_for_test(
&self,
worker_id: &WorkerId,
) -> Result<bool, Error> {
let mut inner = self.get_inner_lock().await;
let worker = inner
.state_manager
.inner
.workers
.workers
.get_mut(worker_id)
.ok_or_else(|| {
make_input_err!("WorkerId '{}' does not exist in workers map", worker_id)
})?;
Ok(worker.can_accept_work())
}

/// A unit test function used to send the keep alive message to the worker from the server.
pub async fn send_keep_alive_to_worker_for_test(
&self,
Expand Down

0 comments on commit 2eefa75

Please sign in to comment.