fix: potential deadlock in async data movement window manager implementation (#2755)

Signed-off-by: Yashash H L <yashashhl25@gmail.com>

Author: yhl25

rust/numaflow-core/src/reduce/reducer/aligned/windower/fixed.rs

@@ -70,14 +70,14 @@ impl FixedWindowManager {
             // Window exists, append message
             AlignedWindowOperation::Append {
                 message: msg,
-                window: window.clone(),
+                window,
             }
         } else {
             // New window, insert it
             active_windows.insert(window.clone());
             AlignedWindowOperation::Open {
                 message: msg,
-                window: window.clone(),
+                window,
             }
         };
 
@@ -93,36 +93,41 @@ impl FixedWindowManager {
     pub(crate) fn close_windows(&self, watermark: DateTime<Utc>) -> Vec<AlignedWindowMessage> {
         let mut result = Vec::new();
 
-        let mut active_windows = self
-            .active_windows
-            .write()
-            .expect("Poisoned lock for active_windows");
-        let mut closed_windows = self
-            .closed_windows
-            .write()
-            .expect("Poisoned lock for closed_windows");
-
-        let mut windows_to_close = Vec::new();
-
-        for window in active_windows.iter() {
-            if window.end_time <= watermark {
-                // Create close message
-                let window_msg = AlignedWindowMessage {
-                    operation: AlignedWindowOperation::Close {
-                        window: window.clone(),
-                    },
-                    pnf_slot: window_pnf_slot(window),
-                };
+        let windows_to_close = {
+            let mut active_windows = self
+                .active_windows
+                .write()
+                .expect("Poisoned lock for active_windows");
+
+            let mut windows_to_close = Vec::new();
+            active_windows.retain(|window| {
+                if window.end_time <= watermark {
+                    windows_to_close.push(window.clone());
+                    false
+                } else {
+                    true
+                }
+            });
+
+            windows_to_close
+        };
 
-                result.push(window_msg);
-                windows_to_close.push(window.clone());
+        // add the windows to closed_windows
+        {
+            let mut closed_windows = self
+                .closed_windows
+                .write()
+                .expect("Poisoned lock for closed_windows");
+            for window in &windows_to_close {
+                closed_windows.insert(window.clone());
             }
         }
 
-        // Move windows from active to closed
         for window in windows_to_close {
-            active_windows.remove(&window);
-            closed_windows.insert(window);
+            result.push(AlignedWindowMessage {
+                pnf_slot: window_pnf_slot(&window),
+                operation: AlignedWindowOperation::Close { window },
+            });
         }
 
         result
@@ -142,20 +147,22 @@ impl FixedWindowManager {
         // get the oldest window from closed_windows, if closed_windows is empty, get the oldest
         // from active_windows
         // NOTE: closed windows will always have a lower end time than active_windows
-        self.closed_windows
+        {
+            let closed_windows = self
+                .closed_windows
+                .read()
+                .expect("Poisoned lock for closed_windows");
+            if let Some(window) = closed_windows.iter().next() {
+                return Some(window.clone());
+            }
+        }
+
+        self.active_windows
             .read()
-            .expect("Poisoned lock for closed_windows")
+            .expect("Poisoned lock for active_windows")
             .iter()
             .next()
             .cloned()
-            .or_else(|| {
-                self.active_windows
-                    .read()
-                    .expect("Poisoned lock for active_windows")
-                    .iter()
-                    .next()
-                    .cloned()
-            })
     }
 }
 

rust/numaflow-core/src/reduce/reducer/aligned/windower/sliding.rs

@@ -156,33 +156,41 @@ impl SlidingWindowManager {
     pub(crate) fn close_windows(&self, watermark: DateTime<Utc>) -> Vec<AlignedWindowMessage> {
         let mut result = Vec::new();
 
-        // Find windows that need to be closed
-        let mut active_windows = self
-            .active_windows
-            .write()
-            .expect("Poisoned lock in active_windows");
-        let mut closed_windows = self
-            .closed_windows
-            .write()
-            .expect("Poisoned lock in active_windows");
+        let windows_to_close = {
+            let mut active_windows = self
+                .active_windows
+                .write()
+                .expect("Poisoned lock for active_windows");
 
-        // get all the windows that have end time less than the watermark
-        let windows_to_close: Vec<_> = active_windows
-            .iter()
-            .filter(|window| window.end_time <= watermark) // window end time is exclusive, hence <=
-            .cloned()
-            .collect();
+            let mut windows_to_close = Vec::new();
+            active_windows.retain(|window| {
+                if window.end_time <= watermark {
+                    windows_to_close.push(window.clone());
+                    false
+                } else {
+                    true
+                }
+            });
+
+            windows_to_close
+        };
+
+        // add the windows to closed_windows
+        {
+            let mut closed_windows = self
+                .closed_windows
+                .write()
+                .expect("Poisoned lock for closed_windows");
+            for window in &windows_to_close {
+                closed_windows.insert(window.clone());
+            }
+        }
 
-        // Move windows from active to closed
         for window in windows_to_close {
             result.push(AlignedWindowMessage {
-                operation: AlignedWindowOperation::Close {
-                    window: window.clone(),
-                },
                 pnf_slot: window_pnf_slot(&window),
+                operation: AlignedWindowOperation::Close { window },
             });
-            active_windows.remove(&window);
-            closed_windows.insert(window);
         }
 
         result
@@ -206,20 +214,22 @@ impl SlidingWindowManager {
         // get the oldest window from closed_windows, if closed_windows is empty, get the oldest
         // from active_windows
         // NOTE: closed windows will always have a lower end time than active_windows
-        self.closed_windows
+        {
+            let closed_windows = self
+                .closed_windows
+                .read()
+                .expect("Poisoned lock for closed_windows");
+            if let Some(window) = closed_windows.iter().next() {
+                return Some(window.clone());
+            }
+        }
+
+        self.active_windows
             .read()
-            .expect("Poisoned lock for closed_windows")
+            .expect("Poisoned lock for active_windows")
             .iter()
             .next()
             .cloned()
-            .or_else(|| {
-                self.active_windows
-                    .read()
-                    .expect("Poisoned lock for active_windows")
-                    .iter()
-                    .next()
-                    .cloned()
-            })
     }
 
     /// Helper method to format sorted window information for logging.
@@ -722,7 +732,7 @@ mod tests {
     #[test]
     fn test_assign_windows_with_small_slide() {
         // prepopulate active windows
-        let active_windows = vec![
+        let active_windows = [
             Window::new(
                 Utc.timestamp_millis_opt(90000).unwrap(),
                 Utc.timestamp_millis_opt(150000).unwrap(),
@@ -771,7 +781,7 @@ mod tests {
         // Window 4: [70000, 130000) - event falls in this window
         // Window 5: [60000, 120000) - event falls in this window
         // Window 6: [50000, 110000) - event falls in this window
-        let expected_windows = vec![
+        let expected_windows = [
             Window::new(
                 Utc.timestamp_millis_opt(100000).unwrap(),
                 Utc.timestamp_millis_opt(160000).unwrap(),

rust/numaflow-core/src/reduce/reducer/unaligned/windower/session.rs

@@ -149,17 +149,16 @@ impl SessionWindowManager {
     /// window.end_time might grow overtime and overlap. At the time of close, we merge these windows
     /// those have overlapping end times.
     pub(crate) fn close_windows(&self, watermark: DateTime<Utc>) -> Vec<UnalignedWindowMessage> {
-        let mut active_windows = self.active_windows.write().expect("Poisoned lock");
-
         // Extract and remove expired windows from active windows
-        let closed_windows_by_key = Self::extract_expired_windows(&mut active_windows, watermark);
+        let closed_windows_by_key = {
+            let mut active_windows = self.active_windows.write().expect("Poisoned lock");
+            Self::extract_expired_windows(&mut active_windows, watermark)
+        };
 
         // Process each key's closed windows
         closed_windows_by_key
             .into_iter()
-            .flat_map(|(key, windows)| {
-                self.process_closed_windows(&mut active_windows, &key, windows)
-            })
+            .flat_map(|(key, windows)| self.process_closed_windows(&key, windows))
             .collect()
     }
 
@@ -198,20 +197,18 @@ impl SessionWindowManager {
     /// and then close the windows that cannot be merged.
     fn process_closed_windows(
         &self,
-        active_windows: &mut HashMap<String, BTreeSet<Window>>,
         key: &str,
         closed_windows: Vec<Window>,
     ) -> Vec<UnalignedWindowMessage> {
         Self::windows_that_can_be_merged(&closed_windows)
             .into_iter()
-            .filter_map(|group| self.process_closing_window_group(active_windows, key, group))
+            .filter_map(|group| self.process_closing_window_group(key, group))
             .collect()
     }
 
     /// Process a group of windows that can be merged
     fn process_closing_window_group(
         &self,
-        active_windows: &mut HashMap<String, BTreeSet<Window>>,
         key: &str,
         closing_group: Vec<Window>,
     ) -> Option<UnalignedWindowMessage> {
@@ -223,7 +220,12 @@ impl SessionWindowManager {
         let window_to_close = Self::merge_windows(&closing_group);
 
         // Try to merge with active windows
-        match Self::try_merge_with_active(active_windows, key, &window_to_close) {
+        let merge_result = {
+            let mut active_windows = self.active_windows.write().expect("Poisoned lock");
+            Self::try_merge_with_active(&mut active_windows, key, &window_to_close)
+        };
+
+        match merge_result {
             Some((old_active, new_merged)) => Some(UnalignedWindowMessage {
                 operation: UnalignedWindowOperation::Merge {
                     windows: vec![window_to_close, old_active, new_merged],
@@ -232,10 +234,10 @@ impl SessionWindowManager {
             }),
             None => {
                 // Move window to closed_windows
-                self.closed_windows
-                    .write()
-                    .expect("Poisoned lock")
-                    .insert(window_to_close.clone());
+                {
+                    let mut closed_windows = self.closed_windows.write().expect("Poisoned lock");
+                    closed_windows.insert(window_to_close.clone());
+                }
 
                 Some(UnalignedWindowMessage {
                     operation: UnalignedWindowOperation::Close {
@@ -353,20 +355,20 @@ impl SessionWindowManager {
         // Get the oldest window from closed_windows first, if closed_windows is empty, get the oldest
         // from active_windows
         // NOTE: closed windows will always have a lower end time than active_windows
-        self.closed_windows
-            .read()
-            .expect("Poisoned lock")
-            .iter()
-            .next()
+
+        {
+            let closed_windows = self.closed_windows.read().expect("Poisoned lock");
+            if let Some(window) = closed_windows.iter().next() {
+                return Some(window.end_time);
+            }
+        }
+
+        let active_windows = self.active_windows.read().expect("Poisoned lock");
+        active_windows
+            .values()
+            .filter_map(|windows| windows.iter().next())
             .map(|window| window.end_time)
-            .or_else(|| {
-                self.active_windows
-                    .read()
-                    .expect("Poisoned lock")
-                    .values()
-                    .flat_map(|window_set| window_set.iter().map(|window| window.end_time))
-                    .min()
-            })
+            .min()
     }
 }