Custom health check and graceful shutdown

[jonaro00]

Describe the feature

Writing down my ideas for what this feature can look like.

---

Expansion of the Service trait:

#[async_trait] // no longer needed?
// addition of Clone would mean the implementing service is a cheap-to-clone struct, similar to axum state
pub trait Service: Send + Clone {
    // rename to "start"?
    async fn bind(self, addr: SocketAddr) -> Result<(), Error>;

    async fn health_check(self) -> Result<(), Error> {
        Ok(())
    }

    async fn shutdown(self) -> Result<(), Error> {
        Ok(())
    }
}

Relevant code locations:

• Where runtime receives health checks

  shuttle/runtime/src/rt.rs

  Line 113 in 9fd65d9

  // TODO: A hook into the `Service` trait can be added here

• Where shutdown handling could be added

  shuttle/runtime/src/rt.rs

  Line 324 in 9fd65d9

CLI would also need to be adjusted to not kill the runtime process, but to send SIGTERM instead, and perhaps SIGKILL after a timeout.

Suggestion or Example of how the feature would be used

The main way would be to manually implement Service.

Some macro magic to allow optionally defining the health check and shutdown methods at the root level would be nice. Something like:

#[shuttle_runtime::main(
    health_check = my_health_check,
    shutdown = my_shutdown,
)]
async fn main() -> MyService {
    let state = (...);
    let router = (...).with_state(state.clone());
    Ok(MyService { state, router })
}

async fn my_health_check(service: MyService) -> Result<()> {
    healthz(State(service.state)).await?;

    Ok(())
}

// axum handler that might also be mounted on `/_healthz`
async fn healthz(State(state): State<MyState>) -> impl IntoResponse {
    ...
}

async fn my_shutdown(service: MyService) -> Result<()> {
    // ...

    Ok(())
}

Duplicate declaration

• I have searched the issues and this feature has not been requested before.

[alvissraghnall]

....