mas-tasks: refactor worker building behind a macro

crates/tasks/src/email.rs

@@ -13,25 +13,14 @@
 // limitations under the License.
 
 use anyhow::Context;
-use apalis_core::{
-    builder::{WorkerBuilder, WorkerFactoryFn},
-    context::JobContext,
-    executor::TokioExecutor,
-    job::Job,
-    monitor::Monitor,
-    storage::builder::WithStorage,
-};
+use apalis_core::{context::JobContext, executor::TokioExecutor, monitor::Monitor};
 use chrono::Duration;
 use mas_email::{Address, EmailVerificationContext, Mailbox};
 use mas_storage::job::{JobWithSpanContext, VerifyEmailJob};
 use rand::{distributions::Uniform, Rng};
 use tracing::info;
 
-use crate::{
-    storage::PostgresStorageFactory,
-    utils::{metrics_layer, trace_layer},
-    JobContextExt, State,
-};
+use crate::{storage::PostgresStorageFactory, JobContextExt, State};
 
 #[tracing::instrument(
     name = "job.verify_email",
@@ -99,16 +88,8 @@ pub(crate) fn register(
     state: &State,
     storage_factory: &PostgresStorageFactory,
 ) -> Monitor<TokioExecutor> {
-    let storage = storage_factory.build();
-    let worker_name = format!("{job}-{suffix}", job = VerifyEmailJob::NAME);
-    let worker = WorkerBuilder::new(worker_name)
-        .layer(state.inject())
-        .layer(trace_layer())
-        .layer(metrics_layer())
-        .with_storage_config(storage, |c| {
-            c.fetch_interval(std::time::Duration::from_secs(1))
-        })
-        .build_fn(verify_email);
+    let verify_email_worker =
+        crate::build!(VerifyEmailJob => verify_email, suffix, state, storage_factory);
 
-    monitor.register(worker)
+    monitor.register(verify_email_worker)
 }

crates/tasks/src/lib.rs

@@ -107,6 +107,32 @@ impl JobContextExt for apalis_core::context::JobContext {
     }
 }
 
+/// Helper macro to build a storage-backed worker.
+macro_rules! build {
+    ($job:ty => $fn:ident, $suffix:expr, $state:expr, $factory:expr) => {{
+        let storage = $factory.build();
+        let worker_name = format!(
+            "{job}-{suffix}",
+            job = <$job as ::apalis_core::job::Job>::NAME,
+            suffix = $suffix
+        );
+
+        let builder = ::apalis_core::builder::WorkerBuilder::new(worker_name)
+            .layer($state.inject())
+            .layer(crate::utils::trace_layer())
+            .layer(crate::utils::metrics_layer());
+
+        let builder = ::apalis_core::storage::builder::WithStorage::with_storage_config(
+            builder,
+            storage,
+            |c| c.fetch_interval(std::time::Duration::from_secs(1)),
+        );
+        ::apalis_core::builder::WorkerFactory::build(builder, ::apalis_core::job_fn::job_fn($fn))
+    }};
+}
+
+pub(crate) use build;
+
 /// Initialise the workers.
 ///
 /// # Errors

crates/tasks/src/matrix.rs

@@ -12,17 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::time::Duration;
-
 use anyhow::Context;
