diff --git a/Cargo.lock b/Cargo.lock index 6df0ff2..25bd04b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -64,6 +64,17 @@ dependencies = [ "webpki-roots", ] +[[package]] +name = "atty" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9b39be18770d11421cdb1b9947a45dd3f37e93092cbf377614828a319d5fee8" +dependencies = [ + "hermit-abi", + "libc", + "winapi", +] + [[package]] name = "autocfg" version = "1.1.0" @@ -209,6 +220,7 @@ dependencies = [ "tokio", "tracing", "tracing-subscriber", + "tracing-tree", ] [[package]] @@ -1743,6 +1755,19 @@ dependencies = [ "tracing-log", ] +[[package]] +name = "tracing-tree" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d07e90b329c621ade432823988574e820212648aa40e7a2497777d58de0fb453" +dependencies = [ + "ansi_term", + "atty", + "tracing-core", + "tracing-log", + "tracing-subscriber", +] + [[package]] name = "trust-dns-proto" version = "0.21.2" diff --git a/Cargo.toml b/Cargo.toml index 704074f..c1ae4fa 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,3 +15,4 @@ tracing-subscriber = { version = "0.3.11", features = ["env-filter"] } serde = { version = "1.0.137", features = ["derive"] } serde_json = "1.0.81" poise = "0.2.2" +tracing-tree = "0.2.1" diff --git a/README.md b/README.md index 83fa392..aab6825 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,7 @@ MONGO_INITDB_ROOT_USERNAME=root MONGO_INITDB_ROOT_PASSWORD=uwu DEV= # SET_GLOBAL= +PRETTY= ``` run mongodb diff --git a/src/handler.rs b/src/handler.rs index 264aa7a..4fafa6b 100644 --- a/src/handler.rs +++ b/src/handler.rs @@ -75,15 +75,7 @@ impl Handler { pub mod lawsuit { use super::*; - #[poise::command( - slash_command, - subcommands( - "create", - "set_category", - "close", - "clear" - ) - )] + #[poise::command(slash_command, subcommands("create", "set_category", "close", "clear"))] pub async fn lawsuit(_: crate::Context<'_>) -> Result<()> { unreachable!() } @@ -140,6 +132,7 @@ pub mod lawsuit { lawsuit_clear_impl(ctx).await.wrap_err("lawsuit_clear") } + #[tracing::instrument(skip(ctx))] async fn lawsuit_create_impl( ctx: crate::Context<'_>, plaintiff: User, @@ -180,6 +173,7 @@ pub mod lawsuit { Ok(()) } + #[tracing::instrument(skip(ctx))] async fn lawsuit_set_category_impl(ctx: crate::Context<'_>, category: Channel) -> Result<()> { let guild_id = ctx.guild_id().wrap_err("guild_id not found")?; @@ -205,6 +199,7 @@ pub mod lawsuit { Ok(()) } + #[tracing::instrument(skip(ctx))] async fn lawsuit_close_impl(ctx: crate::Context<'_>, verdict: String) -> Result<()> { let guild_id = ctx.guild_id().wrap_err("guild_id not found")?; @@ -273,6 +268,7 @@ pub mod lawsuit { Ok(()) } + #[tracing::instrument(skip(ctx))] async fn lawsuit_clear_impl(ctx: crate::Context<'_>) -> Result<()> { let guild_id = ctx.guild_id().wrap_err("guild_id not found")?; @@ -285,10 +281,7 @@ pub mod lawsuit { pub mod prison { use super::*; - #[poise::command( - slash_command, - subcommands("set_role", "arrest", "release") - )] + #[poise::command(slash_command, subcommands("set_role", "arrest", "release"))] pub async fn prison(_: crate::Context<'_>) -> Result<()> { unreachable!() } @@ -326,6 +319,7 @@ pub mod prison { .wrap_err("prison_release") } + #[tracing::instrument(skip(ctx))] async fn prison_set_role_impl(ctx: crate::Context<'_>, role: Role) -> Result<()> { ctx.data() .mongo @@ -340,6 +334,7 @@ pub mod prison { Ok(()) } + #[tracing::instrument(skip(ctx))] async fn prison_arrest_impl(ctx: crate::Context<'_>, user: User) -> Result<()> { let mongo_client = &ctx.data().mongo; let guild_id = ctx.guild_id().wrap_err("guild_id not found")?; @@ -371,6 +366,7 @@ pub mod prison { Ok(()) } + #[tracing::instrument(skip(ctx))] async fn prison_release_impl(ctx: crate::Context<'_>, user: User) -> Result<()> { let mongo_client = &ctx.data().mongo; let guild_id = ctx.guild_id().wrap_err("guild_id not found")?; diff --git a/src/lawsuit.rs b/src/lawsuit.rs index ce3369e..d7035db 100644 --- a/src/lawsuit.rs +++ b/src/lawsuit.rs @@ -2,8 +2,10 @@ use std::sync::Arc; use color_eyre::Result; use mongodb::bson::{doc, Uuid}; -use poise::{serenity::model::prelude::*, serenity_prelude::Http}; -use poise::serenity_prelude::CreateMessage; +use poise::{ + serenity::model::prelude::*, + serenity_prelude::{CreateMessage, Http}, +}; use serde::{Deserialize, Serialize}; use tracing::{error, info}; diff --git a/src/main.rs b/src/main.rs index a87066f..d9a11ff 100644 --- a/src/main.rs +++ b/src/main.rs @@ -12,7 +12,7 @@ use poise::{ serenity_prelude::{Activity, GatewayIntents, GuildId}, }; use tracing::{error, info}; -use tracing_subscriber::EnvFilter; +use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Registry}; use crate::{handler::Handler, model::Mongo}; @@ -24,9 +24,9 @@ async fn main() -> Result<()> { let _ = dotenv::dotenv(); - tracing_subscriber::fmt() - .with_env_filter(EnvFilter::from_default_env()) - .init(); + let pretty = env::var("PRETTY").is_ok(); + + setup_tracing(pretty); info!("Starting up..."); @@ -106,7 +106,11 @@ async fn main() -> Result<()> { }) }) .options(poise::FrameworkOptions { - commands: vec![handler::lawsuit::lawsuit(), handler::prison::prison(), hello()], + commands: vec![ + handler::lawsuit::lawsuit(), + handler::prison::prison(), + hello(), + ], on_error: |err| Box::pin(async { handler::error_handler(err).await }), listener: |ctx, event, ctx2, data| { Box::pin(async move { handler::listener(ctx, event, ctx2, data).await }) @@ -148,3 +152,23 @@ async fn hello(ctx: Context<'_>) -> Result<()> { ctx.say("hoi!").await?; Ok(()) } + +fn setup_tracing(pretty: bool) { + let registry = Registry::default().with(EnvFilter::from_default_env()); + + if pretty { + let tree_layer = tracing_tree::HierarchicalLayer::new(2) + .with_targets(true) + .with_bracketed_fields(true); + + registry.with(tree_layer).init(); + } else { + let fmt_layer = tracing_subscriber::fmt::layer() + .with_level(true) + .with_timer(tracing_subscriber::fmt::time::time()) + .with_ansi(true) + .with_thread_names(true); + + registry.with(fmt_layer).init(); + }; +} diff --git a/src/model.rs b/src/model.rs index 0a7d1d1..d3a9585 100644 --- a/src/model.rs +++ b/src/model.rs @@ -137,6 +137,7 @@ impl Mongo { Ok(Self { db }) } + #[tracing::instrument(skip(self))] pub async fn find_or_insert_state(&self, guild_id: SnowflakeId) -> Result { let coll = self.state_coll(); let state = coll @@ -155,6 +156,7 @@ impl Mongo { Ok(state) } + #[tracing::instrument(skip(self))] pub async fn new_state(&self, guild_id: SnowflakeId) -> Result { let state = State { guild_id, @@ -171,6 +173,7 @@ impl Mongo { Ok(state) } + #[tracing::instrument(skip(self))] pub async fn set_court_category( &self, guild_id: SnowflakeId, @@ -188,6 +191,7 @@ impl Mongo { Ok(()) } + #[tracing::instrument(skip(self))] pub async fn set_prison_role( &self, guild_id: SnowflakeId, @@ -205,6 +209,7 @@ impl Mongo { Ok(()) } + #[tracing::instrument(skip(self))] pub async fn add_court_room(&self, guild_id: SnowflakeId, room: &CourtRoom) -> Result<()> { let _ = self.find_or_insert_state(guild_id).await?; let coll = self.state_coll(); @@ -218,6 +223,7 @@ impl Mongo { Ok(()) } + #[tracing::instrument(skip(self))] pub async fn add_lawsuit(&self, guild_id: SnowflakeId, lawsuit: &Lawsuit) -> Result<()> { let _ = self.find_or_insert_state(guild_id).await?; let coll = self.state_coll(); @@ -233,6 +239,7 @@ impl Mongo { Ok(()) } + #[tracing::instrument(skip(self, value))] pub async fn set_court_room( &self, guild_id: SnowflakeId, @@ -252,6 +259,7 @@ impl Mongo { Ok(()) } + #[tracing::instrument(skip(self, value))] pub async fn set_lawsuit( &self, guild_id: SnowflakeId, @@ -271,6 +279,7 @@ impl Mongo { Ok(()) } + #[tracing::instrument(skip(self))] pub async fn delete_guild(&self, guild_id: SnowflakeId) -> Result<()> { let coll = self.state_coll(); @@ -280,6 +289,7 @@ impl Mongo { Ok(()) } + #[tracing::instrument(skip(self))] pub async fn add_to_prison(&self, guild_id: SnowflakeId, user_id: SnowflakeId) -> Result<()> { let coll = self.prison_coll(); @@ -298,6 +308,7 @@ impl Mongo { Ok(()) } + #[tracing::instrument(skip(self))] pub async fn remove_from_prison( &self, guild_id: SnowflakeId, @@ -312,6 +323,7 @@ impl Mongo { Ok(()) } + #[tracing::instrument(skip(self))] pub async fn find_prison_entry( &self, guild_id: SnowflakeId,