use anyhow::Result;
use async_trait::async_trait;
use axum::Router;
use folk_api::{PluginContext, RpcMethodDef, ServerPlugin};
use tonic::transport::Server;
use tracing::info;
use crate::config::GrpcConfig;
use crate::service::{GrpcState, grpc_handler};
pub struct GrpcPlugin {
config: GrpcConfig,
}
impl GrpcPlugin {
pub fn new(config: GrpcConfig) -> Self {
Self { config }
}
}
#[async_trait]
impl ServerPlugin for GrpcPlugin {
fn name(&self) -> &'static str {
"grpc"
}
async fn run(&self, ctx: PluginContext) -> Result<()> {
let state = GrpcState {
executor: ctx.executor.clone(),
};
let router = Router::new().fallback(grpc_handler).with_state(state);
let routes = tonic::service::Routes::from(router);
let mut sd = ctx.shutdown.clone();
info!(listen = %self.config.listen, "gRPC server listening");
Server::builder()
.add_routes(routes)
.serve_with_shutdown(self.config.listen, async move {
sd.changed().await.ok();
})
.await?;
Ok(())
}
fn rpc_methods(&self) -> Vec<RpcMethodDef> {
vec![RpcMethodDef::new(
"grpc.services",
"list registered gRPC service names",
)]
}
}