Update lambda handler function extraction to work with python 3.9

When running on lambda functions, the CodeGuru Profiler module should be
loaded by lambda bootstrap framework instead of the customer's code.
Then we load the customer's handler function by calling the bootstarp
code directly.

We currently handle python 3.8, 3.7 and 3.6 (with some tweeks), however
for python 3.9 the lambda bootstrap module is now delegating to the
[runtime interface client](https://github.com/aws/aws-lambda-python-runtime-interface-client)
so our code must be changed to also work with this.

This changes moves the code about finding the proper function that
can load a handler function into a separate method. Unit tests are
updated to cope with this.

Author: PapaPedro

codeguru_profiler_agent/aws_lambda/lambda_handler.py

@@ -1,34 +1,24 @@
 import os
-import logging
+import importlib
 from codeguru_profiler_agent.aws_lambda.profiler_decorator import with_lambda_profiler
 from codeguru_profiler_agent.agent_metadata.aws_lambda import HANDLER_ENV_NAME_FOR_CODEGURU_KEY
 HANDLER_ENV_NAME = "_HANDLER"
-logger = logging.getLogger(__name__)
 
 
 def restore_handler_env(original_handler, env=os.environ):
     env[HANDLER_ENV_NAME] = original_handler
 
 
-def load_handler(bootstrap_module, env=os.environ, original_handler_env_key=HANDLER_ENV_NAME_FOR_CODEGURU_KEY):
+def load_handler(handler_extractor, env=os.environ, original_handler_env_key=HANDLER_ENV_NAME_FOR_CODEGURU_KEY):
     try:
         original_handler_name = env.get(original_handler_env_key)
         if not original_handler_name:
             raise ValueError("Could not find module and function name from " + HANDLER_ENV_NAME_FOR_CODEGURU_KEY
                              + " environment variable")
 
-        # Delegate to the lambda code to load the customer's module.
-        if hasattr(bootstrap_module, '_get_handler'):
-            customer_handler_function = bootstrap_module._get_handler(original_handler_name)
-        else:
-            # TODO FIXME Review if the support for python 3.6 bootstrap can be improved.
-            # This returns both a init_handler and the function, we apply the init right away as we are in init process
-            init_handler, customer_handler_function = bootstrap_module._get_handlers(
-                handler=original_handler_name,
-                mode='event',  # with 'event' it will return the function as is (handlerfn in the lambda code)
-                               # 'http' would return wsgi.handle_one(sockfd, ('localhost', 80), handlerfn) instead
-                invokeid='unknown_id')  # FIXME invokeid is used for error handling, need to see if we can get it
-            init_handler()
+        # Delegate to the lambda code to load the customer's module and function.
+        customer_handler_function = handler_extractor(original_handler_name)
+
         restore_handler_env(original_handler_name, env)
         return customer_handler_function
     except:
@@ -39,10 +29,60 @@ def load_handler(bootstrap_module, env=os.environ, original_handler_env_key=HAND
         raise
 
 
-# Load the customer's handler, this should be done at import time which means it is done when lambda frameworks
-# loads our module. We load the bootstrap module by string name so that IDE does not complain
-lambda_bootstrap_module = __import__("bootstrap")
-handler_function = load_handler(lambda_bootstrap_module)
+def _python36_extractor(bootstrap_module, original_handler_name):
+    """
+    The lambda bootstrap code for python 3.6 was different than for later versions, instead of the _get_handler
+    function there was a more complex _get_handlers function with more parameters
+    """
+    # TODO FIXME Review if the support for python 3.6 bootstrap can be improved.
+    # This returns both a init_handler and the function, we apply the init right away as we are in init process
+    init_handler, customer_handler_function = bootstrap_module._get_handlers(
+        handler=original_handler_name,
+        mode='event',  # with 'event' it will return the function as is (handlerfn in the lambda code)
+        # 'http' would return wsgi.handle_one(sockfd, ('localhost', 80), handlerfn) instead
+        invokeid='unknown_id')  # FIXME invokeid is used for error handling, need to see if we can get it
+    init_handler()
+    return customer_handler_function
+
+
+def get_lambda_handler_extractor():
+    """
+    This loads and returns a function from lambda or RIC source code that is able to load the customer's
+    handler function.
+    WARNING !! This is a bit dangerous since we are calling internal functions from other modules that we do not
+    officially depend on. The idea is that this code should run only in a lambda function environment where we can know
+    what is available. However if lambda developers decide to change their internal code it could impact this !
+    """
+    # First try to load the lambda RIC if it is available (i.e. python 3.9)
+    # See https://github.com/aws/aws-lambda-python-runtime-interface-client
+    ric_bootstrap_module = _try_to_load_module("awslambdaric.bootstrap")
+    if ric_bootstrap_module is not None and hasattr(ric_bootstrap_module, '_get_handler'):
+        return ric_bootstrap_module._get_handler
+
+    # If no RIC module is available there should be a bootstrap module available
+    # do not catch ModuleNotFoundError exceptions here as we cannot do anything if this fails.
+    bootstrap_module = importlib.import_module("bootstrap")
+    if hasattr(bootstrap_module, '_get_handler'):
+        return bootstrap_module._get_handler
+    else:
+        return lambda handler_name: _python36_extractor(bootstrap_module, handler_name)
+
+
+def _try_to_load_module(module_name):
+    try:
+        return importlib.import_module(module_name)
+    except ModuleNotFoundError:
+        return None
+
+
+# We need to load the customer's handler function since the lambda framework loaded our function instead.
+# We want to delegate this work to the lambda framework so we need to find the appropriate method that does it
+# (depends on python versions) so we can call it.
+# This should be done at import time which means it is done when lambda frameworks loads our module
+handler_extractor = get_lambda_handler_extractor()
+
+# Now load the actual customer's handler function.
+handler_function = load_handler(handler_extractor)
 
 
 # WARNING: Do not rename this file, this function or HANDLER_ENV_NAME_FOR_CODEGURU without changing the bootstrap script

test/unit/aws_lambda/test_lambda_handler.py

@@ -11,6 +11,7 @@ def handler_function(event, context):
 
 
 init_handler_has_been_called = False
+python36_extractor_has_been_called = False
 
 
 def init_handler():
@@ -27,10 +28,18 @@ def _get_handler(self, handler):
 class BootstrapPython36ModuleMock:
     # for python3.6 version of lambda runtime bootstrap
     def _get_handlers(self, handler, mode, invokeid):
+        global python36_extractor_has_been_called
+        python36_extractor_has_been_called = True
         if handler == "handler_module.handler_function" and mode == "event":
             return init_handler, handler_function
 
 
+class RicBootstrapModuleMock:
+    def _get_handler(self, handler):
+        if handler == "handler_module.handler_function":
+            return handler_function
+
+
 class TestLambdaHandler:
     class TestWhenLambdaHandlerModuleIsLoaded:
         @pytest.fixture(autouse=True)
@@ -53,48 +62,86 @@ def test_call_handler_calls_the_inner_handler(self):
             assert lambda_handler_module.call_handler(event="expected_event",
                                                       context="expected_context") == "expected result"
 
-    class TestLoadModuleFunction:
-        class TestWhenPython38LambdaBootstrapCalls:
-            class TestWhenHandlerEnvIsSetProperly:
-                @before
-                def before(self):
-                    self.bootstrap = BootstrapModuleMock()
-                    self.env = {"HANDLER_ENV_NAME_FOR_CODEGURU": "handler_module.handler_function"}
-
-                def test_it_returns_the_handler_function(self):
-                    from codeguru_profiler_agent.aws_lambda.lambda_handler import load_handler
-                    assert load_handler(self.bootstrap, self.env) == handler_function
-
-                def test_it_resets_handler_env_variable(self):
-                    from codeguru_profiler_agent.aws_lambda.lambda_handler import load_handler
-                    load_handler(self.bootstrap, self.env)
-                    assert self.env['_HANDLER'] == "handler_module.handler_function"
-
-            class TestWhenHandlerEnvIsMissing:
-                @before
-                def before(self):
-                    self.bootstrap = BootstrapModuleMock()
-                    self.env = {}
-
-                def test_it_throws_value_error(self):
-                    with pytest.raises(ValueError):
-                        from codeguru_profiler_agent.aws_lambda.lambda_handler import load_handler
-                        load_handler(self.bootstrap, self.env)
+    class TestGetHandlerExtractor:
+        class TestWhenRicIsAvailable:
+            @pytest.fixture(autouse=True)
+            def around(self):
+                # simulate that we are in a lambda environment where the awslambdaric.bootstrap module is available
+                self.module_available = RicBootstrapModuleMock()
+                sys.modules['awslambdaric.bootstrap'] = self.module_available
+                yield
+                del sys.modules['awslambdaric.bootstrap']
+
+            def test_it_loads_the_ric_module_code(self):
+                from codeguru_profiler_agent.aws_lambda.lambda_handler import get_lambda_handler_extractor
+                result = get_lambda_handler_extractor()
+                assert result == self.module_available._get_handler
+
+        class TestWhenLambdaBootstrapIsAvailable:
+            @pytest.fixture(autouse=True)
+            def around(self):
+                # simulate that we are in a lambda environment where the awslambdaric.bootstrap module is not available
+                # but bootstrap from lambda is available.
+                self.module_available = BootstrapModuleMock()
+                if 'awslambdaric.bootstrap' in sys.modules:
+                    del sys.modules['awslambdaric.bootstrap']
+                sys.modules['bootstrap'] = self.module_available
+                yield
+                del sys.modules['bootstrap']
+
+            def test_it_loads_the_lambda_module_code(self):
+                from codeguru_profiler_agent.aws_lambda.lambda_handler import get_lambda_handler_extractor
+                result = get_lambda_handler_extractor()
+                assert result == self.module_available._get_handler
 
         class TestWhenPython36LambdaBootstrapCalls:
             class TestWhenHandlerEnvIsSetProperly:
-                @before
-                def before(self):
-                    self.bootstrap = BootstrapPython36ModuleMock()
-                    self.env = {"HANDLER_ENV_NAME_FOR_CODEGURU": "handler_module.handler_function"}
+                @pytest.fixture(autouse=True)
+                def around(self):
+                    # simulate that we are in a lambda environment where the awslambdaric.bootstrap module is available
+                    sys.modules['bootstrap'] = BootstrapPython36ModuleMock()
                     global init_handler_has_been_called
                     init_handler_has_been_called = False
+                    global python36_extractor_has_been_called
+                    python36_extractor_has_been_called = False
+                    yield
+                    del sys.modules['bootstrap']
 
-                def test_it_returns_the_handler_function(self):
-                    from codeguru_profiler_agent.aws_lambda.lambda_handler import load_handler
-                    assert load_handler(self.bootstrap, self.env) == handler_function
+                def test_it_uses_the_old_bootstrap_code(self):
+                    from codeguru_profiler_agent.aws_lambda.lambda_handler import get_lambda_handler_extractor
+                    # call extractor
+                    get_lambda_handler_extractor()("handler_module.handler_function")
+                    assert python36_extractor_has_been_called
 
                 def test_it_calls_the_init_handler(self):
-                    from codeguru_profiler_agent.aws_lambda.lambda_handler import load_handler
-                    load_handler(self.bootstrap, self.env)
+                    from codeguru_profiler_agent.aws_lambda.lambda_handler import get_lambda_handler_extractor
+                    # call extractor
+                    get_lambda_handler_extractor()("handler_module.handler_function")
                     assert init_handler_has_been_called
+
+    class TestLoadHandlerFunction:
+        class TestWhenHandlerEnvIsSetProperly:
+            @before
+            def before(self):
+                self.extractor = BootstrapModuleMock()._get_handler
+                self.env = {"HANDLER_ENV_NAME_FOR_CODEGURU": "handler_module.handler_function"}
+
+            def test_it_returns_the_handler_function(self):
+                from codeguru_profiler_agent.aws_lambda.lambda_handler import load_handler
+                assert load_handler(self.extractor, self.env) == handler_function
+
+            def test_it_resets_handler_env_variable(self):
+                from codeguru_profiler_agent.aws_lambda.lambda_handler import load_handler
+                load_handler(self.extractor, self.env)
+                assert self.env['_HANDLER'] == "handler_module.handler_function"
+
+        class TestWhenHandlerEnvIsMissing:
+            @before
+            def before(self):
+                self.extractor = BootstrapModuleMock()._get_handler
+                self.env = {}
+
+            def test_it_throws_value_error(self):
+                with pytest.raises(ValueError):
+                    from codeguru_profiler_agent.aws_lambda.lambda_handler import load_handler
+                    load_handler(self.extractor, self.env)