Skip to main content

diskann_disk/build/chunking/continuation/
continuation_tracker.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use std::time::Duration;
7
8// ContinuationGrant struct used for ochestrating chunkable operations.
9pub enum ContinuationGrant {
10    Continue,
11    Yield(Duration),
12    Stop,
13}
14
15// Contituation grant trait to be implemented by the client.
16// This trait is used to get a continuation grant during chunk intervals while doing chunkable operations.
17pub trait ContinuationTrackerTrait: Send + Sync + ContinuationCheckerTraitClone {
18    fn get_continuation_grant(&self) -> ContinuationGrant;
19}
20
21// Naive implementation of the ContinuationGrantProviderTrait
22// This implementation always returns ContinuationGrant::Continue.
23#[derive(Default, Clone)]
24pub struct NaiveContinuationTracker {}
25
26impl ContinuationTrackerTrait for NaiveContinuationTracker {
27    fn get_continuation_grant(&self) -> ContinuationGrant {
28        ContinuationGrant::Continue
29    }
30}
31
32pub trait ContinuationCheckerTraitClone {
33    fn clone_box(&self) -> Box<dyn ContinuationTrackerTrait>;
34}
35
36impl<T> ContinuationCheckerTraitClone for T
37where
38    T: 'static + ContinuationTrackerTrait + Clone,
39{
40    fn clone_box(&self) -> Box<dyn ContinuationTrackerTrait> {
41        Box::new(self.clone())
42    }
43}
44
45#[cfg(test)]
46mod tests {
47    use super::*;
48
49    #[test]
50    fn test_naive_continuation_tracker_default() {
51        let tracker = NaiveContinuationTracker::default();
52        // Verify it always returns Continue
53        match tracker.get_continuation_grant() {
54            ContinuationGrant::Continue => {}
55            _ => panic!("Expected Continue"),
56        }
57    }
58
59    #[test]
60    fn test_naive_continuation_tracker_clone_box() {
61        let tracker = NaiveContinuationTracker::default();
62        let boxed = tracker.clone_box();
63
64        // The boxed version should also return Continue
65        match boxed.get_continuation_grant() {
66            ContinuationGrant::Continue => {}
67            _ => panic!("Expected Continue"),
68        }
69    }
70}