Description
First off, thanks for these useful crates! It's nice to see official Rust support from AWS.
The lambda-extension crate currently only supports external Lambda extensions (i.e., those running in a separate binary from the main runtime). AWS Lambda also supports internal Lambda extensions that run in the same process as the runtime, and it would be great if this crate also supported those.
The main requirement for supporting internal extensions is that each internal extension must register itself with the Lambda Runtime API before the runtime thread calls Next:
The
Initphase completes after the runtime and each registered extension indicate completion by sending aNextAPI request.
Once the main runtime calls Next, it's too late to register any additional extensions. The Lambda Runtime API returns an error if an extension tries to register itself at that point.
The current implementation of Extension::run() registers the extension with the Lambda Runtime API and then proceeds to call Next. Since Extension::run() returns a single future that includes both registration and invocation of the extension, there's no way for the caller to determine when registration is complete, and that it's safe to call Runtime::run() (which immediately calls Next and ends the Init phase). In other words, calling Extension::run() and Runtime::run() concurrently creates a race condition between registering the extensions and the runtime calling Next.
A straightforward solution would be to replace (or augment) the Extension::run() function with an API that looks like the following (note the new Extension::register() function and a new type representing a registered extension that's ready to be run):
let extension = Extension::new()
.with_events(&["INVOKE"])
.with_events_processor(...);
let registered_extension = extension.register().await;
futures::future::try_join(
registered_extension.run(),
Runtime::run(...),
).awaitA potential downside to this API is that .with_events will allow a user to attempt to register the SHUTDOWN event, which will fail at runtime since it's not permitted for internal extensions. Also, I'm not sure whether the telemetry and logs processing functionality works for internal extensions or not. These issues could be solved by having separate types for internal and external extensions, along with an enum to represent the supported events instead of accepting raw strings. However, that's a more involved API redesign that may not be worth it at this point.
I'm happy to contribute a PR once there's consensus on the desired API!