asdf

Author: jayshrivastava

src/stage_delegation/delegation.rs

@@ -35,10 +35,15 @@ use tokio::time::Instant;
 /// On 2, the `wait_for_delegate_info` call will create an entry in the [DashMap] with a
 /// [oneshot::Sender], and listen on the other end of the channel [oneshot::Receiver] for
 /// the delegate to put something there.
+///
+/// It's possible for [StageContext] to "get lost" if `add_delegate_info` is called without
+/// a corresponding call to `wait_for_delegate_info` or vice versa. In this case, a task will
+/// reap any contexts that live for longer than the `gc_ttl`.
 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.
     notify: Arc<Notify>,
 }
 
@@ -50,11 +55,11 @@ impl Default for StageDelegation {
         let result = Self {
             stage_targets: stage_targets.clone(),
             wait_timeout: Duration::from_secs(5),
-
             notify: notify.clone(),
         };
 
-        tokio::spawn(run_gc_async(
+        // Run the GC task.
+        tokio::spawn(run_gc(
             stage_targets.clone(),
             notify.clone(),
             Duration::from_secs(30), /* gc period */
@@ -64,10 +69,11 @@ impl Default for StageDelegation {
     }
 }
 
-// gc_interval is the period over which gc runs to purge old stage_targets entries which will
-// never be read. This may happen if the actor encounters an error before it can read
-// the delagate info.
-async fn run_gc_async(
+const GC_PERIOD_SECONDS: usize = 30;
+
+// 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>>,
     shutdown: Arc<Notify>,
     period: Duration,
@@ -78,7 +84,7 @@ async fn run_gc_async(
                 break;
             }
             _ = tokio::time::sleep(period) => {
-               // PERF: This iterator is sharded, so it won't lock the whole map.
+               // Performance: This iterator is sharded, so it won't lock the whole map.
                 stage_targets.retain(|_key, value| {
                   value.expiry.gt(&Instant::now())
                   });
@@ -94,9 +100,6 @@ impl Drop for StageDelegation {
 }
 
 impl StageDelegation {
-    fn gc_ttl(&self) -> Duration {
-        self.wait_timeout * 2
-    }
     /// Puts the [StageContext] info so that an actor can pick it up with `wait_for_delegate_info`.
     ///
     /// - If the actor was already waiting for this info, it just puts it on the
@@ -126,20 +129,13 @@ impl StageDelegation {
             Entry::Vacant(entry) => {
                 let (tx, rx) = oneshot::channel();
                 entry.insert(Value {
-                    // Use 2 * the waiter wait duration for now.
                     expiry: Instant::now().add(self.gc_ttl()),
                     value: Oneof::Receiver(rx),
                 });
                 tx
             }
         };
 
-        // TODO: `send` does not wait for the other end of the channel to receive the message,
-        //  so if nobody waits for it, we might leak an entry in `stage_targets` that will never
-        //  be cleaned up. We can either:
-        //  1. schedule a cleanup task that iterates the entries cleaning up old ones
-        //  2. find some other API that allows us to .await until the other end receives the message,
-        //     and on a timeout, cleanup the entry anyway.
         tx.send(next_stage_context)
             .map_err(|_| exec_datafusion_err!("Could not send stage context info"))
     }
@@ -170,7 +166,6 @@ impl StageDelegation {
             Entry::Vacant(entry) => {
                 let (tx, rx) = oneshot::channel();
                 entry.insert(Value {
-                    // Use 2 * the waiter wait duration for now.
                     expiry: Instant::now().add(self.gc_ttl()),
                     value: Oneof::Sender(tx),
                 });
@@ -187,6 +182,12 @@ impl StageDelegation {
                 )
             })
     }
+
+    // gc_ttl is used to set the expiry of elements in the map. Use 2 * the waiter wait duration
+    // to avoid running gc too early.
+    fn gc_ttl(&self) -> Duration {
+        self.wait_timeout * 2
+    }
 }
 
 struct Value {
@@ -384,15 +385,15 @@ mod tests {
         // Wait for expiry time to pass.
         tokio::time::sleep(delegation.gc_ttl()).await;
 
-        // Run GC to cleanup expired entries
-        let gc_task = tokio::spawn(run_gc_async(
+        // Run GC to clean up expired entries
+        let gc_task = tokio::spawn(run_gc(
             stage_targets.clone(),
             shutdown.clone(),
             Duration::from_millis(5),
         ));
 
         // Wait for GC to clear the map
-        for _ in 0..50 {
+        for _ in 0..10 {
             tokio::time::sleep(Duration::from_millis(10)).await;
             if stage_targets.len() == 0 {
                 break;
@@ -403,7 +404,7 @@ mod tests {
         drop(delegation);
         gc_task.await.unwrap();
 
-        // After GC, map should be cleared
+        // After GC, map should be cleared.
         assert_eq!(stage_targets.len(), 0);
     }
 
@@ -431,21 +432,21 @@ mod tests {
         tokio::time::sleep(delegation.gc_ttl()).await;
 
         // Run GC to cleanup expired entries
-        let gc_task = tokio::spawn(run_gc_async(
+        let gc_task = tokio::spawn(run_gc(
             stage_targets.clone(),
             shutdown.clone(),
             Duration::from_millis(10),
         ));
 
         // Wait for GC to clear the map
-        for _ in 0..50 {
+        for _ in 0..10 {
             tokio::time::sleep(Duration::from_millis(20)).await;
             if stage_targets.len() == 0 {
                 break;
             }
         }
 
-        // Stop GC
+        // Stop GC.
         drop(delegation);
         gc_task.await.unwrap();