Allow tracing filters to be changed at runtime

ReloadHandle is taken from conduwuit commit
8a5599adf9eafe9111f3d1597f8fb333b8b76849, authored by Benjamin.

Co-authored-by: Benjamin Lee <benjamin@computer.surgery>

src/service/admin.rs

@@ -1,6 +1,6 @@
 use std::{collections::BTreeMap, fmt::Write, sync::Arc, time::Instant};
 
-use clap::Parser;
+use clap::{Parser, ValueEnum};
 use regex::Regex;
 use ruma::{
     api::appservice::Registration,
@@ -196,6 +196,12 @@ enum AdminCommand {
     // Allowed because the doc comment gets parsed by our code later
     #[allow(clippy::doc_markdown)]
     VerifyJson,
+
+    /// Dynamically change a tracing backend's filter string
+    SetTracingFilter {
+        backend: TracingBackend,
+        filter: String,
+    },
 }
 
 #[derive(Debug)]
@@ -209,6 +215,13 @@ pub(crate) struct Service {
     receiver: Mutex<mpsc::UnboundedReceiver<AdminRoomEvent>>,
 }
 
+#[derive(Debug, Clone, ValueEnum)]
+enum TracingBackend {
+    Log,
+    Flame,
+    Traces,
+}
+
 impl Service {
     pub(crate) fn build() -> Arc<Self> {
         let (sender, receiver) = mpsc::unbounded_channel();
@@ -1081,6 +1094,39 @@ impl Service {
                     )
                 }
             }
+            AdminCommand::SetTracingFilter {
+                backend,
+                filter,
+            } => {
+                let handles = &services().globals.reload_handles;
+                let handle = match backend {
+                    TracingBackend::Log => &handles.log,
+                    TracingBackend::Flame => &handles.flame,
+                    TracingBackend::Traces => &handles.traces,
+                };
+                let Some(handle) = handle else {
+                    return Ok(RoomMessageEventContent::text_plain(
+                        "Backend is disabled",
+                    ));
+                };
+                let filter = match filter.parse() {
+                    Ok(filter) => filter,
+                    Err(e) => {
+                        return Ok(RoomMessageEventContent::text_plain(
+                            format!("Invalid filter string: {e}"),
+                        ));
+                    }
+                };
+                if let Err(e) = handle.reload(filter) {
+                    return Ok(RoomMessageEventContent::text_plain(format!(
+                        "Failed to reload filter: {e}"
+                    )));
+                };
+
+                return Ok(RoomMessageEventContent::text_plain(
+                    "Filter reloaded",
+                ));
+            }
         };
 
         Ok(reply_message_content)

src/service/globals.rs

@@ -32,7 +32,10 @@ use tokio::sync::{broadcast, Mutex, RwLock, Semaphore};
 use tracing::{error, Instrument};
 use trust_dns_resolver::TokioAsyncResolver;
 
-use crate::{api::server_server::FedDest, services, Config, Error, Result};
+use crate::{
+    api::server_server::FedDest, observability::FilterReloadHandles, services,
+    Config, Error, Result,
+};
 
 type WellKnownMap = HashMap<OwnedServerName, (FedDest, String)>;
 type TlsNameMap = HashMap<String, (Vec<IpAddr>, u16)>;
@@ -41,6 +44,7 @@ type RateLimitState = (Instant, u32);
 
 pub(crate) struct Service {
     pub(crate) db: &'static dyn Data,
+    pub(crate) reload_handles: FilterReloadHandles,
 
     // actual_destination, host
     pub(crate) actual_destination_cache: Arc<RwLock<WellKnownMap>>,
@@ -173,7 +177,11 @@ impl Resolve for Resolver {
 
 impl Service {
     #[tracing::instrument(skip_all)]
-    pub(crate) fn load(db: &'static dyn Data, config: Config) -> Result<Self> {
+    pub(crate) fn load(
+        db: &'static dyn Data,
+        config: Config,
+        reload_handles: FilterReloadHandles,
+    ) -> Result<Self> {
         let keypair = db.load_keypair();
 
         let keypair = match keypair {
@@ -227,6 +235,7 @@ impl Service {
         let mut s = Self {
             db,
             config,
+            reload_handles,
             keypair: Arc::new(keypair),
             dns_resolver: TokioAsyncResolver::tokio_from_system_conf()
                 .map_err(|e| {