Function main

examples/extension-internal-flush/src/main.rs:85–119

()

#[tokio::main]
async fn main() -> Result<(), Error> {
    tracing::init_default_subscriber();
    let (request_done_sender, request_done_receiver) = unbounded_channel::<()>();

    let flush_extension = Arc::new(FlushExtension::new(request_done_receiver));
    let extension = Extension::new()
        // Internal extensions only support INVOKE events.
        .with_events(&["INVOKE"])
        .with_events_processor(service_fn(|event| {
            let flush_extension = flush_extension.clone();
            async move { flush_extension.invoke(event).await }
        }))
        // Internal extension names MUST be unique within a given Lambda function.
        .with_extension_name("internal-flush")
        // Extensions MUST be registered before calling lambda_runtime::run(), which ends the Init
        // phase and begins the Invoke phase.
        .register()
        .await?;

    let handler = Arc::new(EventHandler::new(request_done_sender));

    tokio::try_join!(
        // always poll the handler function first before the flush extension,
        // this results in a smaller future due to not needing to track which was polled first
        // each time, and also a tiny latency savings
        biased;
        lambda_runtime::run(service_fn(|event| {
            let handler = handler.clone();
            async move { handler.invoke(event).await }
        })),
        extension.run(),
    )?;

    Ok(())
}