test(lib_core): update task manager tests

refs: #5

wip: add test to task_manager

refs: #5
This commit is contained in:
Alexander Navarro 2025-05-20 16:49:41 -04:00
parent c52a497075
commit cfa2474606
11 changed files with 289 additions and 38 deletions

29
.idea/codeStyles/Project.xml generated Normal file
View file

@ -0,0 +1,29 @@
<component name="ProjectCodeStyleConfiguration">
<code_scheme name="Project" version="173">
<RsCodeStyleSettings>
<option name="INDENT_WHERE_CLAUSE" value="true" />
</RsCodeStyleSettings>
<SqlCodeStyleSettings version="7">
<option name="KEYWORD_CASE" value="2" />
<option name="SUBQUERY_OPENING" value="1" />
<option name="SUBQUERY_CONTENT" value="4" />
<option name="SUBQUERY_CLOSING" value="4" />
<option name="SUBQUERY_PAR_SPACE_BEFORE" value="1" />
<option name="INSERT_INTO_NL" value="2" />
<option name="INSERT_COLLAPSE_MULTI_ROW_VALUES" value="true" />
<option name="INSERT_MATRIX_ALIGN" value="true" />
<option name="INSERT_MATRIX_INCLUDING_HEADER" value="true" />
<option name="FROM_ALIGN_JOIN_TABLES" value="true" />
<option name="FROM_ALIGN_ALIASES" value="true" />
</SqlCodeStyleSettings>
<codeStyleSettings language="Rust">
<indentOptions>
<option name="INDENT_SIZE" value="2" />
<option name="CONTINUATION_INDENT_SIZE" value="2" />
<option name="TAB_SIZE" value="2" />
<option name="USE_TAB_CHARACTER" value="true" />
<option name="SMART_TABS" value="true" />
</indentOptions>
</codeStyleSettings>
</code_scheme>
</component>

View file

@ -1,5 +1,6 @@
<component name="ProjectCodeStyleConfiguration"> <component name="ProjectCodeStyleConfiguration">
<state> <state>
<option name="USE_PER_PROJECT_SETTINGS" value="true" />
<option name="PREFERRED_PROJECT_CODE_STYLE" value="Default" /> <option name="PREFERRED_PROJECT_CODE_STYLE" value="Default" />
</state> </state>
</component> </component>

1
Cargo.lock generated
View file

