asdf

Author: jayshrivastava

src/stage_delegation/delegation.rs

@@ -2,8 +2,13 @@ use super::StageContext;
 use dashmap::{DashMap, Entry};
 use datafusion::common::{exec_datafusion_err, exec_err};
 use datafusion::error::DataFusionError;
-use std::time::Duration;
+use std::ops::Add;
+use std::sync::Arc;
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
 use tokio::sync::oneshot;
+use tokio::sync::Notify;
+use tokio::time;
+use tokio::time::Instant;
 
 /// In each stage of the distributed plan, there will be N workers. All these workers
 /// need to coordinate to pull data from the next stage, which will contain M workers.
@@ -31,20 +36,67 @@ use tokio::sync::oneshot;
 /// [oneshot::Sender], and listen on the other end of the channel [oneshot::Receiver] for
 /// the delegate to put something there.
 pub struct StageDelegation {
-    stage_targets: DashMap<(String, usize), Oneof>,
+    stage_targets: Arc<DashMap<(String, usize), Value>>,
     wait_timeout: Duration,
+
+    notify: Arc<Notify>,
 }
 
 impl Default for StageDelegation {
     fn default() -> Self {
-        Self {
-            stage_targets: DashMap::default(),
+        let stage_targets = Arc::new(DashMap::default());
+        let notify = Arc::new(Notify::new());
+
+        let result = Self {
+            stage_targets: stage_targets.clone(),
             wait_timeout: Duration::from_secs(5),
+
+            notify: notify.clone(),
+        };
+
+        tokio::spawn(run_gc_async(
+            stage_targets.clone(),
+            notify.clone(),
+            Duration::from_secs(30), /* gc period */
+        ));
+
+        result
+    }
+}
+
+// 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(
+    stage_targets: Arc<DashMap<(String, usize), Value>>,
+    shutdown: Arc<Notify>,
+    period: Duration,
+) {
+    loop {
+        tokio::select! {
+            _ = shutdown.notified() => {
+                break;
+            }
+            _ = tokio::time::sleep(period) => {
+               // PERF: This iterator is sharded, so it won't lock the whole map.
+                stage_targets.retain(|_key, value| {
+                  value.expiry.gt(&Instant::now())
+                  });
+            }
         }
     }
 }
 
+impl Drop for StageDelegation {
+    fn drop(&mut self) {
+        self.notify.notify_one();
+    }
+}
+
 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
