Support for running code after the response has been returned

While this request might sound similar to the streaming response feature [see https://github.com/aws/aws-lambda-python-runtime-interface-client/issues/100] now available to Node.js Lambda functions (see https://docs.aws.amazon.com/lambda/latest/dg/configuration-response-streaming.html), it is a separate topic.

We have come across two different use cases where we'd like to run some Python code after returning our response to, say, API Gateway. One use case is running Python's garbage collector "manually" instead of automatically. Rather than have Python decide to run garbage collection in the middle of our processing of an incoming event, and cause our users extra latency, we want to collect the garbage after returning the response.

The other use case involves updating DynamoDb with some bookkeeping about the event just serviced. While we can use an external Lambda Extension to do this, it requires moving a non-trivial amount of data to the Extension process, which seems wasteful assuming we could run this DynamoDb update after returning the response.

As a proof of concept, I wrote a monkeypatch to explore feasibility. I wrote it as a class that I can choose to instance or not at the top of my lambda handler.

class PostResponseHook:
    def __init__(self):
        import awslambdaric.bootstrap

        self.original_handle_event_request = awslambdaric.bootstrap.handle_event_request
        awslambdaric.bootstrap.handle_event_request = self.new_handle_event_request

    def new_handle_event_request(self, *args, **kwargs):
        self.original_handle_event_request(*args, **kwargs)
        self.hook()

    def hook(self, *args, **kwargs):
        # do something after the response was returned to the caller here
        pass

The monkeypatch relies on the current structure of handle_event_request in bootstrap.py. The Lambda function is not finished processing the incoming Event in handle_event_request, but, rather, continues back to the "infinite loop" in run.

        while True:
            event_request = lambda_runtime_client.wait_next_invocation()

            _GLOBAL_AWS_REQUEST_ID = event_request.invoke_id

            update_xray_env_variable(event_request.x_amzn_trace_id)

            handle_event_request(
                lambda_runtime_client,
                request_handler,
                event_request.invoke_id,
                event_request.event_body,
                event_request.content_type,
                event_request.client_context,
                event_request.cognito_identity,
                event_request.invoked_function_arn,
                event_request.deadline_time_in_ms,
                log_sink,
            )

The call to lambda_runtime_client.wait_next_invocation() in this loop is where the "long poll" for the next Lambda Event occurs (and, presumably, is where the Lambda service will freeze or shutdown this Lambda function instance under most circumstances (see https://docs.aws.amazon.com/lambda/latest/dg/lambda-runtime-environment.html for details.)

=============================================

Given that it seems technically feasible to offer this hook, how could we add this as a supported feature instead of relying on monkeypatching? It seems that we'd need a way to optionally "register" our hook method. Which would mean exposing a callable from somewhere in awslambdaric. Since I don't know of an existing API already exposed to user code by awslambdaric, I don't know where this register_post_response_hook method belongs.

Alternatively, maybe opting in for this feature would be by means of some special Layer that awslambdaric would look for, avoiding the need for explicit opt-in registration by the user's handler?

=============================================

But I'm probably getting ahead of myself here! The first question is: would the awslambdaric project consider this kind of hook as desirable, given that it can make a really big difference to latency-sensitive customer use-cases?

aws / aws-lambda-python-runtime-interface-client

Support for running code after the response has been returned #113