diskann_disk/build/chunking/continuation/
utils.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use std::{error::Error, thread::sleep};
7
8use tracing::info;
9
10use super::continuation_tracker::{ContinuationGrant, ContinuationTrackerTrait};
11use crate::build::chunking::checkpoint::Progress;
12
13/// This takes an operation with an iterator of oprands,
14/// and processes the oprands using the operation in a loop,
15/// until the continuation_checker asks it to stop.
16/// The continuation_checker is used to get continuation grants between processing each operation.
17/// The clean_up function is called after the loop is broken and before exit.
18/// The function returns a Progress enum, which indicates the number of operations executed.
19pub fn process_while_resource_is_available<Action, ParamIter, Param, E>(
20    mut action: Action,
21    params: ParamIter,
22    continuation_checker: Box<dyn ContinuationTrackerTrait>,
23) -> Result<Progress, E>
24where
25    ParamIter: Iterator<Item = Param>,
26    Action: FnMut(Param) -> Result<(), E>,
27    E: Error,
28{
29    for (idx, param) in params.enumerate() {
30        loop {
31            match continuation_checker.get_continuation_grant() {
32                ContinuationGrant::Continue => {
33                    info!("Continue processing.");
34                    action(param)?;
35                    break;
36                }
37                ContinuationGrant::Yield(duration) => {
38                    info!(
39                        "Continuation checker asks to yield for {} ms.",
40                        duration.as_millis()
41                    );
42                    sleep(duration);
43                }
44                ContinuationGrant::Stop => {
45                    info!("Continuation checker asks to stop. Breaking the loop.");
46                    return Ok(Progress::Processed(idx));
47                }
48            }
49        }
50    }
51
52    Ok(Progress::Completed)
53}
54
55/// Asynchronous version of [`process_while_resource_is_available`].
56///
57/// Takes an async operation with an iterator of operands and processes them in a loop
58/// until the continuation_checker signals to stop.
59pub async fn process_while_resource_is_available_async<Action, ParamIter, Param, Fut, E>(
60    mut action: Action,
61    params: ParamIter,
62    continuation_checker: Box<dyn ContinuationTrackerTrait>,
63) -> Result<Progress, E>
64where
65    ParamIter: Iterator<Item = Param>,
66    Action: FnMut(Param) -> Fut,
67    Fut: core::future::Future<Output = Result<(), E>>,
68    E: Error,
69{
70    for (idx, param) in params.enumerate() {
71        loop {
72            match continuation_checker.get_continuation_grant() {
73                ContinuationGrant::Continue => {
74                    info!("Continue processing.");
75                    action(param).await?;
76                    break;
77                }
78                ContinuationGrant::Yield(duration) => {
79                    info!(
80                        "Continuation checker asks to yield for {} ms.",
81                        duration.as_millis()
82                    );
83                    sleep(duration);
84                }
85                ContinuationGrant::Stop => {
86                    info!("Continuation checker asks to stop. Breaking the loop.");
87                    return Ok(Progress::Processed(idx));
88                }
89            }
90        }
91    }
92
93    Ok(Progress::Completed)
94}
95
96#[cfg(test)]
97mod tests {
98    use super::super::continuation_tracker::NaiveContinuationTracker;
99    use super::*;
100    use std::fmt;
101
102    #[derive(Debug)]
103    struct TestError;
104
105    impl fmt::Display for TestError {
106        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
107            write!(f, "TestError")
108        }
109    }
110
111    impl Error for TestError {}
112
113    #[test]
114    fn test_process_while_resource_is_available_completes() {
115        let checker = Box::new(NaiveContinuationTracker::default());
116        let items = vec![1, 2, 3, 4, 5];
117        let mut processed = Vec::new();
118
119        let result = process_while_resource_is_available(
120            |item| {
121                processed.push(item);
122                Ok::<(), TestError>(())
123            },
124            items.into_iter(),
125            checker,
126        );
127
128        assert!(result.is_ok());
129        match result.unwrap() {
130            Progress::Completed => assert_eq!(processed, vec![1, 2, 3, 4, 5]),
131            _ => panic!("Expected Completed"),
132        }
133    }
134
135    #[test]
136    fn test_process_while_resource_is_available_empty_iter() {
137        let checker = Box::new(NaiveContinuationTracker::default());
138        let items: Vec<i32> = vec![];
139
140        let result = process_while_resource_is_available(
141            |_item| Ok::<(), TestError>(()),
142            items.into_iter(),
143            checker,
144        );
145
146        assert!(result.is_ok());
147        match result.unwrap() {
148            Progress::Completed => {}
149            _ => panic!("Expected Completed"),
150        }
151    }
152
153    #[tokio::test]
154    async fn test_process_while_resource_is_available_async_completes() {
155        let checker = Box::new(NaiveContinuationTracker::default());
156        let items = vec![1, 2, 3];
157        let processed = std::sync::Arc::new(tokio::sync::Mutex::new(Vec::new()));
158
159        let result = process_while_resource_is_available_async(
160            |item| {
161                let processed = processed.clone();
162                async move {
163                    processed.lock().await.push(item);
164                    Ok::<(), TestError>(())
165                }
166            },
167            items.into_iter(),
168            checker,
169        )
170        .await;
171
172        assert!(result.is_ok());
173        match result.unwrap() {
174            Progress::Completed => {
175                let processed = processed.lock().await;
176                assert_eq!(*processed, vec![1, 2, 3]);
177            }
178            _ => panic!("Expected Completed"),
179        }
180    }
181}
diskann_disk/build/chunking/continuation/utils.rs

diskann_disk/build/chunking/continuation/
utils.rs
