Load function in child process.

yxd-hde · yxd-hde · commit 62563884cbc6 · 2019-12-10T11:42:27.000+09:00
diff --git a/lambda_local/main.py b/lambda_local/main.py
@@ -38,10 +38,32 @@ def filter(self, record):
         return True
 
 
+class FunctionLoader():
+    def __init__(self,
+                 request_id=None,
+                 source=None,
+                 function_name=None,
+                 library_path=None,
+                 func=None):
+        self.request_id = request_id
+        self.source = source
+        self.function_name = function_name
+        self.library_path = library_path
+
+        self.func = func
+
+    def load(self):
+        if self.library_path is not None:
+            load_lib(self.library_path)
+
+        self.func = load_source(
+            self.request_id, self.source, self.function_name)
+
+
 def call(func, event, context, environment_variables={}):
     export_variables(environment_variables)
-
-    return _runner(func, event, context)
+    loader = FunctionLoader(func=func)
+    return _runner(loader, event, context)
 
 
 def run(args):
@@ -53,17 +75,19 @@ def run(args):
         args.timeout,
         invoked_function_arn=args.arn_string,
         function_version=args.version_name)
-    if args.library is not None:
-        load_lib(args.library)
-    func = load(c.aws_request_id, args.file, args.function)
+    loader = FunctionLoader(
+        request_id=c.aws_request_id,
+        source=args.file,
+        function_name=args.function,
+        library_path=args.library)
 
-    (result, err_type) = _runner(func, e, c)
+    (result, err_type) = _runner(loader, e, c)
 
     if err_type is not None:
         sys.exit(EXITCODE_ERR)
 
 
-def _runner(func, event, context):
+def _runner(loader, event, context):
     logger = logging.getLogger()
 
     logger.info("Event: {}".format(event))
@@ -73,7 +97,7 @@ def _runner(func, event, context):
     queue = multiprocessing.Queue()
     p = multiprocessing.Process(
         target=execute_in_process,
-        args=(queue, func, event, context,))
+        args=(queue, loader, event, context,))
     p.start()
     (result, err_type, duration) = queue.get()
     p.join()
@@ -94,7 +118,7 @@ def load_lib(path):
     sys.path.append(os.path.abspath(path))
 
 
-def load(request_id, path, function_name):
+def load_source(request_id, path, function_name):
     mod_name = 'request-' + str(request_id)
 
     file_path = os.path.abspath(path)
@@ -142,9 +166,11 @@ def execute(func, event, context):
     return result, err_type
 
 
-def execute_in_process(queue, func, event, context):
+def execute_in_process(queue, loader, event, context):
+    if loader.func is None:
+        loader.load()
     start_time = timeit.default_timer()
-    result, err_type = execute(func, event, context)
+    result, err_type = execute(loader.func, event, context)
     end_time = timeit.default_timer()
     duration = (end_time - start_time) * 1000
 
