agent: use a timeout when shutting down the runtime

We've observed cases where the agent hangs when shutting down the runtime after
an error occurs. This changes it to use the same pattern that's used elsewhere
in our code base for shutting down the tokio runtime with a timeout, to prevent
it from hanging indefinitely.

crates/agent/src/handlers.rs

@@ -28,9 +28,10 @@ pub struct AgentNotification {
 }
 
 /// Handler is the principal trait implemented by the various task-specific
-/// event handlers that the agent runs.
+/// event handlers that the agent runs. They need to be `Send` because we
+/// spawn the handler invocations on a multithreaded runtime.
 #[async_trait::async_trait]
-pub trait Handler {
+pub trait Handler: Send {
     async fn handle(&mut self, pg_pool: &sqlx::PgPool) -> anyhow::Result<HandlerStatus>;
 
     fn table_name(&self) -> &'static str;

crates/agent/src/main.rs

@@ -48,8 +48,7 @@ struct Args {
     allow_local: bool,
 }
 
-#[tokio::main]
-async fn main() -> Result<(), anyhow::Error> {
+fn main() -> Result<(), anyhow::Error> {
     // Use reasonable defaults for printing structured logs to stderr.
     let subscriber = tracing_subscriber::FmtSubscriber::builder()
         .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
@@ -59,6 +58,19 @@ async fn main() -> Result<(), anyhow::Error> {
     let args = Args::parse();
     tracing::info!(?args, "started!");
 
+    let runtime = tokio::runtime::Builder::new_multi_thread()
+        .enable_all()
+        .build()?;
+
+    let task = runtime.spawn(async move { async_main(args).await });
+    let result = runtime.block_on(task);
+
+    tracing::info!(?result, "main function completed, shutting down runtime");
+    runtime.shutdown_timeout(std::time::Duration::from_secs(5));
+    result?
+}
+
+async fn async_main(args: Args) -> Result<(), anyhow::Error> {
     let bindir = std::fs::canonicalize(args.bindir)
         .context("canonicalize --bin-dir")?
         .into_os_string()