From fc0fba75733aed5e7660d664d0f90c3762dca92d Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 29 Apr 2025 17:34:49 -0400 Subject: [PATCH] task draining / filtering --- Cargo.lock | 1 + Cargo.toml | 1 + crates/desktop/tauri-artcraft/Cargo.toml | 6 +- .../tauri-artcraft/src/state/sora/mod.rs | 2 +- .../src/state/sora/sora_task_queue.rs | 225 +++++++++++++++++- .../src/threads/sora_task_polling_thread.rs | 59 ++++- 6 files changed, 279 insertions(+), 15 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 16d3ceebff..d9ff358e50 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -12066,6 +12066,7 @@ dependencies = [ "filesys", "futures 0.3.29", "image", + "indexmap 2.9.0", "libloading 0.8.6", "log", "memory_caching", diff --git a/Cargo.toml b/Cargo.toml index 170792f44f..f687e3099c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -157,6 +157,7 @@ base64 = "0.22.1" chrono = { version = "0.4.40", features = ["serde"] } cookie = "0.18.1" hex-literal = "1.0.0" +indexmap = "2.9.0" log = "0.4.27" once_cell = "1.21.3" rand = "0.9.0" diff --git a/crates/desktop/tauri-artcraft/Cargo.toml b/crates/desktop/tauri-artcraft/Cargo.toml index 6e3a1b4138..2b92f13206 100644 --- a/crates/desktop/tauri-artcraft/Cargo.toml +++ b/crates/desktop/tauri-artcraft/Cargo.toml @@ -35,7 +35,7 @@ tokens = { workspace = true } # External anyhow = { workspace = true } -base64 = "0.22.1" +base64 = { workspace = true } bytes = "1.10" chrono = { workspace = true } const_format = "0.2.34" @@ -43,11 +43,11 @@ directories = "6.0.0" dirs = "6.0.0" futures = "0.3" image = "0.25.5" +indexmap = { workspace = true } log = { workspace = true } once_cell = { workspace = true } rand = { workspace = true } -# NB(bt,2025-03-10): default-features = false on Windows while I debug openssl / rustls -reqwest = { workspace = true, features = ["multipart", "stream"], default-features = false } +reqwest = { workspace = true, features = ["multipart", "stream"], default-features = false } # NB(bt,2025-03-10): default-features = false on Windows while I debug openssl / rustls serde = { workspace = true } serde_derive = { workspace = true } serde_json = { workspace = true } diff --git a/crates/desktop/tauri-artcraft/src/state/sora/mod.rs b/crates/desktop/tauri-artcraft/src/state/sora/mod.rs index d0ffa3f7c7..f2b2834e3c 100644 --- a/crates/desktop/tauri-artcraft/src/state/sora/mod.rs +++ b/crates/desktop/tauri-artcraft/src/state/sora/mod.rs @@ -2,4 +2,4 @@ pub mod read_sora_credentials_from_disk; pub mod read_sora_credentials_from_disk_with_upgrade; pub mod sora_credential_holder; pub mod sora_credential_manager; -mod sora_task_queue; \ No newline at end of file +pub mod sora_task_queue; \ No newline at end of file diff --git a/crates/desktop/tauri-artcraft/src/state/sora/sora_task_queue.rs b/crates/desktop/tauri-artcraft/src/state/sora/sora_task_queue.rs index 27ae17cd7b..10a5a1c025 100644 --- a/crates/desktop/tauri-artcraft/src/state/sora/sora_task_queue.rs +++ b/crates/desktop/tauri-artcraft/src/state/sora/sora_task_queue.rs @@ -1,5 +1,8 @@ -use chrono::{Date, DateTime, Utc}; -use std::collections::HashMap; +use std::collections::HashSet; +use chrono::{DateTime, Utc}; +use errors::AnyhowResult; +use indexmap::IndexMap; +use openai_sora_client::requests::image_gen::image_gen_status::TaskId; use std::sync::{Arc, RwLock}; #[derive(Clone)] @@ -7,8 +10,222 @@ pub struct TaskDetails { pub enqueue_time: DateTime } +impl TaskDetails { + pub fn new() -> Self { + Self { + enqueue_time: Utc::now() + } + } +} + +#[derive(Clone)] +pub struct TaskIdAndDetails { + pub task_id: TaskId, + pub details: TaskDetails, +} + #[derive(Clone)] pub struct SoraTaskQueue { - - queue: Arc>> + // Insertion-order map of task IDs to task details + queue: Arc>> } + +impl SoraTaskQueue { + pub fn new() -> Self { + Self { + queue: Arc::new(RwLock::new(IndexMap::new())) + } + } + + /// Determine if queue is empty + /// This is an O(1) operation. + pub fn is_empty(&self) -> AnyhowResult { + match self.queue.read() { + Err(err) => { + Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)) + } + Ok(lock) => Ok(lock.is_empty()) + } + } + + /// Get length of set + /// This is an O(1) operation. + pub fn len(&self) -> AnyhowResult { + match self.queue.read() { + Err(err) => { + Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)) + } + Ok(lock) => Ok(lock.len()) + } + } + + /// Determine if the key is in the queue + /// This is an O(1) operation. + pub fn contains_key(&self, task_id: &TaskId) -> AnyhowResult { + match self.queue.read() { + Err(err) => { + Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)) + } + Ok(lock) => Ok(lock.contains_key(task_id)) + } + } + + /// Get the first enqueued task + /// This is an O(1) operation. + pub fn first(&self) -> AnyhowResult> { + match self.queue.read() { + Err(err) => { + Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)) + } + Ok(lock) => match lock.first() { + None => Ok(None), + Some((task_id, task_details)) => Ok(Some(TaskIdAndDetails { + task_id: task_id.clone(), + details: task_details.clone(), + })) + } + } + } + + /// Get the last enqueued task + /// This is an O(1) operation. + pub fn last(&self) -> AnyhowResult> { + match self.queue.read() { + Err(err) => { + Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)) + } + Ok(lock) => match lock.last() { + None => Ok(None), + Some((task_id, task_details)) => Ok(Some(TaskIdAndDetails { + task_id: task_id.clone(), + details: task_details.clone(), + })) + } + } + } + + /// Insert a new task. If a task already exists, we update its value. + /// This is an O(1) operation. + pub fn insert(&self, task_id: &TaskId) -> AnyhowResult> { + match self.queue.write() { + Err(err) => { + Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)) + } + Ok(mut lock) => { + // NB: Probably fine to compute TaskDetails here for now. We may eventually expose this if we include additional info. + let task_details = TaskDetails::new(); + Ok(lock.insert(task_id.clone(), task_details)) + } + } + } + + /// Remove a task by ID. + /// This is an O(n) operation. + pub fn remove(&self, task_id: &TaskId) -> AnyhowResult> { + match self.queue.write() { + Err(err) => { + Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)) + } + Ok(mut lock) => { + Ok(lock.shift_remove(task_id)) + } + } + } + + /// Remove a list of tasks by ID. + /// This is an O(n) operation. + pub fn remove_list(&self, task_ids: &[&TaskId]) -> AnyhowResult<()> { + match self.queue.write() { + Err(err) => { + Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)) + } + Ok(mut lock) => { + // TODO: Cleaner way to do this? + let mut set : HashSet<&TaskId, _> = HashSet::with_capacity(task_ids.len()); + set.extend(task_ids); + lock.retain(|k, v| !set.contains(k)); + Ok(()) + } + } + } +} + +#[cfg(test)] +mod tests { + use errors::AnyhowResult; + use openai_sora_client::requests::image_gen::image_gen_status::TaskId; + use crate::state::sora::sora_task_queue::SoraTaskQueue; + + #[test] + fn remove_list() -> AnyhowResult<()> { + let task_queue = SoraTaskQueue::new(); + + assert!(task_queue.is_empty()?); + assert_eq!(task_queue.len()?, 0); + + task_queue.insert(&TaskId("a".to_string()))?; + task_queue.insert(&TaskId("b".to_string()))?; + task_queue.insert(&TaskId("c".to_string()))?; + task_queue.insert(&TaskId("d".to_string()))?; + task_queue.insert(&TaskId("e".to_string()))?; + task_queue.insert(&TaskId("f".to_string()))?; + + assert!(!task_queue.is_empty()?); + assert_eq!(task_queue.len()?, 6); + + task_queue.insert(&TaskId("a".to_string()))?; + task_queue.insert(&TaskId("b".to_string()))?; + task_queue.insert(&TaskId("c".to_string()))?; + task_queue.insert(&TaskId("d".to_string()))?; + task_queue.insert(&TaskId("e".to_string()))?; + task_queue.insert(&TaskId("f".to_string()))?; + + assert!(!task_queue.is_empty()?); + assert_eq!(task_queue.len()?, 6); + + task_queue.remove_list(&[ + &TaskId("a".to_string()), + &TaskId("c".to_string()), + &TaskId("e".to_string()), + ])?; + + assert!(!task_queue.is_empty()?); + assert_eq!(task_queue.len()?, 3); + + assert!(task_queue.contains_key(&TaskId("b".to_string()))?); + assert!(task_queue.contains_key(&TaskId("d".to_string()))?); + assert!(task_queue.contains_key(&TaskId("f".to_string()))?); + + assert!(!task_queue.contains_key(&TaskId("a".to_string()))?); + assert!(!task_queue.contains_key(&TaskId("c".to_string()))?); + assert!(!task_queue.contains_key(&TaskId("e".to_string()))?); + + task_queue.remove_list(&[ + &TaskId("foo".to_string()), + &TaskId("bar".to_string()), + &TaskId("baz".to_string()), + ])?; + + assert!(!task_queue.is_empty()?); + assert_eq!(task_queue.len()?, 3); + + assert!(task_queue.contains_key(&TaskId("b".to_string()))?); + assert!(task_queue.contains_key(&TaskId("d".to_string()))?); + assert!(task_queue.contains_key(&TaskId("f".to_string()))?); + + task_queue.remove_list(&[ + &TaskId("b".to_string()), + &TaskId("d".to_string()), + &TaskId("f".to_string()), + ])?; + + assert!(task_queue.is_empty()?); + assert_eq!(task_queue.len()?, 0); + + assert!(!task_queue.contains_key(&TaskId("b".to_string()))?); + assert!(!task_queue.contains_key(&TaskId("d".to_string()))?); + assert!(!task_queue.contains_key(&TaskId("f".to_string()))?); + + Ok(()) + } +} \ No newline at end of file diff --git a/crates/desktop/tauri-artcraft/src/threads/sora_task_polling_thread.rs b/crates/desktop/tauri-artcraft/src/threads/sora_task_polling_thread.rs index 011f49ebd0..1557f7fbe2 100644 --- a/crates/desktop/tauri-artcraft/src/threads/sora_task_polling_thread.rs +++ b/crates/desktop/tauri-artcraft/src/threads/sora_task_polling_thread.rs @@ -1,28 +1,44 @@ use crate::state::app_dir::AppDataRoot; use crate::state::sora::sora_credential_manager::SoraCredentialManager; +use crate::state::sora::sora_task_queue::SoraTaskQueue; use errors::AnyhowResult; use log::{error, info}; use openai_sora_client::recipes::list_sora_task_status_with_session_auto_renew::{list_sora_task_status_with_session_auto_renew, StatusRequestArgs}; +use openai_sora_client::requests::image_gen::image_gen_status::{TaskId, TaskStatus}; use tauri::AppHandle; pub async fn sora_session_login_thread( app_data_root: AppDataRoot, - sora_creds_manager: SoraCredentialManager + sora_creds_manager: SoraCredentialManager, + sora_task_queue: SoraTaskQueue, ) -> ! { loop { - if let Err(err) = polling_loop(&sora_creds_manager).await { + let res = polling_loop(&sora_creds_manager, &sora_task_queue).await; + if let Err(err) = res { error!("An error occurred: {:?}", err); } tokio::time::sleep(std::time::Duration::from_millis(30_000)).await; } } -async fn polling_loop(sora_creds_manager: &SoraCredentialManager) -> AnyhowResult<()> { +async fn polling_loop( + sora_creds_manager: &SoraCredentialManager, + sora_task_queue: &SoraTaskQueue, +) -> AnyhowResult<()> { loop { + if sora_task_queue.is_empty()? { + // No need to poll if we don't have pending tasks. + tokio::time::sleep(std::time::Duration::from_millis(1_000)).await; + continue; + } + + info!("Task queue has {} pending tasks.", sora_task_queue.len()?); + let creds = sora_creds_manager.get_credentials_required()?; let (response, maybe_new_creds) = list_sora_task_status_with_session_auto_renew(StatusRequestArgs { limit: None, + // TODO: How can we use the task id to poll better? Our existing code doesn't seem to illuminate this. before: None, credentials: creds, }).await?; @@ -32,13 +48,42 @@ async fn polling_loop(sora_creds_manager: &SoraCredentialManager) -> AnyhowResul sora_creds_manager.set_credentials(&new_creds)?; } + let mut succeeded_tasks = Vec::new(); + let mut failed_tasks = Vec::new(); + // TODO: The cursoring logic likely needs to improve. for task in response.task_responses { - info!("Task ID: {}", task.id); - info!("Task Status: {:?}", task.status); - // Handle the task status as needed + let status = TaskStatus::from_str(&task.status); + + match status { + TaskStatus::Succeeded => { + succeeded_tasks.push(task); + } + TaskStatus::Failed => { + failed_tasks.push(task); + } + TaskStatus::Queued => {} + TaskStatus::Running => {} + TaskStatus::Unknown(_) => {} + } } - tokio::time::sleep(std::time::Duration::from_millis(10_000)).await; + let failed_task_ids: Vec<&TaskId> = failed_tasks + .iter() + .map(|task| &task.id) + .collect(); + + sora_task_queue.remove_list(&failed_task_ids)?; + + // TODO: Handle succeeded tasks. + + let succeeded_task_ids : Vec<&TaskId> = succeeded_tasks + .iter() + .map(|task| &task.id) + .collect(); + + sora_task_queue.remove_list(&succeeded_task_ids)?; + + tokio::time::sleep(std::time::Duration::from_millis(2_000)).await; } }