StateManager will now cleanup actions on client disconnect #1107

allada · 2024-07-08T04:58:27Z

StateManager will now properly remove items from the maps if the client disconnects after a set amount of time. Currently these values are hard codded, but will be easy to transition them to use config variables once we design it out.

This change is

allada

+@zbirenbaum

Reviewable status: 0 of 1 LGTMs obtained, and 1 discussions need to be resolved (waiting on @zbirenbaum)

nativelink-scheduler/src/scheduler_state/workers.rs line 165 at r1 (raw file):

                .err_tip(|| "in update_operation on SimpleScheduler::update_action");
            if let Err(err) = update_operation_res {
                let result = Result::<(), _>::Err(err.clone())

note: A worker can now send updates about an action that the state manager does not know about. This is ok, since the Workers struct now owns the platform properties and resource management of the node instead of state manager.

allada

Reviewable status: 0 of 1 LGTMs obtained, and pending CI: vale, and 1 discussions need to be resolved (waiting on @zbirenbaum)

nativelink-scheduler/src/simple_scheduler.rs line 57 at r2 (raw file):

/// Default time in seconds before a client is evicted.
// TODO!(make this a config and documented)

fyi: My system is to use todo! anywhere that we need to cleanup code before merging into main.

adam-singer

Reviewed 8 of 12 files at r1, 4 of 4 files at r2, all commit messages.
Reviewable status: 0 of 1 LGTMs obtained, and 8 discussions need to be resolved (waiting on @zbirenbaum)

nativelink-scheduler/src/scheduler_state/state_manager.rs line 48 at r2 (raw file):

/// How often the owning database will have the AwaitedAction touched
/// to keep it from being evicted.
const KEEPALIVE_TIMER: Duration = Duration::from_secs(10);

nit: its not really a timer, MAYBE KEEPALIVE_GRACE_PERIOD KEEPALIVE_DURATION ?

nativelink-scheduler/src/scheduler_state/state_manager.rs line 97 at r2 (raw file):

    // ClientAwaitedAction is dropped, but will never actually be
    // None except during the drop.
    client_operation_id: Option<ClientOperationId>,

Would be it overly pedantic to build a custom type for holding a client_operation_id that allows for it to be set once and dropped? That way we are not implying state around Option in corner conditions, also Option doesn't encode if it was set multiple times?

pub struct SetVar<T> {
    value: Option<T>,
    set_count: u8,
    set_count_max: u8,
}

impl<T> SetVar<T> {
    pub fn new(set_count_max: u8) -> Self {
        Self {
            value: None,
            set_count: 0,
            set_count_max: set_count_max
        }
    }

    pub fn set(&mut self, value: T) -> Result<(), &'static str> {
        if self.set_count_max == self.set_count {
           Err("Value can only be set once")
        } else {
            self.value = Some(value);
            self.set_count += 1;
            Ok(())
        }
    }

    pub fn drop_value(&mut self) -> Result<(), &'static str> {
        if self.set_count >= 1 {
            self.value = None;
            self.set_count = self.set_count_max;
            Ok(())
        } else if self.set_count == 0 {
            Err("Value has not been set yet")
        } else {
            Err("Value has already been dropped")
        }
    }

    pub fn get(&self) -> Option<&T> {
        self.value.as_ref()
    }

    pub fn is_set(&self) -> bool {
        self.set_count >= 1
    }
}

// Example:
let client_operation_id = SetVar<ClientOperationId>::new(1).set(a_client_operation_id);

nativelink-scheduler/src/scheduler_state/state_manager.rs line 140 at r2 (raw file):

}

/// Trit to be able to use the [`EvictingMap`] with [`ClientAwaitedAction`].

nit: Trit?

nativelink-scheduler/src/scheduler_state/state_manager.rs line 179 at r2 (raw file):

#[allow(clippy::mutable_key_type)]
impl AwaitedActionDb {
    /// Touches the client operation id to keep it from being evicted.

nit: /// Refreshes/Updates the time to live of the [`ClientOperationId`] in the [`EvictionMap`] by touching the key. ?

nativelink-scheduler/src/scheduler_state/state_manager.rs line 209 at r2 (raw file):

        if !did_remove {
            event!(
                Level::ERROR,

nit: Leave a simple comment to downgrade this event Level if it shows to be noisy and not an error?

nativelink-scheduler/src/scheduler_state/state_manager.rs line 254 at r2 (raw file):

        let sort_info = awaited_action.get_sort_info();
        let sort_key = sort_info.get_previous_sort_key();
        let btree = self.get_sort_map_for_state(&awaited_action.get_current_state().stage);

nit: sorted_awaited_action_btree

nativelink-scheduler/src/scheduler_state/state_manager.rs line 523 at r2 (raw file):

                        client_operation_to_awaited_action: EvictingMap::new(
                            config,
                            SystemTime::now(),

nit: For better control on testing would it make sense to pass in SystemTime::now() in StateManager::new() since could be a fixed / mocked type later? Puts the burden on the caller tho generally not having control of time under test tends to be a pita.

allada

Reviewable status: 0 of 1 LGTMs obtained, and 3 discussions need to be resolved (waiting on @zbirenbaum)

nativelink-scheduler/src/scheduler_state/state_manager.rs line 48 at r2 (raw file):