Extract the job tracing span logic to a layer

crates/tasks/Cargo.toml

@@ -1,9 +1,9 @@
 [package]
 name = "mas-tasks"
 version = "0.1.0"
-authors = ["Quentin Gliech <quenting@element.io>"]
+authors = ["quentin gliech <quenting@element.io>"]
 edition = "2021"
-license = "Apache-2.0"
+license = "apache-2.0"
 
 [dependencies]
 anyhow = "1.0.70"
@@ -19,7 +19,6 @@ thiserror = "1.0.30"
 tower = "0.4.13"
 tracing = "0.1.37"
 tracing-opentelemetry = "0.18.0"
-opentelemetry = "0.18.0"
 ulid = "1.0.0"
 serde = { version = "1.0.159", features = ["derive"] }
 

crates/tasks/src/database.rs

@@ -68,9 +68,14 @@ pub async fn cleanup_expired_tokens(
     Ok(())
 }
 
-pub(crate) fn register(monitor: Monitor<TokioExecutor>, state: &State) -> Monitor<TokioExecutor> {
+pub(crate) fn register(
+    suffix: &str,
+    monitor: Monitor<TokioExecutor>,
+    state: &State,
+) -> Monitor<TokioExecutor> {
     let schedule = apalis_cron::Schedule::from_str("*/15 * * * * *").unwrap();
-    let worker = WorkerBuilder::new("cleanup-expired-tokens")
+    let worker_name = format!("{job}-{suffix}", job = CleanupExpiredTokensJob::NAME);
+    let worker = WorkerBuilder::new(worker_name)
         .stream(CronStream::new(schedule).to_stream())
         .layer(state.inject())
         .build(job_fn(cleanup_expired_tokens));

crates/tasks/src/email.rs

@@ -17,6 +17,7 @@ use apalis_core::{
     builder::{WorkerBuilder, WorkerFactory},
     context::JobContext,
     executor::TokioExecutor,
+    job::Job,
     job_fn::job_fn,
     monitor::Monitor,
     storage::builder::WithStorage,
@@ -25,84 +26,80 @@ use chrono::Duration;
 use mas_email::{Address, EmailVerificationContext, Mailbox};
 use mas_storage::job::{JobWithSpanContext, VerifyEmailJob};
 use rand::{distributions::Uniform, Rng};
-use tracing::{info, info_span, Instrument};
-use tracing_opentelemetry::OpenTelemetrySpanExt;
+use tracing::info;
 
-use crate::{JobContextExt, State};
+use crate::{layers::TracingLayer, JobContextExt, State};
 
+#[tracing::instrument(
+    name = "job.verify_email",
+    fields(user_email.id = %job.user_email_id()),
+    skip_all,
+    err(Debug),
+)]
 async fn verify_email(
     job: JobWithSpanContext<VerifyEmailJob>,
     ctx: JobContext,
 ) -> Result<(), anyhow::Error> {
-    let span = info_span!(
-        "job.verify_email",
-        job.id = %ctx.id(),
-        job.attempts = ctx.attempts(),
-        user_email.id = %job.user_email_id(),
+    let state = ctx.state();
+    let mut repo = state.repository().await?;
+    let mut rng = state.rng();
+    let mailer = state.mailer();
+    let clock = state.clock();
+
+    // Lookup the user email
+    let user_email = repo
+        .user_email()
+        .lookup(job.user_email_id())
+        .await?
+        .context("User email not found")?;
+
+    // Lookup the user associated with the email
+    let user = repo
+        .user()
+        .lookup(user_email.user_id)
+        .await?
+        .context("User not found")?;
+
+    // Generate a verification code
+    let range = Uniform::<u32>::from(0..1_000_000);
+    let code = rng.sample(range);
+    let code = format!("{code:06}");
+
+    let address: Address = user_email.email.parse()?;
+
+    // Save the verification code in the database
+    let verification = repo
+        .user_email()
+        .add_verification_code(&mut rng, &clock, &user_email, Duration::hours(8), code)
+        .await?;
+
+    // And send the verification email
+    let mailbox = Mailbox::new(Some(user.username.clone()), address);
+
+    let context = EmailVerificationContext::new(user.clone(), verification.clone());
+
+    mailer.send_verification_email(mailbox, &context).await?;
+
+    info!(
+        email.id = %user_email.id,
+        "Verification email sent"
     );
 
-    if let Some(context) = job.span_context() {
-        span.add_link(context);
-    }
+    repo.save().await?;
 
-    async move {
-        let state = ctx.state();
-        let mut repo = state.repository().await?;
-        let mut rng = state.rng();
-        let mailer = state.mailer();
-        let clock = state.clock();
-
-        // Lookup the user email
-        let user_email = repo
-            .user_email()
-            .lookup(job.user_email_id())
-            .await?
-            .context("User email not found")?;
-
-        // Lookup the user associated with the email
-        let user = repo
-            .user()
-            .lookup(user_email.user_id)
-            .await?
-            .context("User not found")?;
-
-        // Generate a verification code
-        let range = Uniform::<u32>::from(0..1_000_000);
-        let code = rng.sample(range);
-        let code = format!("{code:06}");
-
-        let address: Address = user_email.email.parse()?;
-
-        // Save the verification code in the database
-        let verification = repo
-            .user_email()
-            .add_verification_code(&mut rng, &clock, &user_email, Duration::hours(8), code)
-            .await?;
-
-        // And send the verification email
-        let mailbox = Mailbox::new(Some(user.username.clone()), address);
-
-        let context = EmailVerificationContext::new(user.clone(), verification.clone());
-
-        mailer.send_verification_email(mailbox, &context).await?;
-
-        info!(
-            email.id = %user_email.id,
-            "Verification email sent"
-        );
-
-        repo.save().await?;
-
-        Ok(())
-    }
-    .instrument(span)
-    .await
+    Ok(())
 }
 
-pub(crate) fn register(monitor: Monitor<TokioExecutor>, state: &State) -> Monitor<TokioExecutor> {
+pub(crate) fn register(
+    suffix: &str,
+    monitor: Monitor<TokioExecutor>,
+    state: &State,
+) -> Monitor<TokioExecutor> {
     let storage = state.store();
-    let worker = WorkerBuilder::new("verify-email")
+    let worker_name = format!("{job}-{suffix}", job = VerifyEmailJob::NAME);
+    let worker = WorkerBuilder::new(worker_name)
         .layer(state.inject())
+        .layer(TracingLayer::new())
         .with_storage(storage)
         .build(job_fn(verify_email));
     monitor.register(worker)

crates/tasks/src/layers.rs

@@ -0,0 +1,70 @@
+// Copyright 2023 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::task::{Context, Poll};
+
+use apalis_core::{job::Job, request::JobRequest};
+use mas_storage::job::JobWithSpanContext;
+use tower::{Layer, Service};
+use tracing::{info_span, instrument::Instrumented, Instrument};
+use tracing_opentelemetry::OpenTelemetrySpanExt;
+
+pub struct TracingLayer;
+
+impl TracingLayer {
+    pub fn new() -> Self {
+        Self
+    }
+}
+
+impl<S> Layer<S> for TracingLayer {
+    type Service = TracingService<S>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        TracingService { inner }
+    }
+}
+
+pub struct TracingService<S> {
+    inner: S,
+}
+
+impl<J, S> Service<JobRequest<JobWithSpanContext<J>>> for TracingService<S>
+where
+    J: Job,
+    S: Service<JobRequest<JobWithSpanContext<J>>>,
+{
+    type Response = S::Response;
+    type Error = S::Error;
+    type Future = Instrumented<S::Future>;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, req: JobRequest<JobWithSpanContext<J>>) -> Self::Future {
+        let span = info_span!(
+            "job.run",
+            job.id = %req.id(),
+            job.attempts = req.attempts(),
+            job.name = J::NAME,
+        );
+
+        if let Some(context) = req.inner().span_context() {
+            span.add_link(context);
+        }
+
+        self.inner.call(req).instrument(span)
+    }
+}

crates/tasks/src/lib.rs

@@ -27,6 +27,7 @@ use tracing::debug;
 
 mod database;
 mod email;
+mod layers;
 
 #[derive(Clone)]
 struct State {
@@ -95,11 +96,11 @@ impl JobContextExt for apalis_core::context::JobContext {
 }
 
 #[must_use]
-pub fn init(pool: &Pool<Postgres>, mailer: &Mailer) -> Monitor<TokioExecutor> {
+pub fn init(name: &str, pool: &Pool<Postgres>, mailer: &Mailer) -> Monitor<TokioExecutor> {
     let state = State::new(pool.clone(), SystemClock::default(), mailer.clone());
     let monitor = Monitor::new();
-    let monitor = self::database::register(monitor, &state);
-    let monitor = self::email::register(monitor, &state);
+    let monitor = self::database::register(name, monitor, &state);
+    let monitor = self::email::register(name, monitor, &state);
     debug!(?monitor, "workers registered");
     monitor
 }