-use apalis_core::{
-    builder::{WorkerBuilder, WorkerFactoryFn},
-    context::JobContext,
-    executor::TokioExecutor,
-    job::Job,
-    monitor::Monitor,
-    storage::builder::WithStorage,
-};
+use apalis_core::{context::JobContext, executor::TokioExecutor, monitor::Monitor};
 use mas_matrix::ProvisionRequest;
 use mas_storage::{
     job::{DeleteDeviceJob, JobWithSpanContext, ProvisionDeviceJob, ProvisionUserJob},
@@ -31,11 +22,7 @@ use mas_storage::{
 };
 use tracing::info;
 
-use crate::{
-    storage::PostgresStorageFactory,
-    utils::{metrics_layer, trace_layer},
-    JobContextExt, State,
-};
+use crate::{storage::PostgresStorageFactory, JobContextExt, State};
 
 /// Job to provision a user on the Matrix homeserver.
 /// This works by doing a PUT request to the /_synapse/admin/v2/users/{user_id}
@@ -163,32 +150,12 @@ pub(crate) fn register(
     state: &State,
     storage_factory: &PostgresStorageFactory,
 ) -> Monitor<TokioExecutor> {
-    let storage = storage_factory.build();
-    let worker_name = format!("{job}-{suffix}", job = ProvisionUserJob::NAME);
-    let provision_user_worker = WorkerBuilder::new(worker_name)
-        .layer(state.inject())
-        .layer(trace_layer())
-        .layer(metrics_layer())
-        .with_storage_config(storage, |c| c.fetch_interval(Duration::from_secs(1)))
-        .build_fn(provision_user);
-
-    let storage = storage_factory.build();
-    let worker_name = format!("{job}-{suffix}", job = ProvisionDeviceJob::NAME);
-    let provision_device_worker = WorkerBuilder::new(worker_name)
-        .layer(state.inject())
-        .layer(trace_layer())
-        .layer(metrics_layer())
-        .with_storage_config(storage, |c| c.fetch_interval(Duration::from_secs(1)))
-        .build_fn(provision_device);
-
-    let storage = storage_factory.build();
-    let worker_name = format!("{job}-{suffix}", job = DeleteDeviceJob::NAME);
-    let delete_device_worker = WorkerBuilder::new(worker_name)
-        .layer(state.inject())
-        .layer(trace_layer())
-        .layer(metrics_layer())
-        .with_storage_config(storage, |c| c.fetch_interval(Duration::from_secs(1)))
-        .build_fn(delete_device);
+    let provision_user_worker =
+        crate::build!(ProvisionUserJob => provision_user, suffix, state, storage_factory);
+    let provision_device_worker =
+        crate::build!(ProvisionDeviceJob => provision_device, suffix, state, storage_factory);
+    let delete_device_worker =
+        crate::build!(DeleteDeviceJob => delete_device, suffix, state, storage_factory);
 
     monitor
         .register(provision_user_worker)

crates/tasks/src/user.rs

@@ -12,29 +12,16 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::time::Duration;
-
 use anyhow::Context;
-use apalis_core::{
-    builder::{WorkerBuilder, WorkerFactoryFn},
-    context::JobContext,
-    executor::TokioExecutor,
-    job::Job,
-    monitor::Monitor,
-    storage::builder::WithStorage,
-};
+use apalis_core::{context::JobContext, executor::TokioExecutor, monitor::Monitor};
 use mas_storage::{
-    job::{DeactivateUserJob, DeleteDeviceJob, JobWithSpanContext},
+    job::{DeactivateUserJob, JobWithSpanContext},
     user::UserRepository,
     RepositoryAccess,
 };
 use tracing::info;
 
-use crate::{
-    storage::PostgresStorageFactory,
-    utils::{metrics_layer, trace_layer},
-    JobContextExt, State,
-};
+use crate::{storage::PostgresStorageFactory, JobContextExt, State};
 
 /// Job to deactivate a user, both locally and on the Matrix homeserver.
 #[tracing::instrument(
@@ -83,14 +70,8 @@ pub(crate) fn register(
     state: &State,
     storage_factory: &PostgresStorageFactory,
 ) -> Monitor<TokioExecutor> {
-    let storage = storage_factory.build();
-    let worker_name = format!("{job}-{suffix}", job = DeleteDeviceJob::NAME);
-    let deactivate_user_worker = WorkerBuilder::new(worker_name)
-        .layer(state.inject())
-        .layer(trace_layer())
-        .layer(metrics_layer())
-        .with_storage_config(storage, |c| c.fetch_interval(Duration::from_secs(1)))
-        .build_fn(deactivate_user);
+    let deactivate_user_worker =
+        crate::build!(DeactivateUserJob => deactivate_user, suffix, state, storage_factory);
 
     monitor.register(deactivate_user_worker)
 }