optimize api state extractor

Signed-off-by: Jason Volk <jason@zemos.net>

src/router/layers.rs

@@ -6,6 +6,7 @@ use axum::{
 };
 use axum_client_ip::SecureClientIpSource;
 use conduit::{error, Result, Server};
+use conduit_api::router::state::Guard;
 use conduit_service::Services;
 use http::{
 	header::{self, HeaderName},
@@ -35,7 +36,7 @@ const CONDUWUIT_CSP: &[&str] = &[
 
 const CONDUWUIT_PERMISSIONS_POLICY: &[&str] = &["interest-cohort=()", "browsing-topics=()"];
 
-pub(crate) fn build(services: &Arc<Services>) -> Result<Router> {
+pub(crate) fn build(services: &Arc<Services>) -> Result<(Router, Guard)> {
 	let server = &services.server;
 	let layers = ServiceBuilder::new();
 
@@ -85,7 +86,8 @@ pub(crate) fn build(services: &Arc<Services>) -> Result<Router> {
 		.layer(body_limit_layer(server))
 		.layer(CatchPanicLayer::custom(catch_panic));
 
-	Ok(router::build(services).layer(layers))
+	let (router, guard) = router::build(services);
+	Ok((router.layer(layers), guard))
 }
 
 #[cfg(any(feature = "zstd_compression", feature = "gzip_compression", feature = "brotli_compression"))]

src/router/router.rs

@@ -2,19 +2,20 @@ use std::sync::Arc;
 
 use axum::{response::IntoResponse, routing::get, Router};
 use conduit::Error;
-use conduit_api::State;
+use conduit_api::router::{state, state::Guard};
 use conduit_service::Services;
 use http::{StatusCode, Uri};
 use ruma::api::client::error::ErrorKind;
 
-pub(crate) fn build(services: &Arc<Services>) -> Router {
-	let router = Router::<State>::new();
-	let state = State::new(services.clone());
-
-	conduit_api::router::build(router, &services.server)
+pub(crate) fn build(services: &Arc<Services>) -> (Router, Guard) {
+	let router = Router::<state::State>::new();
+	let (state, guard) = state::create(services.clone());
+	let router = conduit_api::router::build(router, &services.server)
 		.route("/", get(it_works))
 		.fallback(not_found)
-		.with_state(state)
+		.with_state(state);
+
+	(router, guard)
 }
 
 async fn not_found(_uri: Uri) -> impl IntoResponse {

src/router/serve/mod.rs

@@ -18,7 +18,7 @@ pub(super) async fn serve(
 	let server = &services.server;
 	let config = &server.config;
 	let addrs = config.get_bind_addrs();
-	let app = layers::build(&services)?;
+	let (app, _guard) = layers::build(&services)?;
 
 	if cfg!(unix) && config.unix_socket_path.is_some() {
 		unix::serve(server, app, shutdown).await