Batch loading for tasks

kamu-data · Jan 7, 2025 · 69f318c · 69f318c
1 parent aa80b66
commit 69f318c
Show file tree

Hide file tree

Showing 4 changed files with 79 additions and 9 deletions.
diff --git a/src/domain/task-system/services/src/task_agent_impl.rs b/src/domain/task-system/services/src/task_agent_impl.rs
@@ -86,19 +86,21 @@ impl TaskAgentImpl {
                 })
                 .try_collect()
                 .await?;
+            let batch_size = running_task_ids.len();
 
-            for running_task_id in &running_task_ids {
-                // TODO: batch loading of tasks
-                let mut task = Task::load(*running_task_id, task_event_store.as_ref())
-                    .await
-                    .int_err()?;
+            let tasks = Task::load_multi(running_task_ids, task_event_store.as_ref())
+                .await
+                .int_err()?;
+
+            for task in tasks {
+                let mut t = task.int_err()?;
 
                 // Requeue
-                task.requeue(self.time_source.now()).int_err()?;
-                task.save(task_event_store.as_ref()).await.int_err()?;
+                t.requeue(self.time_source.now()).int_err()?;
+                t.save(task_event_store.as_ref()).await.int_err()?;
             }
 
-            processed_running_tasks += running_task_ids.len();
+            processed_running_tasks += batch_size;
         }
 
         Ok(())

diff --git a/src/infra/flow-system/postgres/src/postgres_flow_event_store.rs b/src/infra/flow-system/postgres/src/postgres_flow_event_store.rs
@@ -388,7 +388,7 @@ impl EventStore<FlowState> for PostgresFlowEventStore {
                 let event = serde_json::from_value::<FlowEvent>(event_row.event_payload)
                     .map_err(|e| sqlx::Error::Decode(Box::new(e)))?;
 
-                Ok((FlowID::try_from(event_row.flow_id).unwrap(), // todo: handle error
+                Ok((FlowID::try_from(event_row.flow_id).unwrap(), // ids are always > 0
                     EventID::new(event_row.event_id),
                     event))
             })

diff --git a/...ostgres/.sqlx/query-4ab9fb67546df9ac01da967d2ac00ddbbbcf898c7646c89aebc06cf83045686e.json b/...ostgres/.sqlx/query-4ab9fb67546df9ac01da967d2ac00ddbbbcf898c7646c89aebc06cf83045686e.json
diff --git a/src/infra/task-system/postgres/src/postgres_task_event_store.rs b/src/infra/task-system/postgres/src/postgres_task_event_store.rs
@@ -184,6 +184,40 @@ impl EventStore<TaskState> for PostgresTaskEventStore {
         })
     }
 
+    fn get_events_multi(&self, queries: Vec<TaskID>) -> MultiEventStream<TaskID, TaskEvent> {
+        let task_ids: Vec<i64> = queries.iter().map(|id| (*id).try_into().unwrap()).collect();
+
+        Box::pin(async_stream::stream! {
+            let mut tr = self.transaction.lock().await;
+            let connection_mut = tr
+                .connection_mut()
+                .await?;
+
+            let mut query_stream = sqlx::query!(
+                r#"
+                SELECT task_id, event_id, event_payload
+                FROM task_events
+                WHERE task_id = ANY($1)
+                ORDER BY event_id ASC
+                "#,
+                &task_ids,
+            ).try_map(|event_row| {
+                let event = serde_json::from_value::<TaskEvent>(event_row.event_payload)
+                    .map_err(|e| sqlx::Error::Decode(Box::new(e)))?;
+
+                Ok((TaskID::try_from(event_row.task_id).unwrap(), // ids are always > 0
+                    EventID::new(event_row.event_id),
+                    event))
+            })
+            .fetch(connection_mut)
+            .map_err(|e| GetEventsError::Internal(e.int_err()));
+
+            while let Some((task_id, event_id, event)) = query_stream.try_next().await? {
+                yield Ok((task_id, event_id, event));
+            }
+        })
+    }
+
     async fn save_events(
         &self,
         task_id: &TaskID,