diskann_disk/build/builder/
tokio.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use diskann::{ANNError, ANNResult};
7
8/// Creates a new multi-threaded tokio runtime with the specified number of worker threads.
9/// If `num_threads` is 0, it defaults to the number of logical CPUs.
10pub fn create_runtime(num_threads: usize) -> ANNResult<tokio::runtime::Runtime> {
11    let mut builder = tokio::runtime::Builder::new_multi_thread();
12
13    if num_threads != 0 {
14        builder.worker_threads(num_threads);
15    }
16
17    builder.build().map_err(|err| {
18        ANNError::log_index_error(format!("Failed to initialize tokio runtime: {}", err))
19    })
20}
21
22#[cfg(test)]
23mod tests {
24    use super::*;
25
26    fn get_logical_cpu_count() -> usize {
27        std::thread::available_parallelism()
28            .map(|n| n.get())
29            .unwrap_or(1)
30    }
31
32    #[test]
33    fn test_create_runtime_with_zero_threads_no_panic() {
34        // This test ensures that passing 0 threads doesn't panic
35        // and properly defaults to the number of logical CPUs
36        let result = create_runtime(0);
37
38        // Should not panic and should succeed
39        assert!(result.is_ok(), "create_runtime(0) should not panic or fail");
40
41        let runtime = result.unwrap();
42
43        // Verify the runtime was created successfully by executing a simple task
44        let result = runtime.block_on(async { tokio::spawn(async { 42 }).await });
45
46        assert!(result.is_ok(), "Runtime should be functional");
47        assert_eq!(result.unwrap(), 42);
48    }
49
50    #[test]
51    fn test_create_runtime_with_specific_threads() {
52        // Test that specifying a specific number of threads works
53        let result = create_runtime(2);
54        assert!(result.is_ok(), "create_runtime(2) should succeed");
55
56        let runtime = result.unwrap();
57
58        // Verify the runtime works
59        let result = runtime.block_on(async { tokio::spawn(async { "test" }).await });
60
61        assert!(result.is_ok(), "Runtime should be functional");
62        assert_eq!(result.unwrap(), "test");
63    }
64
65    #[test]
66    fn test_create_runtime_with_one_thread() {
67        // Test edge case with 1 thread
68        let result = create_runtime(1);
69        assert!(result.is_ok(), "create_runtime(1) should succeed");
70
71        let runtime = result.unwrap();
72
73        // Verify the runtime works even with just 1 thread
74        let result = runtime.block_on(async { tokio::spawn(async { true }).await });
75
76        assert!(
77            result.is_ok(),
78            "Single-threaded runtime should be functional"
79        );
80        assert!(result.unwrap());
81    }
82
83    #[test]
84    fn test_zero_threads_defaults_to_cpu_count() {
85        // Test that 0 threads actually uses the logical CPU count
86        let expected_cpu_count = get_logical_cpu_count();
87
88        // We can't directly inspect the runtime's thread count easily,
89        // but we can ensure it doesn't panic and works correctly
90        let result = create_runtime(0);
91        assert!(
92            result.is_ok(),
93            "create_runtime(0) should default to {} CPUs",
94            expected_cpu_count
95        );
96
97        let runtime = result.unwrap();
98
99        // Test that the runtime can handle multiple concurrent tasks
100        // which would fail if it only had 1 thread and we expected more
101        let result = runtime.block_on(async {
102            let tasks = (0..expected_cpu_count.min(4))
103                .map(|i| tokio::spawn(async move { i * 2 }))
104                .collect::<Vec<_>>();
105
106            let mut results = Vec::new();
107            for task in tasks {
108                results.push(task.await.unwrap());
109            }
110            results
111        });
112
113        assert!(
114            result.len() <= 4,
115            "Should handle concurrent tasks successfully"
116        );
117    }
118}
diskann_disk/build/builder/tokio.rs

diskann_disk/build/builder/
tokio.rs
