Extract the job tracing span logic to a layer

Cargo.lock

@@ -124,7 +124,7 @@ checksum = "7de8ce5e0f9f8d88245311066a578d72b7af3e7088f32783804676302df237e4"
 [[package]]
 name = "apalis-core"
 version = "0.4.0-alpha.4"
-source = "git+https://github.com/geofmureithi/apalis.git?rev=ead6f840b92a3590a8bf46398eaf65ed55aa92dc#ead6f840b92a3590a8bf46398eaf65ed55aa92dc"
+source = "git+https://github.com/geofmureithi/apalis.git?rev=1cc600571c0dc4f5bd6cbb90fa9eef7f34888ba8#1cc600571c0dc4f5bd6cbb90fa9eef7f34888ba8"
 dependencies = [
  "async-stream",
  "async-trait",
@@ -147,7 +147,7 @@ dependencies = [
 [[package]]
 name = "apalis-cron"
 version = "0.4.0-alpha.4"
-source = "git+https://github.com/geofmureithi/apalis.git?rev=ead6f840b92a3590a8bf46398eaf65ed55aa92dc#ead6f840b92a3590a8bf46398eaf65ed55aa92dc"
+source = "git+https://github.com/geofmureithi/apalis.git?rev=1cc600571c0dc4f5bd6cbb90fa9eef7f34888ba8#1cc600571c0dc4f5bd6cbb90fa9eef7f34888ba8"
 dependencies = [
  "apalis-core",
  "async-stream",
@@ -161,7 +161,7 @@ dependencies = [
 [[package]]
 name = "apalis-sql"
 version = "0.4.0-alpha.4"
-source = "git+https://github.com/geofmureithi/apalis.git?rev=ead6f840b92a3590a8bf46398eaf65ed55aa92dc#ead6f840b92a3590a8bf46398eaf65ed55aa92dc"
+source = "git+https://github.com/geofmureithi/apalis.git?rev=1cc600571c0dc4f5bd6cbb90fa9eef7f34888ba8#1cc600571c0dc4f5bd6cbb90fa9eef7f34888ba8"
 dependencies = [
  "apalis-core",
  "async-stream",
@@ -3506,7 +3506,6 @@ dependencies = [
  "mas-email",
  "mas-storage",
  "mas-storage-pg",
- "opentelemetry",
  "rand 0.8.5",
  "rand_chacha 0.3.1",
  "serde",

Cargo.toml

@@ -15,12 +15,12 @@ rev = "0b9295c2db2114cd87aa19abcc1fc00c16b272db"
 
 [patch.crates-io.apalis-core]
 git = "https://github.com/geofmureithi/apalis.git"
-rev = "ead6f840b92a3590a8bf46398eaf65ed55aa92dc"
+rev = "1cc600571c0dc4f5bd6cbb90fa9eef7f34888ba8"
 
 [patch.crates-io.apalis-sql]
 git = "https://github.com/geofmureithi/apalis.git"
-rev = "ead6f840b92a3590a8bf46398eaf65ed55aa92dc"
+rev = "1cc600571c0dc4f5bd6cbb90fa9eef7f34888ba8"
 
 [patch.crates-io.apalis-cron]
 git = "https://github.com/geofmureithi/apalis.git"
-rev = "ead6f840b92a3590a8bf46398eaf65ed55aa92dc"
+rev = "1cc600571c0dc4f5bd6cbb90fa9eef7f34888ba8"

crates/cli/src/commands/server.rs

@@ -22,6 +22,10 @@ use mas_handlers::{AppState, HttpClientFactory, MatrixHomeserver};
 use mas_listener::{server::Server, shutdown::ShutdownStream};
 use mas_router::UrlBuilder;
 use mas_storage_pg::MIGRATOR;
+use rand::{
+    distributions::{Alphanumeric, DistString},
+    thread_rng,
+};
 use tokio::signal::unix::SignalKind;
 use tracing::{info, info_span, warn, Instrument};
 
@@ -52,7 +56,7 @@ impl Options {
         let config: RootConfig = root.load_config()?;
 
         // Connect to the database
-        info!("Conntecting to the database");
+        info!("Connecting to the database");
         let pool = database_from_config(&config.database).await?;
 
         if self.migrate {
@@ -87,8 +91,12 @@ impl Options {
             let mailer = mailer_from_config(&config.email, &templates).await?;
             mailer.test_connection().await?;
 
-            info!("Starting task worker");
-            let monitor = mas_tasks::init(&pool, &mailer);
+            #[allow(clippy::disallowed_methods)]
+            let mut rng = thread_rng();
+            let worker_name = Alphanumeric.sample_string(&mut rng, 10);
+
+            info!(worker_name, "Starting task worker");
+            let monitor = mas_tasks::init(&worker_name, &pool, &mailer);
             // TODO: grab the handle
             tokio::spawn(monitor.run());
         }

crates/cli/src/commands/worker.rs

@@ -15,7 +15,11 @@
 use clap::Parser;
 use mas_config::RootConfig;
 use mas_router::UrlBuilder;
-use tracing::{info_span, log::info};
+use rand::{
+    distributions::{Alphanumeric, DistString},
+    thread_rng,
+};
+use tracing::{info, info_span};
 
 use crate::util::{database_from_config, mailer_from_config, templates_from_config};
 
@@ -28,7 +32,7 @@ impl Options {
         let config: RootConfig = root.load_config()?;
 
         // Connect to the database
-        info!("Conntecting to the database");
+        info!("Connecting to the database");
         let pool = database_from_config(&config.database).await?;
 
         let url_builder = UrlBuilder::new(config.http.public_base.clone());
@@ -40,8 +44,12 @@ impl Options {
         mailer.test_connection().await?;
         drop(config);
 
-        info!("Starting task scheduler");
-        let monitor = mas_tasks::init(&pool, &mailer);
+        #[allow(clippy::disallowed_methods)]
+        let mut rng = thread_rng();
+        let worker_name = Alphanumeric.sample_string(&mut rng, 10);
+
+        info!(worker_name, "Starting task scheduler");
+        let monitor = mas_tasks::init(&worker_name, &pool, &mailer);
 
         span.exit();
 

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,27 +26,20 @@ 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(),
-    );
-
-    if let Some(context) = job.span_context() {
-        span.add_link(context);
-    }
-
-    async move {
     let state = ctx.state();
     let mut repo = state.repository().await?;
     let mut rng = state.rng();
@@ -95,14 +89,17 @@ async fn verify_email(
 
     Ok(())
 }
-    .instrument(span)
-    .await
-}
 
-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
 }