chore(rust): tidier libpk::main macro impl

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",

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");
+#[libpk::main]
-async fn real_main() -> anyhow::Result<()> {
+async fn main() -> anyhow::Result<()> {
     let db = libpk::db::init_data_db().await?;
     let redis = libpk::db::init_redis().await?;
 

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");
+#[libpk::main]
-async fn real_main() -> anyhow::Result<()> {
+async fn main() -> anyhow::Result<()> {
     let config = libpk::config
         .avatars
         .as_ref()

crates/avatars/src/main.rs

@@ -170,8 +170,8 @@ pub struct AppState {
     config: Arc<AvatarsConfig>,
 }
 
-libpk::main!("avatars");
+#[libpk::main]
-async fn real_main() -> anyhow::Result<()> {
+async fn main() -> anyhow::Result<()> {
     let config = libpk::config
         .avatars
         .as_ref()

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 }

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
+#[libpk::main]
-
-#[tokio::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::<UdpSocket>::with_timeout(address, Duration::from_secs(3));
     let (client, bg) = AsyncClient::connect(stream).await?;

crates/gateway/src/main.rs

@@ -25,8 +25,8 @@ mod logger;
 
 const RUNTIME_CONFIG_KEY_EVENT_TARGET: &'static str = "event_target";
 
-libpk::main!("gateway");
+#[libpk::main]
-async fn real_main() -> anyhow::Result<()> {
+async fn main() -> anyhow::Result<()> {
     let redis = libpk::db::init_redis().await?;
 
     let runtime_config = Arc::new(

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");
+#[libpk::main]
-async fn real_main() -> anyhow::Result<()> {
+async fn main() -> anyhow::Result<()> {
     let db = libpk::db::init_messages_db().await?;
 
     let mut client_builder = twilight_http::Client::builder()

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 }

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(())
-        }
-    };
-}

crates/macros/Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "model_macros"
+name = "pk_macros"
 version = "0.1.0"
 edition = "2021"
 

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();
+}

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)
+}

crates/macros/src/model.rs

@@ -84,8 +84,7 @@ fn parse_field(field: syn::Field) -> ModelField {
     f
 }
 
-#[proc_macro_attribute]
+pub fn macro_impl(
-pub fn pk_model(
     _args: proc_macro::TokenStream,
     input: proc_macro::TokenStream,
 ) -> proc_macro::TokenStream {

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"] }

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};

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;

crates/scheduled_tasks/src/main.rs

@@ -20,8 +20,8 @@ pub struct AppCtx {
     pub discord: Arc<twilight_http::Client>,
 }
 
-libpk::main!("scheduled_tasks");
+#[libpk::main]
-async fn real_main() -> anyhow::Result<()> {
+async fn main() -> anyhow::Result<()> {
     let mut client_builder = twilight_http::Client::builder().token(
         libpk::config
             .discord