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

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

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

crates/avatars/src/main.rs

@@ -170,8 +170,8 @@ pub struct AppState {
     config: Arc<AvatarsConfig>,
 }
 
-libpk::main!("avatars");
-async fn real_main() -> anyhow::Result<()> {
+#[libpk::main]
+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
-
-#[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::<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");
-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(

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

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 pk_model(
+pub fn macro_impl(
     _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");
-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