Allow setting a custom clock on the QueueWorker & add one to the

TestState

Author: sandhose

Cargo.lock

@@ -173,8 +173,9 @@ impl Options {
             test_mailer_in_background(&mailer, Duration::from_secs(30));
 
             info!("Starting task worker");
-            mas_tasks::init(
+            mas_tasks::init_and_run(
                 PgRepositoryFactory::new(pool.clone()),
+                SystemClock::default(),
                 &mailer,
                 homeserver_connection.clone(),
                 url_builder.clone(),

crates/cli/src/commands/server.rs

@@ -10,6 +10,7 @@ use clap::Parser;
 use figment::Figment;
 use mas_config::{AppConfig, ConfigurationSection};
 use mas_router::UrlBuilder;
+use mas_storage::SystemClock;
 use mas_storage_pg::PgRepositoryFactory;
 use tracing::{info, info_span};
 
@@ -63,8 +64,9 @@ impl Options {
         drop(config);
 
         info!("Starting task scheduler");
-        mas_tasks::init(
+        mas_tasks::init_and_run(
             PgRepositoryFactory::new(pool.clone()),
+            SystemClock::default(),
             &mailer,
             conn,
             url_builder,

crates/cli/src/commands/worker.rs

@@ -72,6 +72,7 @@ mas-axum-utils.workspace = true
 mas-config.workspace = true
 mas-context.workspace = true
 mas-data-model.workspace = true
+mas-email.workspace = true
 mas-http.workspace = true
 mas-i18n.workspace = true
 mas-iana.workspace = true
@@ -83,6 +84,7 @@ mas-policy.workspace = true
 mas-router.workspace = true
 mas-storage.workspace = true
 mas-storage-pg.workspace = true
+mas-tasks.workspace = true
 mas-templates.workspace = true
 oauth2-types.workspace = true
 zxcvbn.workspace = true

crates/handlers/Cargo.toml

@@ -29,6 +29,7 @@ use mas_axum_utils::{
 };
 use mas_config::RateLimitingConfig;
 use mas_data_model::SiteConfig;
+use mas_email::{MailTransport, Mailer};
 use mas_i18n::Translator;
 use mas_keystore::{Encrypter, JsonWebKey, JsonWebKeySet, Keystore, PrivateKey};
 use mas_matrix::{HomeserverConnection, MockHomeserverConnection};
@@ -39,6 +40,7 @@ use mas_storage::{
     clock::MockClock,
 };
 use mas_storage_pg::PgRepositoryFactory;
+use mas_tasks::QueueWorker;
 use mas_templates::{SiteConfigExt, Templates};
 use oauth2_types::{registration::ClientRegistrationResponse, requests::AccessTokenResponse};
 use rand::SeedableRng;
@@ -113,6 +115,7 @@ pub(crate) struct TestState {
     pub rng: Arc<Mutex<ChaChaRng>>,
     pub http_client: reqwest::Client,
     pub task_tracker: TaskTracker,
+    queue_worker: Arc<Mutex<QueueWorker>>,
 
     #[allow(dead_code)] // It is used, as it will cancel the CancellationToken when dropped
     cancellation_drop_guard: Arc<DropGuard>,
@@ -235,6 +238,27 @@ impl TestState {
             shutdown_token.child_token(),
         );
 
+        let mailer = Mailer::new(
+            templates.clone(),
+            MailTransport::blackhole(),
+            "[email protected]".parse().unwrap(),
+            "[email protected]".parse().unwrap(),
+        );
+
+        let queue_worker = mas_tasks::init(
+            PgRepositoryFactory::new(pool.clone()),
+            Arc::clone(&clock),
+            &mailer,
+            homeserver_connection.clone(),
+            url_builder.clone(),
+            &site_config,
+            shutdown_token.child_token(),
+        )
+        .await
+        .unwrap();
+
+        let queue_worker = Arc::new(Mutex::new(queue_worker));
+
         Ok(Self {
             repository_factory: PgRepositoryFactory::new(pool),
             templates,
@@ -254,6 +278,7 @@ impl TestState {
             rng,
             http_client,
             task_tracker,
+            queue_worker,
             cancellation_drop_guard: Arc::new(shutdown_token.drop_guard()),
         })
     }

crates/handlers/src/test_utils.rs

@@ -10,14 +10,16 @@ use mas_data_model::SiteConfig;
 use mas_email::Mailer;
 use mas_matrix::HomeserverConnection;
 use mas_router::UrlBuilder;
-use mas_storage::{BoxRepository, Clock, RepositoryError, RepositoryFactory, SystemClock};
+use mas_storage::{BoxRepository, Clock, RepositoryError, RepositoryFactory};
 use mas_storage_pg::PgRepositoryFactory;
 use new_queue::QueueRunnerError;
 use opentelemetry::metrics::Meter;
 use rand::SeedableRng;
 use sqlx::{Pool, Postgres};
 use tokio_util::{sync::CancellationToken, task::TaskTracker};
 
+pub use crate::new_queue::QueueWorker;
+
 mod database;
 mod email;
 mod matrix;
@@ -99,29 +101,31 @@ impl State {
     }
 }
 
-/// Initialise the workers.
+/// Initialise the worker, without running it.
+///
+/// This is mostly useful for tests.
 ///
 /// # Errors
 ///
 /// This function can fail if the database connection fails.
 pub async fn init(
     repository_factory: PgRepositoryFactory,
+    clock: impl Clock + 'static,
     mailer: &Mailer,
     homeserver: impl HomeserverConnection + 'static,
     url_builder: UrlBuilder,
     site_config: &SiteConfig,
     cancellation_token: CancellationToken,
-    task_tracker: &TaskTracker,
-) -> Result<(), QueueRunnerError> {
+) -> Result<QueueWorker, QueueRunnerError> {
     let state = State::new(
         repository_factory,
-        SystemClock::default(),
+        clock,
         mailer.clone(),
         homeserver,
         url_builder,
         site_config.clone(),
     );
-    let mut worker = self::new_queue::QueueWorker::new(state, cancellation_token).await?;
+    let mut worker = QueueWorker::new(state, cancellation_token).await?;
 
     worker
         .register_handler::<mas_storage::queue::CleanupExpiredTokensJob>()
@@ -157,6 +161,35 @@ pub async fn init(
             mas_storage::queue::PruneStalePolicyDataJob,
         );
 
+    Ok(worker)
+}
+
+/// Initialise the worker and run it.
+///
+/// # Errors
+///
+/// This function can fail if the database connection fails.
+pub async fn init_and_run(
+    repository_factory: PgRepositoryFactory,
+    clock: impl Clock + 'static,
+    mailer: &Mailer,
+    homeserver: impl HomeserverConnection + 'static,
+    url_builder: UrlBuilder,
+    site_config: &SiteConfig,
+    cancellation_token: CancellationToken,
+    task_tracker: &TaskTracker,
+) -> Result<(), QueueRunnerError> {
+    let worker = init(
+        repository_factory,
+        clock,
+        mailer,
+        homeserver,
+        url_builder,
+        site_config,
+        cancellation_token,
+    )
+    .await?;
+
     task_tracker.spawn(worker.run());
 
     Ok(())

crates/tasks/src/lib.rs

@@ -214,7 +214,7 @@ impl QueueWorker {
         skip_all,
         fields(worker.id)
     )]
-    pub async fn new(
+    pub(crate) async fn new(
         state: State,
         cancellation_token: CancellationToken,
     ) -> Result<Self, QueueRunnerError> {
@@ -289,7 +289,7 @@ impl QueueWorker {
         })
     }
 
-    pub fn register_handler<T: RunnableJob + InsertableJob>(&mut self) -> &mut Self {
+    pub(crate) fn register_handler<T: RunnableJob + InsertableJob>(&mut self) -> &mut Self {
         // There is a potential panic here, which is fine as it's going to be caught
         // within the job task
         let factory = |payload: JobPayload| {
@@ -302,7 +302,7 @@ impl QueueWorker {
         self
     }
 
-    pub fn add_schedule<T: InsertableJob>(
+    pub(crate) fn add_schedule<T: InsertableJob>(
         &mut self,
         schedule_name: &'static str,
         expression: Schedule,
@@ -320,7 +320,7 @@ impl QueueWorker {
         self
     }
 
-    pub async fn run(mut self) {
+    pub(crate) async fn run(mut self) {
         if let Err(e) = self.run_inner().await {
             tracing::error!(
                 error = &e as &dyn std::error::Error,
@@ -344,7 +344,7 @@ impl QueueWorker {
     }
 
     #[tracing::instrument(name = "worker.setup_schedules", skip_all)]
-    pub async fn setup_schedules(&mut self) -> Result<(), QueueRunnerError> {
+    pub(crate) async fn setup_schedules(&mut self) -> Result<(), QueueRunnerError> {
         let schedules: Vec<_> = self.schedules.iter().map(|s| s.schedule_name).collect();
 
         // Start a transaction on the existing PgListener connection