delegation: add gc task to clean up abandoned contexts

[jayshrivastava]

This change addresses a TODO in the StageDelegation struct. Previously, it was possible for StageContexts to leak if a context was produced and not consumed (or if a waiter created a map entry and timed out without any corresponding producer call). This change introduces a GC task to clear contexts from the map if they have been there for longer than a certain TTL. The TTL proposed is 2 * wait_timeout.

When the StageDelegation is dropped, the GC task is terminated.

Testing

• Adds two unit tests (one for add_delegate_info and one for wait_for_delegate_info) which test that abandoned map entries are garbage collected appropriately.

[gabotechs]

Keep in mind that, as discussed in our private channels, the full delegation system is going to go away, so unfortunately this code will probably not reach the main branch.

However, this is a really good exercise, as this pattern of garbage collecting abandoned tasks is going to appear more times in different places across the project, so it's going to be very important to nail it down and find a solution as performant and robust as possible to this.

[gabotechs]

There should be no need for a Notify here, there's already tooling available in the tokio and DataFusion ecosystem for cancelling tasks on Drop, for example you could simply do

pub struct StageDelegation {
    stage_targets: Arc<DashMap<(String, usize), Value>>,
    wait_timeout: Duration,
    /// notify is used to shut down the garbage collection task when the StageDelegation is dropped.
    _task: SpawnedTask<()>,
}

impl Default for StageDelegation {
    fn default() -> Self {
        let stage_targets = Arc::new(DashMap::default());
        Self {
            stage_targets: stage_targets.clone(),
            wait_timeout: Duration::from_secs(5),
            _task: SpawnedTask::spawn(run_gc(
                stage_targets.clone(),
                Duration::from_secs(30), /* gc period */
            )),
        }
    }
}

// run_gc will continuously clear expired entries from the map, checking every `period`. The
// function terminates if `shutdown` is signalled.
async fn run_gc(stage_targets: Arc<DashMap<(String, usize), Value>>, period: Duration) {
    loop {
        tokio::time::sleep(period).await;
        stage_targets.retain(|_key, value| value.expiry.gt(&Instant::now()));
    }
}

And it should work the same

[gabotechs]

Looks like this const is unused?

[gabotechs]

.retain() iterates the whole map, so it do will end up locking the whole map unfortunately.

I think there should be ways of just garbage collecting abandoned tasks without iterating the full map, maybe spawning a task with a big tokio::sleep() at the beginning for each for each added entry that gets either cancelled due to a drop (the task was not abandoned) or completely executed included the tokio::sleep (the task was abandoned)

[NGA-TRAN]

Nice work, Jay. Big first step to add GC. Hopefully the time waiting for the new delegation in place, you will have sometime to polish this work.

[jayshrivastava]

Closing in favor of #92