111 lines
3.1 KiB
Rust
111 lines
3.1 KiB
Rust
|
//! Internal batch mode tracking using tokio task-local storage.
|
||
|
|
//!
|
||
|
|
//! This module provides a way to track whether we're in batch processing mode
|
||
|
|
//! without exposing it in the public API. Extractors check this flag to decide
|
||
|
|
//! whether to use `spawn_blocking` for CPU-intensive work.
|
||
|
|
|
||
|
|
use std::cell::Cell;
|
||
|
|
use tokio::task_local;
|
||
|
|
|
||
|
|
task_local! {
|
||
|
|
/// Task-local flag indicating batch processing mode.
|
||
|
|
///
|
||
|
|
/// When true, extractors use `spawn_blocking` for CPU-intensive work to enable
|
||
|
|
/// parallelism. When false (single-file mode), extractors run directly to avoid
|
||
|
|
/// spawn overhead.
|
||
|
|
static BATCH_MODE: Cell<bool>;
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Check if we're currently in batch processing mode.
|
||
|
|
///
|
||
|
|
/// Returns `false` if the task-local is not set (single-file mode).
|
||
|
|
#[cfg(any(
|
||
|
|
feature = "pdf",
|
||
|
|
feature = "office",
|
||
|
|
feature = "excel",
|
||
|
|
feature = "excel-wasm",
|
||
|
|
feature = "archives"
|
||
|
|
))]
|
||
|
|
pub(crate) fn is_batch_mode() -> bool {
|
||
|
|
BATCH_MODE.try_with(|cell| cell.get()).unwrap_or(false)
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Run a future with batch mode enabled.
|
||
|
|
///
|
||
|
|
/// This sets the task-local BATCH_MODE flag for the duration of the future.
|
||
|
|
pub(crate) async fn with_batch_mode<F, T>(future: F) -> T
|
||
|
|
where
|
||
|
|
F: std::future::Future<Output = T>,
|
||
|
|
{
|
||
|
|
BATCH_MODE.scope(Cell::new(true), future).await
|
||
|
|
}
|
||
|
|
|
||
|
|
#[cfg(all(
|
||
|
|
test,
|
||
|
|
any(
|
||
|
|
feature = "pdf",
|
||
|
|
feature = "office",
|
||
|
|
feature = "excel",
|
||
|
|
feature = "excel-wasm",
|
||
|
|
feature = "archives"
|
||
|
|
)
|
||
|
|
))]
|
||
|
|
mod tests {
|
||
|
|
use super::*;
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn test_batch_mode_not_set_by_default() {
|
||
|
|
let result = is_batch_mode();
|
||
|
|
assert!(!result, "batch mode should be false by default");
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn test_with_batch_mode_sets_flag() {
|
||
|
|
let result = with_batch_mode(async { is_batch_mode() }).await;
|
||
|
|
|
||
|
|
assert!(result, "batch mode should be true inside with_batch_mode");
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn test_batch_mode_scoped_to_future() {
|
||
|
|
assert!(!is_batch_mode(), "batch mode should be false before");
|
||
|
|
|
||
|
|
with_batch_mode(async {
|
||
|
|
assert!(is_batch_mode(), "batch mode should be true inside");
|
||
|
|
})
|
||
|
|
.await;
|
||
|
|
|
||
|
|
assert!(!is_batch_mode(), "batch mode should be false after future completes");
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn test_nested_batch_mode_calls() {
|
||
|
|
let result = with_batch_mode(async {
|
||
|
|
let outer = is_batch_mode();
|
||
|
|
let inner = with_batch_mode(async { is_batch_mode() }).await;
|
||
|
|
(outer, inner)
|
||
|
|
})
|
||
|
|
.await;
|
||
|
|
|
||
|
|
assert!(result.0, "outer batch mode should be true");
|
||
|
|
assert!(result.1, "inner batch mode should be true");
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn test_batch_mode_unaffected_after_with_batch_mode() {
|
||
|
|
with_batch_mode(async {
|
||
|
|
assert!(is_batch_mode(), "first call should set batch mode");
|
||
|
|
})
|
||
|
|
.await;
|
||
|
|
|
||
|
|
assert!(!is_batch_mode(), "batch mode should be false between calls");
|
||
|
|
|
||
|
|
with_batch_mode(async {
|
||
|
|
assert!(is_batch_mode(), "second call should set batch mode");
|
||
|
|
})
|
||
|
|
.await;
|
||
|
|
|
||
|
|
assert!(!is_batch_mode(), "batch mode should be false after all calls");
|
||
|
|
}
|
||
|
|
}
|