@ -1066,6 +1066,7 @@ dependencies = [
"tabled", "tabled",
"thiserror", "thiserror",
"tokio", "tokio",
"tokio-stream",
"tracing", "tracing",
"tracing-core", "tracing-core",
"tracing-subscriber", "tracing-subscriber",

View file

@ -6,6 +6,7 @@ edition = "2024"
[dependencies] [dependencies]
directories = "6.0.0" directories = "6.0.0"
tokio = { version = "1.45.0", features = ["default", "rt", "rt-multi-thread", "macros"] } tokio = { version = "1.45.0", features = ["default", "rt", "rt-multi-thread", "macros"] }
tokio-stream = "0.1.17"
sqlx = { version = "0.8", features = ["runtime-tokio", "sqlite", "chrono", "migrate", "uuid"] } sqlx = { version = "0.8", features = ["runtime-tokio", "sqlite", "chrono", "migrate", "uuid"] }
clap = { version = "4.5.37", features = ["derive"] } clap = { version = "4.5.37", features = ["derive"] }
serde = { version = "1.0.219", features = ["derive"] } serde = { version = "1.0.219", features = ["derive"] }

View file

@ -15,10 +15,10 @@ impl TaskPagination {
} }
pub fn next(&self) -> Self { pub fn next(&self) -> Self {
Self { Self {
offset: self.offset.saturating_add(self.limit.unwrap_or(0)), offset: self.offset.saturating_add(self.limit.unwrap_or(0)),
..self.clone() ..self.clone()
} }
} }
pub fn prev(&self) -> Self { pub fn prev(&self) -> Self {
@ -45,20 +45,17 @@ impl TaskPagination {
} }
pub struct TasksPage<T: TaskPayload> { pub struct TasksPage<T: TaskPayload> {
tasks: Vec<Task<T>>, tasks: Vec<Task<T>>,
page: TaskPagination page: TaskPagination,
} }
impl<T: TaskPayload> TasksPage<T> { impl<T: TaskPayload> TasksPage<T> {
fn new(tasks: Vec<Task<T>>, page: TaskPagination) -> Self { fn new(tasks: Vec<Task<T>>, page: TaskPagination) -> Self {
Self { Self { tasks, page }
tasks,
page
}
} }
pub fn next(&self) -> TaskPagination { pub fn next(&self) -> TaskPagination {
self.page.next() self.page.next()
} }
pub fn prev(&self) -> TaskPagination { pub fn prev(&self) -> TaskPagination {
@ -67,9 +64,13 @@ impl<T: TaskPayload> TasksPage<T> {
} }
pub trait TaskStorage<T: TaskPayload> { pub trait TaskStorage<T: TaskPayload> {
async fn insert_tasks<'a, I: IntoIterator<Item=&'a Task<T>>>(&self, tasks: I) -> crate::Result<()>; async fn insert_tasks<'a, I: IntoIterator<Item = &'a Task<T>>>(
fn get_tasks(&self, options: TaskStatus) -> impl Stream<Item = crate::Result<Task<T>>>; &self,
tasks: I,
) -> crate::Result<()>;
fn get_tasks(&self, task_status: TaskStatus) -> impl Stream<Item = crate::Result<Task<T>>>;
fn listen_tasks(&self, task_status: TaskStatus) -> impl Stream<Item = crate::Result<Task<T>>>;
async fn get_paginated_tasks(&self, page: TaskPagination) -> crate::Result<TasksPage<T>>; async fn get_paginated_tasks(&self, page: TaskPagination) -> crate::Result<TasksPage<T>>;
} }

View file

@ -108,6 +108,11 @@ impl<T: TaskPayload> TaskStorage<T> for Sqlite {
query.fetch(&self.pool).err_into::<crate::Error>() query.fetch(&self.pool).err_into::<crate::Error>()
} }
fn listen_tasks(&self, task_status: TaskStatus) -> impl Stream<Item=crate::error::Result<Task<T>>> {
futures::stream::empty()
}
async fn get_paginated_tasks(&self, page: TaskPagination) -> crate::Result<TasksPage<T>> { async fn get_paginated_tasks(&self, page: TaskPagination) -> crate::Result<TasksPage<T>> {
let mut builder: QueryBuilder<'_, sqlx::Sqlite> = QueryBuilder::new( let mut builder: QueryBuilder<'_, sqlx::Sqlite> = QueryBuilder::new(
"select id, payload_key, payload, status_id, created_at, updated_at from tasks ", "select id, payload_key, payload, status_id, created_at, updated_at from tasks ",
@ -170,6 +175,7 @@ mod tests {
} }
#[sqlx::test(migrator = "MIGRATIONS")] #[sqlx::test(migrator = "MIGRATIONS")]
#[traced_test]
async fn it_save_tasks(pool: SqlitePool) -> sqlx::Result<()> { async fn it_save_tasks(pool: SqlitePool) -> sqlx::Result<()> {
let (sqlite, tasks) = setup(pool.clone(), 100); let (sqlite, tasks) = setup(pool.clone(), 100);

View file

@ -5,6 +5,9 @@ use std::fmt::Display;
use tabled::Tabled; use tabled::Tabled;
mod manager; mod manager;
mod jobs;
mod worker;
mod bus;
#[derive(sqlx::Type, Debug, Clone)] #[derive(sqlx::Type, Debug, Clone)]
#[repr(u8)] #[repr(u8)]
@ -43,8 +46,6 @@ impl<T: Serialize + DeserializeOwned + Send + Unpin + 'static + std::fmt::Debug>
{ {
} }
pub type TaskJob<T> = fn(&Task<T>) -> TaskStatus;
#[derive(sqlx::FromRow, Tabled, Debug)] #[derive(sqlx::FromRow, Tabled, Debug)]
pub struct Task<T: TaskPayload> { pub struct Task<T: TaskPayload> {
id: u32, id: u32,
@ -76,9 +77,7 @@ impl<T: TaskPayload> Task<T> {
pub fn payload(&self) -> &T { pub fn payload(&self) -> &T {
&self.payload &self.payload
} }
}
impl<T: TaskPayload> Task<T> {
pub fn get_key(&self) -> String { pub fn get_key(&self) -> String {
self.payload_key.clone() self.payload_key.clone()
} }

View file

@ -0,0 +1,4 @@
#[derive(Clone)]
pub enum Bus {
Local,
}

View file

@ -0,0 +1,3 @@
use crate::tasks::{Task, TaskStatus};
pub type TaskJob<T> = fn(&Task<T>) -> TaskStatus;

View file

@ -1,32 +1,190 @@
use crate::database::TaskStorage;
use crate::tasks::bus::Bus;
use crate::tasks::jobs::TaskJob;
use crate::tasks::{Task, TaskPayload, TaskStatus};
use futures::StreamExt; use futures::StreamExt;
use std::marker::PhantomData; use std::marker::PhantomData;
use crate::database::TaskStorage; use std::pin::pin;
use crate::tasks::{Task, TaskJob, TaskPayload, TaskStatus}; use tokio::sync::mpsc::Receiver;
use tokio::sync::{mpsc, oneshot};
use tokio::sync::oneshot::Sender;
use crate::tasks::worker::TaskMessage;
struct TaskManager<S: TaskPayload, T: TaskStorage<S>> pub enum RateLimit {
{ Buffer(usize),
Rate(usize),
Ticks(usize),
None,
}
pub struct ManagerOptions {
rate_limit: RateLimit,
bus: Bus,
}
impl ManagerOptions {
pub fn new() -> Self {
Self::default()
}
pub fn with_rate_limit(mut self, rate_limit: RateLimit) -> Self {
self.rate_limit = rate_limit;
self
}
}
impl Default for ManagerOptions {
fn default() -> Self {
Self {
rate_limit: RateLimit::None,
bus: Bus::Local,
}
}
}
struct TaskManager<S: TaskPayload, T: TaskStorage<S>> {
storage: T, storage: T,
options: ManagerOptions,
_marker: PhantomData<S>, _marker: PhantomData<S>,
} }
impl<S: TaskPayload, T: TaskStorage<S>> TaskManager<S, T> { impl<S: TaskPayload, T: TaskStorage<S>> TaskManager<S, T> {
pub fn new(storage: T) -> Self { pub fn new(storage: T, options: ManagerOptions) -> Self {
Self { Self {
storage, storage,
_marker: PhantomData, options,
} _marker: PhantomData,
}
} }
pub async fn run_tasks(&self, func: TaskJob<S>) -> crate::Result<()> { pub async fn run_tasks(&self, mut task_sink: TaskMessage<S>) -> crate::Result<()> {
let rows = self.storage.get_tasks(TaskStatus::Pending); let rows = self.storage.get_tasks(TaskStatus::Pending);
let listener = self.storage.listen_tasks(TaskStatus::Pending);
let result: Vec<(Task<S>, TaskStatus)> = rows.map(|x| { let mut queue = pin!(rows.chain(listener));
let task = x.unwrap();
let status = func(&task);
(task, status) while let Some(task) = queue.next().await {
}).collect().await; let task = match task {
Ok(task) => task,
Err(e) => {
continue
}
};
let sink = match task_sink.recv().await {
Some(s) => s,
None => break, // sink has stoped requesting tasks
};
if let Err(_) = sink.send(task) {
continue;
}
// (task, status)
}
Ok(()) Ok(())
} }
} }
#[cfg(test)]
mod tests {
use super::*;
use crate::database::{TaskPagination, TasksPage};
use async_stream::stream;
use fake::{Dummy, Fake, Faker};
use futures::Stream;
use serde::{Deserialize, Serialize};
use sqlx::types::Uuid;
use sync::mpsc;
use tokio::sync;
use tokio_stream::wrappers::ReceiverStream;
use tracing_test::traced_test;
use crate::error::Error;
use crate::tasks::worker::{Worker, WorkerManager};
#[derive(Dummy, Serialize, Deserialize, Debug)]
struct DummyTaskPayload {
key: Uuid,
_foo: String,
_bar: String,
}
struct DummyTaskStorage {}
impl TaskStorage<DummyTaskPayload> for DummyTaskStorage {
async fn insert_tasks<'a, I: IntoIterator<Item = &'a Task<DummyTaskPayload>>>(
&self,
_: I,
) -> crate::error::Result<()> {
todo!()
}
fn get_tasks(
&self,
task_status: TaskStatus,
) -> impl Stream<Item = crate::Result<Task<DummyTaskPayload>>> {
let payloads: Vec<DummyTaskPayload> = Faker.fake();
let tasks = payloads.into_iter().enumerate().map(move |(i, item)| {
Ok(Task::new((i + 1).to_string(), item, task_status.clone()))
});
futures::stream::iter(tasks)
}
fn listen_tasks(
&self,
task_status: TaskStatus,
) -> impl Stream<Item = crate::Result<Task<DummyTaskPayload>>> {
let (tx, rx) = mpsc::channel::<crate::Result<Task<DummyTaskPayload>>>(10);
tokio::spawn(async move {
for _ in 0..10 {
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
let payload: DummyTaskPayload = Faker.fake();
let task_status: TaskStatus = task_status.clone();
let task = Ok(Task::new(payload.key.to_string(), payload, task_status));
if let Err(_) = tx.send(task).await {
break;
}
}
});
ReceiverStream::new(rx)
}
async fn get_paginated_tasks(
&self,
_: TaskPagination,
) -> crate::error::Result<TasksPage<DummyTaskPayload>> {
todo!()
}
}
struct DummyWorker;
impl Worker<DummyTaskPayload> for DummyWorker {
fn process_job(task: &Task<DummyTaskPayload>) -> crate::error::Result<()> {
println!("{:#?}", task);
Ok(())
}
async fn on_job_failure(task: &Task<DummyTaskPayload>, error: Error) -> crate::error::Result<()> {
println!("{:#?} {:?}", task, error);
Ok(())
}
}
#[tokio::test]
#[traced_test]
async fn manager_runs() {
let execute_options = ManagerOptions::new();
let local_worker_sink = WorkerManager::get_listener_sink::<DummyTaskPayload, DummyWorker>(execute_options.bus.clone());
let task_manager = TaskManager::new(DummyTaskStorage {}, execute_options);
task_manager.run_tasks(local_worker_sink).await.unwrap()
}
}

View file

@ -0,0 +1,48 @@
use crate::error::Error;
use crate::tasks::bus::Bus;
use crate::tasks::{Task, TaskPayload};
use tokio::sync::mpsc::Receiver;
use tokio::sync::oneshot::Sender;
use tokio::sync::{mpsc, oneshot};
pub type TaskMessage<T> = Receiver<Sender<Task<T>>>;
pub struct WorkerManager;
impl WorkerManager {
pub fn get_listener_sink<T: TaskPayload, W: Worker<T>>(bus: Bus) -> TaskMessage<T> {
match bus {
Bus::Local => {
let (bus_tx, bus_rx) = mpsc::channel(100);
tokio::spawn(async move {
loop {
// TODO: properly catch errors
let (tx, rx) = oneshot::channel();
// Request a task
bus_tx.send(tx).await.unwrap();
// Wait for a task to be returned
let task = rx.await.unwrap();
W::process_job(&task).unwrap();
}
});
bus_rx
}
}
}
}
pub trait Worker<T: TaskPayload> {
async fn pre_process_job(task: &Task<T>) -> crate::Result<()> {
Ok(())
}
fn process_job(task: &Task<T>) -> crate::Result<()>;
async fn post_process_job(task: &Task<T>) -> crate::Result<()> {
Ok(())
}
async fn on_job_failure(task: &Task<T>, error: Error) -> crate::Result<()>;
}