Kept the runtime API responsive during blocking work (#1043)

This commit is contained in:
Thorsten Sommer authored and GitHub committed 2026-10-10 21:39:13 +02:00
1 parent 4616d1a4b2
commit fb756dd767
11 files changed
+389 -77

No files matched your search

+19
View File
@@ -130,6 +130,25 @@ Key modules:
- `pandoc.rs` - Integration with Pandoc for document conversion
- `log.rs` - Logging infrastructure using `flexi_logger`
**Runtime API handlers never block.** All calls of the .NET app share one HTTP/2 connection, and the
task driving it may wait on exactly the Tokio worker which a blocking handler occupies, so a single
blocking handler can hold up the whole app. Therefore:
- File and disk access, the OS keyring, waiting for a `std::sync::Mutex`, and CPU-bound work such as
loading a tokenizer or scanning text run in `tokio::task::spawn_blocking`. Follow
`run_qdrant_edge_request` in `qdrant_edge_database.rs`, `prepare_image` and `prepare_image_sync` in
`image.rs`, or `sanitize_batch` in `prompt_injection/api.rs`.
- When an API offers a callback, as the file dialogs do, await the callback instead of calling the
blocking variant; see `await_dialog` in `file_actions.rs`.
- Short, bounded work, such as a single metadata lookup or writing a log line, may stay on the worker.
- When the blocking task fails, answer with an error, never with an empty value the app could take
for a valid answer.
- Never keep a `std::sync::MutexGuard` alive across an `.await`, not even with an explicit `drop`
before it: the future of the handler is then no longer `Send`, and Axum refuses it. Scope the guard
in a block instead.
- `runtime_api::test_support::assert_runtime_stays_free` tests a handler which waits for a lock. Run
such a test with `#[tokio::test(flavor = "multi_thread", worker_threads = 1)]`.
### .NET App (`app/MindWork AI Studio/`)
**Entry point:** `app/MindWork AI Studio/Program.cs`
@@ -26,6 +26,7 @@
- Fixed AI Studio sometimes rebuilding the index of a data source from scratch when the configuration of your organization was applied while the data source was being updated.
- Fixed justified texts being hyphenated by the rules of English even when AI Studio shows another language, such as German. Screen readers now also know which language AI Studio uses.
- Fixed the log file of AI Studio showing the wrong time for some entries when AI Studio was busy. Each entry now shows when it actually happened.
- Fixed AI Studio pausing for a moment while it opened the local index of a data source, such as a mailbox. Other work, for example applying the configuration of your organization, had to wait until then.
- Updated the code contributions on the supporters page, which now thank everyone who has contributed code to AI Studio so far.
- Upgraded several libraries to improve security.
- Upgraded to Rust v1.99.0
+20
View File
@@ -77,6 +77,26 @@ fn release_clipboard<B>(clipboard: &mut Option<B>) -> bool {
/// Sets the clipboard text to the provided encrypted text.
pub async fn set_clipboard(_token: APIToken, encrypted_text: String) -> Json<SetClipboardResponse> {
//
// The clipboard backend of the operating system may take a moment to answer, a failed write is
// retried with a fresh backend, and the lock waits for a write still running. None of this
// belongs on a runtime worker, because all calls of the app share one HTTP/2 connection.
//
match tokio::task::spawn_blocking(move || set_clipboard_sync(encrypted_text)).await {
Ok(response) => response,
Err(error) => {
let issue = format!("The clipboard worker failed: {error}");
error!(Source = "Clipboard"; "{issue}");
Json(SetClipboardResponse {
success: false,
issue,
})
},
}
}
/// Performs the blocking part of [`set_clipboard`].
fn set_clipboard_sync(encrypted_text: String) -> Json<SetClipboardResponse> {
let encrypted_text = EncryptedText::new(encrypted_text);
// Decrypt this text first:
+25 -24
View File
@@ -1,4 +1,5 @@
use crate::api_token::APIToken;
use axum::http::StatusCode;
use axum::Json;
use log::{debug, error, info, warn};
use serde::Serialize;
@@ -641,35 +642,35 @@ struct EnterpriseSourceData {
encryption_secret: String,
}
pub async fn read_enterprise_env_config_id(_token: APIToken) -> String {
debug!("Trying to read the effective enterprise configuration ID.");
resolve_effective_enterprise_config_source()
.configs
.into_iter()
.next()
.map(|config| config.id)
.unwrap_or_default()
}
pub async fn read_enterprise_env_config_server_url(_token: APIToken) -> String {
debug!("Trying to read the effective enterprise configuration server URL.");
resolve_effective_enterprise_config_source()
.configs
.into_iter()
.next()
.map(|config| config.server_url)
.unwrap_or_default()
}
pub async fn read_enterprise_env_config_encryption_secret(_token: APIToken) -> String {
pub async fn read_enterprise_env_config_encryption_secret(_token: APIToken) -> Result<String, (StatusCode, String)> {
debug!("Trying to read the effective enterprise configuration encryption secret.");
resolve_effective_enterprise_secret_source().encryption_secret
read_enterprise_sources(|| resolve_effective_enterprise_secret_source().encryption_secret).await
}
/// Returns all enterprise configurations from the effective source.
pub async fn read_enterprise_configs(_token: APIToken) -> Json<Vec<EnterpriseConfig>> {
pub async fn read_enterprise_configs(_token: APIToken) -> Result<Json<Vec<EnterpriseConfig>>, (StatusCode, String)> {
info!("Trying to read the effective enterprise configurations.");
Json(resolve_effective_enterprise_config_source().configs)
read_enterprise_sources(|| Json(resolve_effective_enterprise_config_source().configs)).await
}
/// Reads the enterprise sources on the blocking pool.
///
/// Every call reads the policy files and, on Windows, the registry again. That is quick most of
/// the time, but a slow disk must not hold up every other call of the app, because they all share
/// one HTTP/2 connection.
///
/// A failure is answered as an error, never as an empty value: the app would read an empty value
/// as "no configuration of your organization" and act on it.
async fn read_enterprise_sources<T, F>(read: F) -> Result<T, (StatusCode, String)>
where
T: Send + 'static,
F: FnOnce() -> T + Send + 'static,
{
tokio::task::spawn_blocking(read).await.map_err(|error| {
let issue = format!("The enterprise configuration worker failed: {error}");
error!("{issue}");
(StatusCode::INTERNAL_SERVER_ERROR, issue)
})
}
pub fn resolve_external_http_custom_root_certificate_policy() -> ExternalHttpCustomRootCertificatePolicy {
+51 -17
View File
@@ -123,26 +123,25 @@ pub async fn select_directory(
Query(query): Query<SelectDirectoryQuery>,
previous_directory: Option<Json<PreviousDirectory>>,
) -> Json<DirectorySelectionResponse> {
let main_window_lock = MAIN_WINDOW.lock().unwrap();
let main_window = match main_window_lock.as_ref() {
Some(window) => window,
None => {
error!(Source = "Tauri"; "Cannot open directory dialog: main window not available.");
return Json(DirectorySelectionResponse {
user_cancelled: true,
selected_directory: String::from(""),
});
}
let dialog = MAIN_WINDOW
.lock()
.unwrap()
.as_ref()
.map(|w| w.dialog().file().set_parent(w).set_title(&query.title));
let Some(mut dialog) = dialog else {
error!(Source = "Tauri"; "Cannot open directory dialog: main window not available.");
return Json(DirectorySelectionResponse {
user_cancelled: true,
selected_directory: String::from(""),
});
};
let mut dialog = main_window.dialog().file().set_parent(main_window).set_title(&query.title);
if let Some(previous) = previous_directory {
dialog = dialog.set_directory(previous.path.clone());
}
drop(main_window_lock);
let folder_path = dialog.blocking_pick_folder();
let folder_path = await_dialog(|answer| dialog.pick_folder(answer)).await;
match folder_path {
Some(path) => {
match path.into_path() {
@@ -203,7 +202,7 @@ pub async fn select_file(
}
// Show the file dialog and get the selected file path:
let file_path = file_dialog.blocking_pick_file();
let file_path = await_dialog(|answer| file_dialog.pick_file(answer)).await;
match file_path {
Some(path) => match path.into_path() {
Ok(pb) => {
@@ -262,7 +261,7 @@ pub async fn select_files(
}
// Show the file dialog and get the selected file path:
let file_paths = file_dialog.blocking_pick_files();
let file_paths = await_dialog(|answer| file_dialog.pick_files(answer)).await;
match file_paths {
Some(paths) => {
let converted: Vec<String> = paths.into_iter().filter_map(|p| p.into_path().ok()).map(|pb| pb.to_string_lossy().to_string()).collect();
@@ -314,7 +313,7 @@ pub async fn save_file(_token: APIToken, payload: Json<SaveFileOptions>) -> Json
}
// Displays the file dialogue box and select the file:
let file_path = file_dialog.blocking_save_file();
let file_path = await_dialog(|answer| file_dialog.save_file(answer)).await;
match file_path {
Some(path) => match path.into_path() {
Ok(pb) => {
@@ -882,6 +881,26 @@ fn apply_filter<R: tauri::Runtime>(file_dialog: FileDialogBuilder<R>, filter: &O
}
}
/// Shows a file dialog and waits for the user's answer without holding up the runtime.
///
/// A dialog stays open until the user answers it, which can take minutes. The blocking variants of
/// the dialog API would hold a runtime worker that long, and with it every other call of the app,
/// because they all share one HTTP/2 connection. The plugin shows the dialog on the main thread and
/// hands the answer to a callback, so this function awaits that callback instead.
async fn await_dialog<T: Send + 'static>(show: impl FnOnce(Box<dyn FnOnce(Option<T>) + Send>)) -> Option<T> {
let (sender, receiver) = tokio::sync::oneshot::channel();
show(Box::new(move |answer| {
// The receiver is only gone when the request was dropped, and then nobody waits for the answer:
let _ = sender.send(answer);
}));
receiver.await.unwrap_or_else(|_| {
// The plugin drops the callback without calling it when it cannot reach the main thread:
error!(Source = "Tauri"; "The file dialog ended without an answer.");
None
})
}
fn split_save_file_path(file_path: &str) -> (Option<PathBuf>, Option<String>) {
let path = Path::new(file_path);
let directory = path
@@ -1047,6 +1066,21 @@ mod tests {
use super::*;
use std::fs;
#[tokio::test]
async fn a_dialog_answer_arrives_through_its_callback() {
let answer = await_dialog(|answer| {
std::thread::spawn(move || answer(Some("Quarterly briefing.html")));
}).await;
assert_eq!(answer, Some("Quarterly briefing.html"));
}
#[tokio::test]
async fn a_dialog_without_an_answer_counts_as_cancelled() {
let answer = await_dialog::<String>(drop).await;
assert_eq!(answer, None);
}
#[test]
fn save_file_options_accept_the_previous_file_contract() {
let options: SaveFileOptions = serde_json::from_str(
+1 -1
View File
@@ -524,7 +524,7 @@ pub async fn extract_data(
// whole file, and a failure here means we cannot answer the request at all.
//
let tokenizer = if query.include_token_count {
match crate::tokenizer::get_tokenizer(&query.tokenizer_path) {
match crate::tokenizer::get_tokenizer_off_worker(query.tokenizer_path.clone()).await {
Ok(tokenizer) => Some(tokenizer),
Err(e) => {
let error = ExtractionError::new(ExtractionErrorCode::InvalidRequest, format!("The tokenizer could not be loaded: {e}"));
+23 -6
View File
@@ -45,13 +45,30 @@ pub struct SanitizeBatchResponse {
pub results: Vec<SanitizeResponse>,
}
pub async fn sanitize(_token: APIToken, Json(request): Json<SanitizeRequest>) -> Json<SanitizeResponse> {
let (sanitized_text, report) = sanitize_text(&request.text);
pub async fn sanitize(
_token: APIToken,
Json(request): Json<SanitizeRequest>,
) -> Result<Json<SanitizeResponse>, (StatusCode, String)> {
//
// A single web page or retrieval context can still be large, and scanning it is CPU-bound. Like
// a batch, it goes to the blocking pool, so it does not stall every other call the app makes
// meanwhile.
//
tokio::task::spawn_blocking(move || {
let (sanitized_text, report) = sanitize_text(&request.text);
Json(SanitizeResponse {
sanitized_text,
findings: report.findings,
redacted_count: report.redacted_count,
Json(SanitizeResponse {
sanitized_text,
findings: report.findings,
redacted_count: report.redacted_count,
})
})
.await
.map_err(|error| {
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("The prompt injection filter failed: {error}"),
)
})
}
+86 -20
View File
@@ -473,15 +473,25 @@ fn qdrant_edge_base_path() -> QdrantEdgeResult<PathBuf> {
}
pub async fn qdrant_edge_info(_token: APIToken) -> Json<QdrantEdgeServiceInfo> {
let status = QDRANT_EDGE_STATUS.lock().unwrap();
let current_status = status.status;
let unavailable_reason = status.unavailable_reason.clone();
drop(status);
let (current_status, mut unavailable_reason) = {
let status = QDRANT_EDGE_STATUS.lock().unwrap();
(status.status, status.unavailable_reason.clone())
};
let database_guard = QDRANT_EDGE_DATABASE.lock().unwrap();
let database_info = database_guard
.as_ref()
.and_then(|database| database.info().ok());
//
// The info needs the lock of the database, which another request may hold for seconds while
// it loads a store, and it reads the stores directory. Both belong on the blocking pool, see
// run_qdrant_edge_request.
//
let database_info = match tokio::task::spawn_blocking(read_qdrant_edge_info).await {
Ok(database_info) => database_info,
Err(error) => {
let reason = format!("The Qdrant Edge worker failed: {error}");
error!(Source = "Qdrant Edge"; "Reading the Qdrant Edge info failed: {reason}");
unavailable_reason = Some(reason);
None
},
};
let is_available = current_status == QdrantEdgeStatus::Available && database_info.is_some();
Json(QdrantEdgeServiceInfo {
@@ -495,40 +505,48 @@ pub async fn qdrant_edge_info(_token: APIToken) -> Json<QdrantEdgeServiceInfo> {
})
}
fn read_qdrant_edge_info() -> Option<QdrantEdgeInfo> {
QDRANT_EDGE_DATABASE
.lock()
.unwrap()
.as_ref()
.and_then(|database| database.info().ok())
}
pub async fn ensure_qdrant_edge_store(_token: APIToken, Json(request): Json<EnsureQdrantEdgeStoreRequest>) -> Json<QdrantEdgeResponse<QdrantEdgeEnsureStoreResult>> {
execute_qdrant_edge_request(|database| {
run_qdrant_edge_request(move |database| {
database.ensure_store_exists(&request.store_name, &request.data_source_name, request.vector_size)
})
}).await
}
pub async fn insert_qdrant_edge_embedding(_token: APIToken, Json(request): Json<InsertQdrantEdgeEmbeddingRequest>) -> Json<QdrantEdgeResponse<()>> {
execute_qdrant_edge_request(|database| {
run_qdrant_edge_request(move |database| {
database.insert_embedding(&request.store_name, request.points)
})
}).await
}
pub async fn search_qdrant_edge_embeddings(_token: APIToken, Json(request): Json<SearchQdrantEdgeEmbeddingRequest>) -> Json<QdrantEdgeResponse<Vec<QdrantEdgeSearchResult>>> {
execute_qdrant_edge_request(|database| {
run_qdrant_edge_request(move |database| {
database.search_embedding(&request.store_name, request.vector, request.max_matches, request.point_ids)
})
}).await
}
pub async fn delete_qdrant_edge_embedding_by_file(_token: APIToken, Json(request): Json<DeleteQdrantEdgeEmbeddingByFileRequest>) -> Json<QdrantEdgeResponse<()>> {
execute_qdrant_edge_request(|database| {
run_qdrant_edge_request(move |database| {
database.delete_embedding_by_file(&request.store_name, &request.file_path)
})
}).await
}
pub async fn optimize_qdrant_edge_store(_token: APIToken, Json(request): Json<OptimizeQdrantEdgeStoreRequest>) -> Json<QdrantEdgeResponse<()>> {
execute_qdrant_edge_request(|database| {
run_qdrant_edge_request(move |database| {
database.optimize_store(&request.store_name)
})
}).await
}
pub async fn delete_qdrant_edge_store(_token: APIToken, Json(request): Json<DeleteQdrantEdgeStoreRequest>) -> Json<QdrantEdgeResponse<()>> {
execute_qdrant_edge_request(|database| {
run_qdrant_edge_request(move |database| {
database.delete_store(&request.store_name)
})
}).await
}
pub fn start_qdrant_edge_database<R: tauri::Runtime>(app_handle: tauri::AppHandle<R>) {
@@ -569,6 +587,32 @@ pub fn stop_qdrant_edge_database() {
set_qdrant_edge_unavailable("Qdrant Edge was stopped.".to_string());
}
/// Runs one request against the database on the blocking pool.
///
/// A request waits for the lock of the database and may load a store from disk, which together
/// can take seconds. On a runtime worker, that would hold up every other call of the app as well:
/// they all share one HTTP/2 connection, and the task driving it may be waiting in the LIFO slot of
/// exactly this worker, where no other worker can take it over.
async fn run_qdrant_edge_request<T, F>(operation: F) -> Json<QdrantEdgeResponse<T>>
where
T: Serialize + Send + 'static,
F: FnOnce(&mut QdrantEdgeDatabase) -> QdrantEdgeResult<T> + Send + 'static,
{
match tokio::task::spawn_blocking(move || execute_qdrant_edge_request(operation)).await {
Ok(response) => response,
Err(error) => {
let issue = format!("The Qdrant Edge worker failed: {error}");
error!(Source = "Qdrant Edge"; "Qdrant Edge request failed: {issue}");
Json(QdrantEdgeResponse {
success: false,
issue,
issue_code: "",
data: None,
})
},
}
}
fn execute_qdrant_edge_request<T, F>(operation: F) -> Json<QdrantEdgeResponse<T>>
where
T: Serialize,
@@ -1033,6 +1077,7 @@ fn validate_store_name(store_name: &str) -> QdrantEdgeResult<()> {
#[cfg(test)]
mod tests {
use super::*;
use crate::runtime_api::test_support::assert_runtime_stays_free;
use qdrant_edge::WalOptions;
#[test]
@@ -1254,6 +1299,27 @@ mod tests {
fs::remove_dir_all(test_directory).unwrap();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn a_request_waiting_for_the_database_leaves_the_runtime_free() {
let request = optimize_qdrant_edge_store(
APIToken::from_hex_text("test"),
Json(OptimizeQdrantEdgeStoreRequest { store_name: SEARCH_STORE.to_string() }),
);
let response = assert_runtime_stays_free(&QDRANT_EDGE_DATABASE, request).await;
// No database runs in the tests, so the request ends as soon as it gets the lock.
assert!(!response.success);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn the_info_waiting_for_the_database_leaves_the_runtime_free() {
let request = qdrant_edge_info(APIToken::from_hex_text("test"));
let info = assert_runtime_stays_free(&QDRANT_EDGE_DATABASE, request).await;
assert!(!info.is_available);
}
const SEARCH_STORE: &str = "rag_6cc665a82b1e4d42bc748015b7b391ec";
const POINT_X: &str = "0b5f1e8a-3c2d-4e6f-9a1b-7c8d9e0f1a2b";
const POINT_Y: &str = "1c6a2f9b-4d3e-4f70-8b2c-8d9e0f1a2b3c";
+48 -2
View File
@@ -64,8 +64,6 @@ pub fn start_runtime_api() {
.route("/system/runtime/info", get(crate::environment::get_runtime_info))
.route("/system/language", get(crate::environment::read_user_language))
.route("/system/username", get(crate::environment::read_user_name))
.route("/system/enterprise/config/id", get(crate::environment::read_enterprise_env_config_id))
.route("/system/enterprise/config/server", get(crate::environment::read_enterprise_env_config_server_url))
.route("/system/enterprise/config/encryption_secret", get(crate::environment::read_enterprise_env_config_encryption_secret))
.route("/system/enterprise/configs", get(crate::environment::read_enterprise_configs))
.route("/retrieval/fs/extract", get(crate::file_data::extract_data))
@@ -113,3 +111,51 @@ fn install_rustls_crypto_provider() {
let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
});
}
#[cfg(test)]
pub(crate) mod test_support {
use std::sync::Mutex;
use std::time::{Duration, Instant};
/// How long the lock stays held while the request waits for it.
const LOCK_DURATION: Duration = Duration::from_secs(1);
/// How long the probe may wait for the runtime. Far below the lock duration, so a request which
/// blocks the worker while it waits fails the test every time.
const PROBE_LIMIT: Duration = Duration::from_millis(200);
/// Starts a request while another thread holds the given lock, as if a long operation were
/// running, and checks that a probe still gets through the runtime meanwhile. Returns the
/// response once the lock is free again.
///
/// Call it from a test with a single worker, `#[tokio::test(flavor = "multi_thread",
/// worker_threads = 1)]`. A request which blocks that worker while it waits then holds up the
/// probe, just as in the app it holds up the one HTTP/2 connection which all calls share.
pub(crate) async fn assert_runtime_stays_free<T, F>(lock: &'static Mutex<T>, request: F) -> F::Output
where
T: Send + 'static,
F: Future + Send + 'static,
F::Output: Send + 'static,
{
let (locked_sender, locked_receiver) = std::sync::mpsc::channel();
let lock_holder = std::thread::spawn(move || {
let guard = lock.lock().unwrap();
locked_sender.send(()).unwrap();
std::thread::sleep(LOCK_DURATION);
drop(guard);
});
locked_receiver.recv().unwrap();
let request = tokio::spawn(request);
let probe_started = Instant::now();
tokio::spawn(async {}).await.unwrap();
let probe_duration = probe_started.elapsed();
lock_holder.join().unwrap();
let response = request.await.unwrap();
assert!(probe_duration < PROBE_LIMIT, "The probe waited {probe_duration:?} for the runtime.");
response
}
}
+55
View File
@@ -78,8 +78,37 @@ pub fn init_secret_store() {
}
}
/// Runs one request against the credential store on the blocking pool.
///
/// The credential store of the operating system may take a moment to answer, and on macOS and Linux
/// it may wait until the user unlocks it, which can take minutes. On a runtime worker, that would
/// hold up every other call of the app, because they all share one HTTP/2 connection.
async fn run_secret_store_request<T, F>(request: F) -> Result<T, String>
where
T: Send + 'static,
F: FnOnce() -> T + Send + 'static,
{
tokio::task::spawn_blocking(request).await.map_err(|error| {
let issue = format!("The secret store worker failed: {error}");
error!(Source = "Secret Store"; "{issue}");
issue
})
}
/// Stores a secret in the secret store using the operating system's keyring.
pub async fn store_secret(_token: APIToken, request: Json<StoreSecret>) -> Json<StoreSecretResponse> {
match run_secret_store_request(move || store_secret_sync(request)).await {
Ok(response) => response,
Err(issue) => Json(StoreSecretResponse {
success: false,
issue,
issue_code: SecretStoreIssueCode::Unknown,
}),
}
}
/// Performs the blocking part of [`store_secret`].
fn store_secret_sync(request: Json<StoreSecret>) -> Json<StoreSecretResponse> {
let user_name = request.user_name.as_str();
let decrypted_text = match ENCRYPTION.decrypt(&request.secret) {
Ok(text) => text,
@@ -145,6 +174,19 @@ pub struct StoreSecretResponse {
/// Retrieves a secret from the secret store using the operating system's keyring.
pub async fn get_secret(_token: APIToken, request: Json<RequestSecret>) -> Json<RequestedSecret> {
match run_secret_store_request(move || get_secret_sync(request)).await {
Ok(response) => response,
Err(issue) => Json(RequestedSecret {
success: false,
secret: EncryptedText::new(String::from("")),
issue,
issue_code: SecretStoreIssueCode::Unknown,
}),
}
}
/// Performs the blocking part of [`get_secret`].
fn get_secret_sync(request: Json<RequestSecret>) -> Json<RequestedSecret> {
let user_name = request.user_name.as_str();
let service = format!("mindwork-ai-studio::{}", request.destination);
let entry = match Entry::new(service.as_str(), user_name) {
@@ -227,6 +269,19 @@ pub struct RequestedSecret {
/// Deletes a secret from the secret store using the operating system's keyring.
pub async fn delete_secret(_token: APIToken, request: Json<RequestSecret>) -> Json<DeleteSecretResponse> {
match run_secret_store_request(move || delete_secret_sync(request)).await {
Ok(response) => response,
Err(issue) => Json(DeleteSecretResponse {
success: false,
was_entry_found: false,
issue,
issue_code: SecretStoreIssueCode::Unknown,
}),
}
}
/// Performs the blocking part of [`delete_secret`].
fn delete_secret_sync(request: Json<RequestSecret>) -> Json<DeleteSecretResponse> {
let user_name = request.user_name.as_str();
let service = format!("mindwork-ai-studio::{}", request.destination);
let entry = match Entry::new(service.as_str(), user_name) {
+60 -7
View File
@@ -104,33 +104,58 @@ pub fn set_default_tokenizer_path(app_handle: tauri::AppHandle) {
}
pub async fn token_count(_token: APIToken, req: Json<SetTokenText>) -> Json<TokenizerResponse> {
match get_token_count(&req.tokenizer_path, &req.text) {
match run_tokenizer_work(move || get_token_count(&req.tokenizer_path, &req.text)).await {
Ok(count) => Json(TokenizerResponse::available(count)),
Err(e) => Json(TokenizerResponse::unavailable(e)),
}
}
pub async fn validate_tokenizer(_token: APIToken, payload: Json<TokenizerPath>) -> Json<TokenizerResponse> {
match handle_tokenizer_validate(&PathBuf::from(payload.file_path.clone())) {
match run_tokenizer_work(move || handle_tokenizer_validate(&PathBuf::from(payload.file_path.clone()))).await {
Ok(count) => Json(TokenizerResponse::available(count)),
Err(e) => Json(TokenizerResponse::unavailable(e)),
}
}
pub async fn store_tokenizer(_token: APIToken, payload: Json<TokenizerStorage>) -> Json<TokenizerResponse> {
match handle_tokenizer_store(&payload) {
match run_tokenizer_work(move || handle_tokenizer_store(&payload).map_err(|e| e.to_string())).await {
Ok(dest_path) => Json(TokenizerResponse::stored(dest_path)),
Err(e) => Json(TokenizerResponse::unavailable(e.to_string())),
Err(e) => Json(TokenizerResponse::unavailable(e)),
}
}
pub async fn delete_tokenizer(_token: APIToken, payload: Json<TokenizerDelete>) -> Json<TokenizerResponse> {
match handle_tokenizer_delete(&payload) {
match run_tokenizer_work(move || handle_tokenizer_delete(&payload).map_err(|e| e.to_string())).await {
Ok(_) => Json(TokenizerResponse::stored(String::new())),
Err(e) => Json(TokenizerResponse::unavailable(e.to_string())),
Err(e) => Json(TokenizerResponse::unavailable(e)),
}
}
/// Gets the tokenizer for the given path, loading it on the blocking pool when it is not cached yet.
pub async fn get_tokenizer_off_worker(path: String) -> Result<Arc<Tokenizer>, String> {
run_tokenizer_work(move || get_tokenizer(&path)).await
}
/// Runs tokenizer work on the blocking pool.
///
/// Loading a tokenizer reads and parses its file, which takes most of a second for a large one, and
/// it waits for the storage lock while another request copies or deletes a tokenizer. Encoding a
/// long text is CPU-bound as well. On a runtime worker, any of this would hold up every other call
/// of the app, because they all share one HTTP/2 connection.
async fn run_tokenizer_work<T, F>(work: F) -> Result<T, String>
where
T: Send + 'static,
F: FnOnce() -> Result<T, String> + Send + 'static,
{
tokio::task::spawn_blocking(work)
.await
.unwrap_or_else(|error| {
let issue = format!("The tokenizer worker failed: {error}");
error!(Source = "Tokenizer"; "{issue}");
Err(issue)
})
}
fn handle_tokenizer_validate(path: &PathBuf) -> Result<usize, String> {
validate_tokenizer_file(path)
}
@@ -240,7 +265,7 @@ fn tokenizer_cache() -> &'static RwLock<HashMap<PathBuf, Arc<Tokenizer>>> {
TOKENIZERS.get_or_init(|| RwLock::new(HashMap::new()))
}
pub fn get_tokenizer(path: &str) -> Result<Arc<Tokenizer>, String> {
fn get_tokenizer(path: &str) -> Result<Arc<Tokenizer>, String> {
let resolved_path = resolve_tokenizer_path(path)?;
let tokenizer_path = fs::canonicalize(&resolved_path)
.map_err(|e| format!("Could not resolve tokenizer file '{}': {e}", resolved_path.display()))?;
@@ -303,3 +328,31 @@ fn load_tokenizer_from_file(path: &PathBuf) -> Result<Tokenizer, String> {
Tokenizer::from_file(path)
.map_err(|e| format!("Failed to load tokenizer from '{}': {e}", path.display()))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::runtime_api::test_support::assert_runtime_stays_free;
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn a_token_count_waiting_for_the_storage_leaves_the_runtime_free() {
//
// Any file will do: the request waits for the storage lock before it reads the tokenizer.
// Since this file is no tokenizer, the request fails once it gets the lock.
//
let tokenizer_path = std::env::temp_dir().join(format!("ai-studio-tokenizer-{}.json", std::process::id()));
fs::write(&tokenizer_path, "This is not a tokenizer.").unwrap();
let request = token_count(
APIToken::from_hex_text("test"),
Json(SetTokenText {
text: "Hello, world!".to_string(),
tokenizer_path: tokenizer_path.to_string_lossy().to_string(),
}),
);
let response = assert_runtime_stays_free(&TOKENIZER_STORAGE_LOCK, request).await;
fs::remove_file(&tokenizer_path).unwrap();
assert!(!response.success);
}
}