collective
Browse / github.com/storytold/artcraft / patch
fix(grok,fal): don't hold std lock across await; poll FAL tasks concurrently
fix(grok,fal): don't hold std lock across await; poll FAL tasks concurrently Addresses two of the small items from #2065 (section 4). ## Changes - **`GrokWrappedWebsocket`**: `send`, `get_response_with_timeout` and `try_next_timeout` held a `std::sync::RwLock` write guard across `.await` (clippy `await_holding_lock`). That can block the runtime or deadlock. Replaced with `tokio::sync::Mutex`. - **`poll_fal_tasks`**: pending FAL jobs were awaited one after another, so a pass cost the sum of all round trips. They are now polled concurrently, capped at 4 (`MAX_CONCURRENT_POLLS`). ## Testing - `SQLX_OFFLINE=true cargo check -p artcraft` passes. - `cargo clippy` no long
- Status
- open
- Score
- 0.00
- Adopters
- 0 other operators
- Author
- HighTecno (imported from pull request #2078, open)
- Base
8abde75133d2d66fe73b8184eaec8372efae2abf- Size
- +73 −73 in 2 files
- Review
- pending
- Shared
- 2026-10-28
- Id
832b33945b1f2e946439b762
How it went
No other operator has reported using it yet.
Diff
raw · crates/api_clients/grok_consumer_client/src/requests/image_websocket/grok_wrapped_websocket.rs crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/fal/poll_fal_tasks.rs
@@ -5,36 +5,34 @@ use futures::stream::FusedStream; use futures::{Stream, TryStreamExt}; use log::{info, warn}; use serde::Serialize;-use std::sync::{Arc, RwLock};+use std::sync::Arc; use std::time::Duration;+use tokio::sync::Mutex; use wreq::ws::message::Message; use wreq::ws::WebSocket; #[derive(Clone)] pub struct GrokWrappedWebsocket {- pub(crate) websocket: Arc<RwLock<WebSocket>>,+ pub(crate) websocket: Arc<Mutex<WebSocket>>, } impl GrokWrappedWebsocket { pub fn new(websocket: WebSocket) -> Self { Self {- websocket: Arc::new(RwLock::new(websocket)),+ websocket: Arc::new(Mutex::new(websocket)), } } pub async fn send(&self, message: String) -> Result<(), GrokError> {- match self.websocket.write() {- Err(err) => Err(GrokClientError::WebsocketLockError.into()),- Ok(mut websocket) => {- let message = Message::text(message);- websocket.send(message)- .await- .map_err(|err| {- GrokClientError::WreqClientError(err)- })?;- Ok(())- },- }+ // NB: tokio's Mutex can be held across `.await`, unlike std's guards.+ let mut websocket = self.websocket.lock().await;+ let message = Message::text(message);+ websocket.send(message)+ .await+ .map_err(|err| {+ GrokClientError::WreqClientError(err)+ })?;+ Ok(()) } pub async fn send_serializable<T: Serialize>(&self, message: T) -> Result<(), GrokError> {@@ -47,53 +45,45 @@ impl GrokWrappedWebsocket { } pub async fn get_response_with_timeout(&self, duration: Duration) -> Result<Option<WebsocketServerMessage>, GrokError> {- match self.websocket.write() {- Err(_) => Err(GrokClientError::WebsocketLockError.into()),- Ok(mut websocket) => {- let result = tokio::time::timeout(duration, websocket.try_next()).await;+ let mut websocket = self.websocket.lock().await;+ let result = tokio::time::timeout(duration, websocket.try_next()).await; - match result {- Err(elapsed) => {- info!("Websocket try_next() elapsed without receiving a message: {:?}", elapsed);- Ok(None) // Timeout elapsed.- }- Ok(inner) => {- let maybe_message = inner.map_err(GrokClientError::WebsocketReadError)?;+ match result {+ Err(elapsed) => {+ info!("Websocket try_next() elapsed without receiving a message: {:?}", elapsed);+ Ok(None) // Timeout elapsed.+ }+ Ok(inner) => {+ let maybe_message = inner.map_err(GrokClientError::WebsocketReadError)?; - match maybe_message {- None => Ok(None),- Some(Message::Text(text)) => {- let maybe_message = WebsocketServerMessage::from_json_str(&text)?;- Ok(Some(maybe_message))- },- Some(_) => {- warn!("Received non-text websocket message.");- Ok(None)- },- }- }+ match maybe_message {+ None => Ok(None),+ Some(Message::Text(text)) => {+ let maybe_message = WebsocketServerMessage::from_json_str(&text)?;+ Ok(Some(maybe_message))+ },+ Some(_) => {+ warn!("Received non-text websocket message.");+ Ok(None)+ }, }- },+ } } } pub async fn try_next_timeout(&self, duration: Duration) -> Result<Option<Message>, GrokClientError> {- match self.websocket.write() {- Err(_) => Err(GrokClientError::WebsocketLockError),- Ok(mut websocket) => {- let result = tokio::time::timeout(duration, websocket.try_next()).await;+ let mut websocket = self.websocket.lock().await;+ let result = tokio::time::timeout(duration, websocket.try_next()).await; - match result {- Err(elapsed) => {- info!("Websocket try_next() elapsed without receiving a message: {:?}", elapsed);- Ok(None) // Timeout elapsed.- }- Ok(inner) => {- let maybe_message = inner.map_err(GrokClientError::WebsocketReadError)?;- Ok(maybe_message)- }- }- },+ match result {+ Err(elapsed) => {+ info!("Websocket try_next() elapsed without receiving a message: {:?}", elapsed);+ Ok(None) // Timeout elapsed.+ }+ Ok(inner) => {+ let maybe_message = inner.map_err(GrokClientError::WebsocketReadError)?;+ Ok(maybe_message)+ } } } @@ -11,11 +11,15 @@ use enums::tauri::tasks::task_status::TaskStatus; use fal_client::creds::fal_api_key::FalApiKey; use fal_client::polling::poll_job_response::poll_job_response::{poll_job_response, PollJobResponseArgs}; use fal_client::polling::poll_job_status::poll_job_status::{poll_job_status, FalJobStatus, PollJobStatusArgs};+use futures::stream::{self, StreamExt}; use log::{error, info, warn}; use sqlite_tasks::queries::task::Task; use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs}; use tauri::AppHandle; +/// Maximum number of FAL jobs polled concurrently in one polling pass.+const MAX_CONCURRENT_POLLS: usize = 4;+ pub async fn poll_fal_tasks( app_handle: &AppHandle, app_env_configs: &AppEnvConfigs,@@ -33,25 +37,31 @@ pub async fn poll_fal_tasks( } }; - for task in fal_tasks {- let result = poll_single_fal_task(- app_handle,- app_env_configs,- app_data_root,- task_database,- storyteller_creds_manager,- &api_key,- task,- ).await;-- if let Err(err) = result {- error!(- "[FalPolling] Error processing task {}: {:?}",- task.id.as_str(),- err,- );- }- }+ let api_key = &api_key;++ // Each task is independent, so poll them concurrently rather than paying+ // the sum of every round trip.+ stream::iter(fal_tasks.iter().copied())+ .for_each_concurrent(MAX_CONCURRENT_POLLS, |task| async move {+ let result = poll_single_fal_task(+ app_handle,+ app_env_configs,+ app_data_root,+ task_database,+ storyteller_creds_manager,+ api_key,+ task,+ ).await;++ if let Err(err) = result {+ error!(+ "[FalPolling] Error processing task {}: {:?}",+ task.id.as_str(),+ err,+ );+ }+ })+ .await; } async fn poll_single_fal_task(