Skip to main content

toolkit/
contracts.rs

1use async_trait::async_trait;
2use axum::Router;
3use tokio_util::sync::CancellationToken;
4
5pub use crate::api::openapi_registry::OpenApiRegistry;
6
7/// System capability: receives runtime internals before init.
8///
9/// This trait is internal to toolkit and only used by gears with the "system" capability.
10/// Normal user gears don't implement this.
11#[async_trait]
12pub trait SystemCapability: Send + Sync {
13    /// Optional pre-init hook for system gears.
14    ///
15    /// This runs BEFORE `init()` has completed for ALL gears, and only for system gears.
16    ///
17    /// Default implementation is a no-op so most gears don't need to implement it.
18    ///
19    /// # Errors
20    /// Returns an error if system wiring fails.
21    fn pre_init(&self, _sys: &crate::runtime::SystemContext) -> anyhow::Result<()> {
22        Ok(())
23    }
24
25    /// Optional post-init hook for system gears.
26    ///
27    /// This runs AFTER `init()` has completed for ALL gears, and only for system gears.
28    ///
29    /// Default implementation is a no-op so most gears don't need to implement it.
30    async fn post_init(&self, _sys: &crate::runtime::SystemContext) -> anyhow::Result<()> {
31        Ok(())
32    }
33}
34
35/// Core gear: DI/wiring; do not rely on migrated schema here.
36#[async_trait]
37pub trait Gear: Send + Sync + 'static {
38    async fn init(&self, ctx: &crate::context::GearCtx) -> anyhow::Result<()>;
39}
40
41/// Database capability: gears provide migrations, runtime executes them.
42///
43/// # Security
44///
45/// Gears MUST NOT receive raw database connections. They only return migration definitions.
46#[cfg(feature = "db")]
47pub trait DatabaseCapability: Send + Sync {
48    fn migrations(&self) -> Vec<Box<dyn sea_orm_migration::MigrationTrait>>;
49}
50
51/// REST API capability: Pure wiring; must be sync. Runs AFTER DB migrations.
52pub trait RestApiCapability: Send + Sync {
53    /// Register REST routes for this gear.
54    ///
55    /// # Errors
56    /// Returns an error if route registration fails.
57    fn register_rest(
58        &self,
59        ctx: &crate::context::GearCtx,
60        router: Router,
61        openapi: &dyn OpenApiRegistry,
62    ) -> anyhow::Result<Router>;
63}
64
65/// API Gateway capability: handles gateway hosting with prepare/finalize phases.
66/// Must be sync. Runs during REST phase, but doesn't start the server.
67#[allow(dead_code)]
68pub trait ApiGatewayCapability: Send + Sync + 'static {
69    /// Prepare a base Router (e.g., global middlewares, /healthz) and optionally touch `OpenAPI` meta.
70    /// Do NOT start the server here.
71    ///
72    /// # Errors
73    /// Returns an error if router preparation fails.
74    fn rest_prepare(&self, ctx: &crate::context::GearCtx, router: Router)
75    -> anyhow::Result<Router>;
76
77    /// Finalize before start: attach /openapi.json, /docs, persist the Router internally if needed.
78    /// Do NOT start the server here.
79    ///
80    /// # Errors
81    /// Returns an error if router finalization fails.
82    fn rest_finalize(
83        &self,
84        ctx: &crate::context::GearCtx,
85        router: Router,
86    ) -> anyhow::Result<Router>;
87
88    // Return OpenAPI registry of the gear, e.g., to register endpoints
89    fn as_registry(&self) -> &dyn OpenApiRegistry;
90}
91
92/// Capability for gears that have a long-running background task.
93///
94/// # Shutdown Contract
95///
96/// The `stop` method receives a **deadline token** that implements two-phase shutdown:
97///
98/// 1. **Graceful stop request**: When `stop(deadline_token)` is called, the `deadline_token`
99///    is *not* cancelled. This is the signal to begin graceful shutdown.
100///
101/// 2. **Hard-stop deadline**: After the runtime's `shutdown_deadline` expires (default 30s),
102///    the `deadline_token` is cancelled. This signals that graceful shutdown time is over
103///    and the gear should abort immediately.
104///
105/// ## Recommended Implementation Pattern
106///
107/// ```ignore
108/// async fn stop(&self, deadline_token: CancellationToken) -> anyhow::Result<()> {
109///     // 1. Request cooperative shutdown of child tasks
110///     self.request_graceful_shutdown();
111///
112///     // 2. Wait for graceful completion OR hard-stop deadline
113///     tokio::select! {
114///         _ = self.wait_for_graceful_completion() => {
115///             // Graceful shutdown succeeded
116///         }
117///         _ = deadline_token.cancelled() => {
118///             // Deadline reached, force abort
119///             self.force_abort();
120///         }
121///     }
122///     Ok(())
123/// }
124/// ```
125///
126/// ## Important Notes
127///
128/// - The `deadline_token` passed to `stop()` is a **fresh token**, not the root cancellation
129///   token that triggered the shutdown. This allows gears to implement real graceful shutdown.
130/// - Gears should NOT assume the token is already cancelled when `stop()` is called.
131/// - The `WithLifecycle` wrapper handles this contract automatically via its `stop_timeout`.
132#[async_trait]
133pub trait RunnableCapability: Send + Sync {
134    /// Start the gear's background task.
135    ///
136    /// The `cancel` token is a child of the runtime's root cancellation token.
137    /// When cancelled, the gear should stop its background work.
138    async fn start(&self, cancel: CancellationToken) -> anyhow::Result<()>;
139
140    /// Stop the gear's background task.
141    ///
142    /// The `deadline_token` implements two-phase shutdown:
143    /// - Initially not cancelled: begin graceful shutdown
144    /// - When cancelled: graceful period expired, abort immediately
145    ///
146    /// See trait-level documentation for the full shutdown contract.
147    async fn stop(&self, deadline_token: CancellationToken) -> anyhow::Result<()>;
148}
149
150/// Represents a gRPC service registration callback used by the gRPC hub.
151///
152/// Each gear that exposes gRPC services provides one or more of these.
153/// The `register` closure adds the service into the provided `RoutesBuilder`.
154#[cfg(feature = "otel")]
155pub struct RegisterGrpcServiceFn {
156    pub service_name: &'static str,
157    pub register: Box<dyn Fn(&mut tonic::service::RoutesBuilder) + Send + Sync>,
158}
159
160#[cfg(not(feature = "otel"))]
161pub struct RegisterGrpcServiceFn {
162    pub service_name: &'static str,
163}
164
165/// gRPC Service capability: gears that export gRPC services.
166///
167/// The runtime will call this during the gRPC registration phase to collect
168/// all services that should be exposed on the shared gRPC server.
169#[async_trait]
170pub trait GrpcServiceCapability: Send + Sync {
171    /// Returns all gRPC services this gear wants to expose.
172    ///
173    /// Each installer adds one service to the `tonic::Server` builder.
174    async fn get_grpc_services(
175        &self,
176        ctx: &crate::context::GearCtx,
177    ) -> anyhow::Result<Vec<RegisterGrpcServiceFn>>;
178}
179
180/// gRPC Hub capability: hosts the gRPC server.
181///
182/// This trait is implemented by the single gear responsible for hosting
183/// the `tonic::Server` instance. Only one gear per process should implement this.
184pub trait GrpcHubCapability: Send + Sync {
185    /// Returns the bound endpoint after the server starts listening.
186    ///
187    /// Examples:
188    /// - TCP: `http://127.0.0.1:50652`
189    /// - Unix socket: `unix:///path/to/socket`
190    /// - Named pipe: `pipe://\\.\pipe\name`
191    ///
192    /// Returns `None` if the server hasn't started listening yet.
193    fn bound_endpoint(&self) -> Option<String>;
194}