From f9a6c2cefe5545a2584275d45d40f3b97f9d9302 Mon Sep 17 00:00:00 2001 From: alyssa Date: Sat, 17 May 2025 16:19:30 +0000 Subject: [PATCH] chore(rust): tidier libpk::main macro impl --- Cargo.lock | 22 +++++----- crates/api/src/main.rs | 4 +- crates/avatars/src/cleanup.rs | 4 +- crates/avatars/src/main.rs | 4 +- crates/dispatch/Cargo.toml | 1 + crates/dispatch/src/main.rs | 11 +---- crates/gateway/src/main.rs | 4 +- crates/gdpr_worker/src/main.rs | 4 +- crates/libpk/Cargo.toml | 1 + crates/libpk/src/lib.rs | 29 ++----------- crates/{model_macros => macros}/Cargo.toml | 2 +- crates/macros/src/entrypoint.rs | 41 +++++++++++++++++++ crates/macros/src/lib.rs | 14 +++++++ .../src/lib.rs => macros/src/model.rs} | 3 +- crates/models/Cargo.toml | 2 +- crates/models/src/system.rs | 2 +- crates/models/src/system_config.rs | 2 +- crates/scheduled_tasks/src/main.rs | 4 +- 18 files changed, 90 insertions(+), 64 deletions(-) rename crates/{model_macros => macros}/Cargo.toml (85%) create mode 100644 crates/macros/src/entrypoint.rs create mode 100644 crates/macros/src/lib.rs rename crates/{model_macros/src/lib.rs => macros/src/model.rs} (99%) diff --git a/Cargo.lock b/Cargo.lock index 0b9666c6..2228d94b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -861,6 +861,7 @@ dependencies = [ "anyhow", "axum 0.7.9", "hickory-client", + "libpk", "reqwest 0.12.15", "serde", "serde_json", @@ -2011,6 +2012,7 @@ dependencies = [ "lazy_static", "metrics", "metrics-exporter-prometheus", + "pk_macros", "sentry", "sentry-tracing", "serde", @@ -2213,15 +2215,6 @@ dependencies = [ "windows-sys 0.52.0", ] -[[package]] -name = "model_macros" -version = "0.1.0" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - [[package]] name = "nibble_vec" version = "0.1.0" @@ -2518,6 +2511,15 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" +[[package]] +name = "pk_macros" +version = "0.1.0" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "pkcs1" version = "0.7.5" @@ -2550,7 +2552,7 @@ name = "pluralkit_models" version = "0.1.0" dependencies = [ "chrono", - "model_macros", + "pk_macros", "sea-query", "serde", "serde_json", diff --git a/crates/api/src/main.rs b/crates/api/src/main.rs index 24ed58c5..590064eb 100644 --- a/crates/api/src/main.rs +++ b/crates/api/src/main.rs @@ -130,8 +130,8 @@ fn router(ctx: ApiContext) -> Router { .route("/", get(|| async { axum::response::Redirect::to("https://pluralkit.me/api") })) } -libpk::main!("api"); -async fn real_main() -> anyhow::Result<()> { +#[libpk::main] +async fn main() -> anyhow::Result<()> { let db = libpk::db::init_data_db().await?; let redis = libpk::db::init_redis().await?; diff --git a/crates/avatars/src/cleanup.rs b/crates/avatars/src/cleanup.rs index 9b5ed249..cbf5ef1e 100644 --- a/crates/avatars/src/cleanup.rs +++ b/crates/avatars/src/cleanup.rs @@ -4,8 +4,8 @@ use sqlx::prelude::FromRow; use std::{sync::Arc, time::Duration}; use tracing::{error, info}; -libpk::main!("avatar_cleanup"); -async fn real_main() -> anyhow::Result<()> { +#[libpk::main] +async fn main() -> anyhow::Result<()> { let config = libpk::config .avatars .as_ref() diff --git a/crates/avatars/src/main.rs b/crates/avatars/src/main.rs index b551cc41..863619d6 100644 --- a/crates/avatars/src/main.rs +++ b/crates/avatars/src/main.rs @@ -170,8 +170,8 @@ pub struct AppState { config: Arc, } -libpk::main!("avatars"); -async fn real_main() -> anyhow::Result<()> { +#[libpk::main] +async fn main() -> anyhow::Result<()> { let config = libpk::config .avatars .as_ref() diff --git a/crates/dispatch/Cargo.toml b/crates/dispatch/Cargo.toml index 81e35811..c76856d1 100644 --- a/crates/dispatch/Cargo.toml +++ b/crates/dispatch/Cargo.toml @@ -6,6 +6,7 @@ edition = "2021" [dependencies] anyhow = { workspace = true } axum = { workspace = true } +libpk = { path = "../libpk" } reqwest = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } diff --git a/crates/dispatch/src/main.rs b/crates/dispatch/src/main.rs index 5c355d71..6570cf19 100644 --- a/crates/dispatch/src/main.rs +++ b/crates/dispatch/src/main.rs @@ -19,17 +19,8 @@ use axum::{extract::State, http::Uri, routing::post, Json, Router}; mod logger; -// this package does not currently use libpk - -#[tokio::main] +#[libpk::main] async fn main() -> anyhow::Result<()> { - tracing_subscriber::fmt() - .json() - .with_env_filter(EnvFilter::from_default_env()) - .init(); - - info!("hello world"); - let address = std::env::var("DNS_UPSTREAM").unwrap().parse().unwrap(); let stream = UdpClientStream::::with_timeout(address, Duration::from_secs(3)); let (client, bg) = AsyncClient::connect(stream).await?; diff --git a/crates/gateway/src/main.rs b/crates/gateway/src/main.rs index 5cf08f46..ed9ccca1 100644 --- a/crates/gateway/src/main.rs +++ b/crates/gateway/src/main.rs @@ -25,8 +25,8 @@ mod logger; const RUNTIME_CONFIG_KEY_EVENT_TARGET: &'static str = "event_target"; -libpk::main!("gateway"); -async fn real_main() -> anyhow::Result<()> { +#[libpk::main] +async fn main() -> anyhow::Result<()> { let redis = libpk::db::init_redis().await?; let runtime_config = Arc::new( diff --git a/crates/gdpr_worker/src/main.rs b/crates/gdpr_worker/src/main.rs index e801b9fe..b40557c0 100644 --- a/crates/gdpr_worker/src/main.rs +++ b/crates/gdpr_worker/src/main.rs @@ -11,8 +11,8 @@ use twilight_model::id::{ // create table messages_gdpr_jobs (mid bigint not null references messages(mid) on delete cascade, channel bigint not null); -libpk::main!("messages_gdpr_worker"); -async fn real_main() -> anyhow::Result<()> { +#[libpk::main] +async fn main() -> anyhow::Result<()> { let db = libpk::db::init_messages_db().await?; let mut client_builder = twilight_http::Client::builder() diff --git a/crates/libpk/Cargo.toml b/crates/libpk/Cargo.toml index 40d430c8..30d77ae0 100644 --- a/crates/libpk/Cargo.toml +++ b/crates/libpk/Cargo.toml @@ -8,6 +8,7 @@ anyhow = { workspace = true } fred = { workspace = true } lazy_static = { workspace = true } metrics = { workspace = true } +pk_macros = { path = "../macros" } sentry = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } diff --git a/crates/libpk/src/lib.rs b/crates/libpk/src/lib.rs index 8de69d92..55031bf3 100644 --- a/crates/libpk/src/lib.rs +++ b/crates/libpk/src/lib.rs @@ -14,7 +14,9 @@ pub mod state; pub mod _config; pub use crate::_config::CONFIG as config; -// functions in this file are only used by the main function below +// functions in this file are only used by the main function in macros/entrypoint.rs + +pub use pk_macros::main; pub fn init_logging(component: &str) { let sentry_layer = @@ -68,28 +70,3 @@ pub fn init_sentry() -> sentry::ClientInitGuard { ..Default::default() }) } - -#[macro_export] -macro_rules! main { - ($component:expr) => { - fn main() -> anyhow::Result<()> { - let _sentry_guard = libpk::init_sentry(); - // we might also be able to use env!("CARGO_CRATE_NAME") here - libpk::init_logging($component); - tokio::runtime::Builder::new_multi_thread() - .enable_all() - .build() - .unwrap() - .block_on(async { - if let Err(error) = libpk::init_metrics() { - tracing::error!(?error, "failed to init metrics collector"); - }; - tracing::info!("hello world"); - if let Err(error) = real_main().await { - tracing::error!(?error, "failed to run service"); - }; - }); - Ok(()) - } - }; -} diff --git a/crates/model_macros/Cargo.toml b/crates/macros/Cargo.toml similarity index 85% rename from crates/model_macros/Cargo.toml rename to crates/macros/Cargo.toml index d2d0e009..8090798f 100644 --- a/crates/model_macros/Cargo.toml +++ b/crates/macros/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "model_macros" +name = "pk_macros" version = "0.1.0" edition = "2021" diff --git a/crates/macros/src/entrypoint.rs b/crates/macros/src/entrypoint.rs new file mode 100644 index 00000000..1e10012a --- /dev/null +++ b/crates/macros/src/entrypoint.rs @@ -0,0 +1,41 @@ +use proc_macro::{Delimiter, TokenTree}; +use quote::quote; + +pub fn macro_impl( + _args: proc_macro::TokenStream, + input: proc_macro::TokenStream, +) -> proc_macro::TokenStream { + // yes, this ignores everything except the codeblock + // it's fine. + let body = match input.into_iter().last().expect("empty") { + TokenTree::Group(group) if group.delimiter() == Delimiter::Brace => group.stream(), + _ => panic!("invalid function"), + }; + + let body = proc_macro2::TokenStream::from(body); + + return quote! { + fn main() { + let _sentry_guard = libpk::init_sentry(); + libpk::init_logging(env!("CARGO_CRATE_NAME")); + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + if let Err(error) = libpk::init_metrics() { + tracing::error!(?error, "failed to init metrics collector"); + }; + + tracing::info!("hello world"); + + let result: anyhow::Result<()> = async { #body }.await; + + if let Err(error) = result { + tracing::error!(?error, "failed to run service"); + }; + }); + } + } + .into(); +} diff --git a/crates/macros/src/lib.rs b/crates/macros/src/lib.rs new file mode 100644 index 00000000..db5a55b7 --- /dev/null +++ b/crates/macros/src/lib.rs @@ -0,0 +1,14 @@ +use proc_macro::TokenStream; + +mod entrypoint; +mod model; + +#[proc_macro_attribute] +pub fn main(args: TokenStream, input: TokenStream) -> TokenStream { + entrypoint::macro_impl(args, input) +} + +#[proc_macro_attribute] +pub fn pk_model(args: TokenStream, input: TokenStream) -> TokenStream { + model::macro_impl(args, input) +} diff --git a/crates/model_macros/src/lib.rs b/crates/macros/src/model.rs similarity index 99% rename from crates/model_macros/src/lib.rs rename to crates/macros/src/model.rs index 77f286e2..bfd219a7 100644 --- a/crates/model_macros/src/lib.rs +++ b/crates/macros/src/model.rs @@ -84,8 +84,7 @@ fn parse_field(field: syn::Field) -> ModelField { f } -#[proc_macro_attribute] -pub fn pk_model( +pub fn macro_impl( _args: proc_macro::TokenStream, input: proc_macro::TokenStream, ) -> proc_macro::TokenStream { diff --git a/crates/models/Cargo.toml b/crates/models/Cargo.toml index da90d79a..0fbc358c 100644 --- a/crates/models/Cargo.toml +++ b/crates/models/Cargo.toml @@ -5,7 +5,7 @@ edition = "2021" [dependencies] chrono = { workspace = true, features = ["serde"] } -model_macros = { path = "../model_macros" } +pk_macros = { path = "../macros" } sea-query = "0.32.1" serde = { workspace = true } serde_json = { workspace = true, features = ["preserve_order"] } diff --git a/crates/models/src/system.rs b/crates/models/src/system.rs index d59d5957..42b61fe4 100644 --- a/crates/models/src/system.rs +++ b/crates/models/src/system.rs @@ -1,6 +1,6 @@ use std::error::Error; -use model_macros::pk_model; +use pk_macros::pk_model; use chrono::NaiveDateTime; use sqlx::{postgres::PgTypeInfo, Database, Decode, Postgres, Type}; diff --git a/crates/models/src/system_config.rs b/crates/models/src/system_config.rs index d6b58a58..04b1995b 100644 --- a/crates/models/src/system_config.rs +++ b/crates/models/src/system_config.rs @@ -1,4 +1,4 @@ -use model_macros::pk_model; +use pk_macros::pk_model; use sqlx::{postgres::PgTypeInfo, Database, Decode, Postgres, Type}; use std::error::Error; diff --git a/crates/scheduled_tasks/src/main.rs b/crates/scheduled_tasks/src/main.rs index 56d58c7c..c689a99e 100644 --- a/crates/scheduled_tasks/src/main.rs +++ b/crates/scheduled_tasks/src/main.rs @@ -20,8 +20,8 @@ pub struct AppCtx { pub discord: Arc, } -libpk::main!("scheduled_tasks"); -async fn real_main() -> anyhow::Result<()> { +#[libpk::main] +async fn main() -> anyhow::Result<()> { let mut client_builder = twilight_http::Client::builder().token( libpk::config .discord