@@ -57,9 +109,13 @@ impl StageDelegation {
         actor_idx: usize,
         next_stage_context: StageContext,
     ) -> Result<(), DataFusionError> {
+        let now = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_secs();
         let tx = match self.stage_targets.entry((stage_id, actor_idx)) {
-            Entry::Occupied(entry) => match entry.get() {
-                Oneof::Sender(_) => match entry.remove() {
+            Entry::Occupied(entry) => match entry.get().value {
+                Oneof::Sender(_) => match entry.remove().value {
                     Oneof::Sender(tx) => tx,
                     Oneof::Receiver(_) => unreachable!(),
                 },
@@ -69,7 +125,11 @@ impl StageDelegation {
             },
             Entry::Vacant(entry) => {
                 let (tx, rx) = oneshot::channel();
-                entry.insert(Oneof::Receiver(rx));
+                entry.insert(Value {
+                    // Use 2 * the waiter wait duration for now.
+                    expiry: Instant::now().add(self.gc_ttl()),
+                    value: Oneof::Receiver(rx),
+                });
                 tx
             }
         };
@@ -95,16 +155,25 @@ impl StageDelegation {
         actor_idx: usize,
     ) -> Result<StageContext, DataFusionError> {
         let rx = match self.stage_targets.entry((stage_id.clone(), actor_idx)) {
-            Entry::Occupied(entry) => match entry.get() {
-                Oneof::Sender(_) => return exec_err!("Programming error: while waiting for delegate info the entry in the StageDelegation target map cannot be a Sender"),
-                Oneof::Receiver(_) => match entry.remove() {
+            Entry::Occupied(entry) => match entry.get().value {
+                Oneof::Sender(_) => {
+                    return exec_err!(
+                        "Programming error: while waiting for delegate info the entry in the \
+                    StageDelegation target map cannot be a Sender"
+                    )
+                }
+                Oneof::Receiver(_) => match entry.remove().value {
                     Oneof::Sender(_) => unreachable!(),
-                    Oneof::Receiver(rx) => rx
+                    Oneof::Receiver(rx) => rx,
                 },
             },
             Entry::Vacant(entry) => {
                 let (tx, rx) = oneshot::channel();
-                entry.insert(Oneof::Sender(tx));
+                entry.insert(Value {
+                    // Use 2 * the waiter wait duration for now.
+                    expiry: Instant::now().add(self.gc_ttl()),
+                    value: Oneof::Sender(tx),
+                });
                 rx
             }
         };
@@ -120,6 +189,11 @@ impl StageDelegation {
     }
 }
 
+struct Value {
+    expiry: Instant,
+    value: Oneof,
+}
+
 enum Oneof {
     Sender(oneshot::Sender<StageContext>),
     Receiver(oneshot::Receiver<StageContext>),
@@ -129,6 +203,7 @@ enum Oneof {
 mod tests {
     use super::*;
     use crate::stage_delegation::StageContext;
+    use futures::TryFutureExt;
     use std::sync::Arc;
     use uuid::Uuid;
 
@@ -222,6 +297,7 @@ mod tests {
 
         let received_context = wait_task1.await.unwrap().unwrap();
         assert_eq!(received_context.id, stage_context.id);
+        assert_eq!(0, delegation.stage_targets.len())
     }
 
     #[tokio::test]
@@ -287,4 +363,93 @@ mod tests {
             .unwrap();
         assert_eq!(received_context, stage_context);
     }
+
+    #[tokio::test]
+    async fn test_waiter_timeout_and_gc_cleanup() {
+        let stage_targets = Arc::new(DashMap::default());
+        let shutdown = Arc::new(Notify::new());
+        let delegation = StageDelegation {
+            stage_targets: stage_targets.clone(),
+            wait_timeout: Duration::from_millis(1),
+            notify: shutdown.clone(),
+        };
+        let stage_id = Uuid::new_v4().to_string();
+
+        // Actor waits but times out
+        let result = delegation.wait_for_delegate_info(stage_id, 0).await;
+
+        assert!(result.is_err());
+        assert!(result.unwrap_err().to_string().contains("Timeout"));
+
+        // 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(
+            stage_targets.clone(),
+            shutdown.clone(),
+            Duration::from_millis(5),
+        ));
+
+        // Wait for GC to clear the map
+        for _ in 0..50 {
+            tokio::time::sleep(Duration::from_millis(10)).await;
+            if stage_targets.len() == 0 {
+                break;
+            }
+        }
+
+        // Stop GC by dropping the delegation. Assert that it has shutdown.
+        drop(delegation);
+        gc_task.await.unwrap();
+
+        // After GC, map should be cleared
+        assert_eq!(stage_targets.len(), 0);
+    }
+
+    #[tokio::test]
+    async fn test_writer_only_and_gc_cleanup() {
+        let stage_targets = Arc::new(DashMap::default());
+        let shutdown = Arc::new(Notify::new());
+        let delegation = StageDelegation {
+            stage_targets: stage_targets.clone(),
+            wait_timeout: Duration::from_millis(1),
+            notify: shutdown.clone(),
+        };
+        let stage_id = Uuid::new_v4().to_string();
+        let stage_context = create_test_stage_context();
+
+        // Writer adds info without anyone waiting
+        let result = delegation.add_delegate_info(stage_id, 0, stage_context);
+
+        assert!(result.is_ok());
+
+        // Entry should be in map
+        assert_eq!(stage_targets.len(), 1);
+
+        // Wait for expiry time to pass (gc_ttl is 2 * wait_timeout)
+        tokio::time::sleep(delegation.gc_ttl()).await;
+
+        // Run GC to cleanup expired entries
+        let gc_task = tokio::spawn(run_gc_async(
+            stage_targets.clone(),
+            shutdown.clone(),
+            Duration::from_millis(10),
+        ));
+
+        // Wait for GC to clear the map
+        for _ in 0..50 {
+            tokio::time::sleep(Duration::from_millis(20)).await;
+            if stage_targets.len() == 0 {
+                break;
+            }
+        }
+
+        // Stop GC
+        drop(delegation);
+        gc_task.await.unwrap();
+
+        // After GC, map should be cleared
+        assert_eq!(stage_targets.len(), 0);
+    }
 }