Support for internal Lambda extensions

[ramosbugs]

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 Init phase completes after the runtime and each registered extension indicate completion by sending a Next API 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(...),
).await

A 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!

[calavera]

I don't know much about internal extensions, so my suggestion would be to start with the most simple thing we can do to address the immediate problem.

Let's create that register function that only registers the extension. If a user sends events that are not supported by an internal extension, or tries unsupported apis, they'll receive an error. If the error is not clear, we can improve the api.

[bnusunny]

@ramosbugs Nice suggestion. I can work with you on this feature. I was too busy for the re:Invent this week. I will add comments next week.

[ramosbugs]

Thanks @bnusunny! I should have a PR out this weekend; just doing some more testing before submitting it.

[github-actions]

⚠️COMMENT VISIBILITY WARNING⚠️

Comments on closed issues are hard for the maintainers of this repository to see.
If you need more assistance, please open a new issue that references this one.
If you wish to keep having a conversation with other community members under this issue feel free to do so.