Lumigo extension (#151)

Author: saartochner-lumigo

scripts/prepare_extension_layer.sh

@@ -0,0 +1,25 @@
+echo "Extension layer"
+python -W ignore setup.py --quiet bdist_wheel  # create .egg-info
+pushd src > /dev/null || exit
+echo "-create directories to zip"
+echo "--extension"
+mkdir extensions
+cp lumigo_tracer/extension/bootstrap/lumigo extensions/
+echo "--tracer"
+mkdir extension-python-modules
+cp -R lumigo_tracer.egg-info extension-python-modules/
+cp -R lumigo_tracer extension-python-modules/
+echo "--python runtime"
+aws s3 cp --quiet s3://lumigo-runtimes/python/lean-python-runtime-37.zip runtime.zip
+unzip -q runtime.zip
+echo "--special temp file"
+touch preview-extensions-ggqizro707
+echo "-zipping"
+zip -qr "extensions.zip" "extensions" "extension-python-modules" "python-runtime" "preview-extensions-ggqizro707"  # take all the directory
+echo "-publish"
+../utils/common_bash/create_layer.sh --layer-name "extensions-layer" --region ALL --package-folder "extensions.zip" --version $(git describe --abbrev=0 --tags) --runtimes  ""
+rm -rf extensions extension-python-modules extensions.zip runtime.zip python-runtime __MACOSX
+popd > /dev/null || exit
+
+echo "\nDone.\n"
+

src/lumigo_tracer/extension/__init__.py

@@ -0,0 +1,11 @@
+#!/bin/sh
+
+set -eo pipefail
+
+# python runtime initialization
+if [[ $LD_LIBRARY_PATH != *python* ]]; then
+  export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/python-runtime/var/lang/lib
+  export PYTHONPATH=$PYTHONPATH:/opt/extension-python-modules:/opt/python-runtime/var/runtime
+fi
+
+/opt/python-runtime/var/lang/bin/python3.7 /opt/extension-python-modules/lumigo_tracer/extension/main.py &

src/lumigo_tracer/extension/bootstrap/lumigo

@@ -0,0 +1,74 @@
+import os
+from dataclasses import dataclass, asdict
+
+from typing import Dict, Optional, List, Union
+from datetime import datetime
+
+from lumigo_tracer.extension.extension_utils import get_current_bandwidth, get_extension_logger
+from lumigo_tracer import utils
+from lumigo_tracer.extension.lambda_service import LambdaService
+from lumigo_tracer.extension.sampler import Sampler
+from lumigo_tracer.utils import lumigo_safe_execute
+
+SPAN_TYPE = "extensionExecutionEnd"
+
+
+@dataclass
+class ExtensionEvent:
+    token: str
+    started: int
+    requestId: str
+    networkBytesUsed: int
+    cpuUsageTime: List[Dict[str, Union[float, int]]]
+    type: str = SPAN_TYPE
+
+
+class LumigoExtension:
+    def __init__(self, lambda_service: LambdaService):
+        self.lambda_service: LambdaService = lambda_service
+        self.sampler: Sampler = Sampler()
+        self.start_time: Optional[datetime] = None
+        self.request_id: Optional[str] = None
+        self.bandwidth: Optional[int] = None
+
+    def start_new_invocation(self, event: Dict[str, str]):
+        with lumigo_safe_execute("Extension: start_new_invocation"):
+            current_bandwidth = get_current_bandwidth()
+            if self.request_id:
+                self._finish_previous_invocation(current_bandwidth)
+            self.sampler.start_sampling()
+            self.lambda_service.ready_for_next_event()
+            self.request_id = event.get("requestId")
+            self.start_time = datetime.now()
+            self.bandwidth = current_bandwidth
+
+    def shutdown(self):
+        with lumigo_safe_execute("Extension: shutdown"):
+            current_bandwidth = get_current_bandwidth()
+            self._finish_previous_invocation(current_bandwidth)
+
+    def _finish_previous_invocation(self, current_bandwidth: Optional[int]):
+        self.sampler.stop_sampling()
+        token = os.environ.get(utils.LUMIGO_TOKEN_KEY)
+        if not token:
+            get_extension_logger().warning(
+                f"Skip sending data: No token was found. Request id: {self.request_id}"
+            )
+            return
+        if not (
+            self.request_id
+            and self.start_time  # noqa: W503
+            and self.sampler.get_samples()  # noqa: W503
+            and self.bandwidth  # noqa: W503
+            and current_bandwidth  # noqa: W503
+        ):
+            get_extension_logger().warning("Skip sending data: unable retrieving all data")
+            return
+        span = ExtensionEvent(
+            token=token,
+            started=int(self.start_time.timestamp() * 1000),
+            requestId=self.request_id,
+            networkBytesUsed=current_bandwidth - self.bandwidth,
+            cpuUsageTime=[s.dump() for s in self.sampler.get_samples()],
+        )
+        utils.report_json(os.environ.get("AWS_REGION", "us-east-1"), msgs=[asdict(span)])

src/lumigo_tracer/extension/extension.py

@@ -0,0 +1,38 @@
+import json
+import os
+import urllib.request
+from typing import Optional, Dict
+
+from lumigo_tracer.utils import get_logger, lumigo_safe_execute
+
+
+def get_current_cpu_time() -> Optional[int]:
+    """
+    :return: the total number of milliseconds that being used by the CPU.
+    """
+    with lumigo_safe_execute("Extension: get cpu time"):
+        total = 0
+        with open("/proc/stat", "r") as stats:
+            for line in stats.readlines():
+                if line.startswith("cpu "):
+                    parts = line.split()
+                    total += (int(parts[1]) + int(parts[3])) * 10
+        return total
+
+
+def get_current_bandwidth() -> Optional[int]:
+    with lumigo_safe_execute("Extension: get bandwidth"):
+        with open("/proc/net/netstat", "r") as stats:
+            last = stats.read().splitlines()[-1]
+            parts = last.split()
+        return int(parts[7]) + int(parts[8])
+
+
+def request_event(extension_id: str) -> Dict[str, str]:
+    url = f"http://{os.environ['AWS_LAMBDA_RUNTIME_API']}/2020-01-01/extension/event/next"
+    headers = {"Lambda-Extension-Identifier": extension_id}
+    return json.loads(urllib.request.urlopen(urllib.request.Request(url, headers=headers)).read())
+
+
+def get_extension_logger():
+    return get_logger("lumigo-extension")

src/lumigo_tracer/extension/extension_utils.py

@@ -0,0 +1,31 @@
+from concurrent.futures import ThreadPoolExecutor
+from typing import Dict, Generator, Optional
+
+from lumigo_tracer.extension.extension_utils import get_extension_logger, request_event
+
+
+_lambda_service: Optional["LambdaService"] = None
+
+
+class LambdaService:
+    def __init__(self, extension_id):
+        self.extension_id: str = extension_id
+        self.thread: ThreadPoolExecutor = ThreadPoolExecutor()
+        self.next_event_future = None
+
+    def ready_for_next_event(self):
+        if not self.next_event_future:
+            self.next_event_future = self.thread.submit(request_event, self.extension_id)
+
+    def block_until_next_event(self) -> Dict[str, str]:
+        self.ready_for_next_event()
+        result = self.next_event_future.result()
+        self.next_event_future = None
+        return result
+
+    def events_generator(self) -> Generator[Dict[str, str], None, None]:
+        try:
+            while True:
+                yield self.block_until_next_event()
+        except Exception:
+            get_extension_logger().exception("Extension:: failed retrieving events")

src/lumigo_tracer/extension/lambda_service.py

@@ -0,0 +1,73 @@
+import json
+import os
+import http.client
+
+from lumigo_tracer.extension.extension_utils import get_extension_logger
+from lumigo_tracer.extension.extension import LumigoExtension
+from lumigo_tracer.extension.lambda_service import LambdaService
+from lumigo_tracer.utils import is_kill_switch_on, lumigo_safe_execute, config
+
+STOP_EXTENSION_KEY = "LUMIGO_EXTENSION_STOP"
+LUMIGO_EXTENSION_NAME = "lumigo"
+
+REQUESTED_ENVS = [
+    "LUMIGO_DEBUG",
+    "LUMIGO_TRACER_TOKEN",
+    "LUMIGO_TRACER_HOST",
+    "LUMIGO_EXTENSION_ON",
+    "LUMIGO_SWITCH_OFF",
+]
+EVENTS = ["INVOKE", "SHUTDOWN"]
+
+
+def register():
+    body = json.dumps({"events": EVENTS})
+    headers = {"Lambda-Extension-Name": LUMIGO_EXTENSION_NAME}
+    conn = http.client.HTTPConnection(os.environ["AWS_LAMBDA_RUNTIME_API"])
+    conn.request("POST", "/2020-01-01/extension/register", body, headers=headers)
+    response = conn.getresponse()
+    response.read()
+
+    return response.headers["Lambda-Extension-Identifier"]
+
+
+def start_extension_loop(lambda_service: LambdaService):
+    with lumigo_safe_execute("Extension main initialization"):
+        get_extension_logger().debug(
+            f"Extension started running with extension id: {lambda_service.extension_id}"
+        )
+        config()
+        extension = LumigoExtension(lambda_service)
+    for event in lambda_service.events_generator():
+        get_extension_logger().debug(f"Extension got event: {event}")
+        if event.get("eventType") == "INVOKE":
+            with lumigo_safe_execute("Extension: start new invocation"):
+                extension.start_new_invocation(event)
+        elif event.get("eventType") == "SHUTDOWN":
+            with lumigo_safe_execute("Extension: shutdown"):
+                extension.shutdown()
+            break
+        else:
+            get_extension_logger().error(f"Extension got unknown event: {event}")
+    get_extension_logger().debug("Extension finished running")
+
+
+def start_empty_extension_loop(lambda_service: LambdaService):
+    get_extension_logger().debug("Lumigo extension is disables")
+    for event in lambda_service.events_generator():
+        if event.get("eventType") == "SHUTDOWN":
+            break
+
+
+def main():
+    extension_id = register()
+
+    lambda_service = LambdaService(extension_id)
+    if is_kill_switch_on() or (os.environ.get(STOP_EXTENSION_KEY, "").lower() == "true"):
+        start_empty_extension_loop(lambda_service)
+    else:
+        start_extension_loop(lambda_service)
+
+
+if __name__ == "__main__":
+    main()

src/lumigo_tracer/extension/main.py

@@ -0,0 +1,60 @@
+from datetime import datetime
+
+from typing import Optional, List, Union, Dict
+
+from dataclasses import dataclass
+
+import signal
+
+from lumigo_tracer.extension.extension_utils import get_current_cpu_time
+
+DEFAULT_SAMPLING_INTERVAL = 500
+
+
+@dataclass
+class CpuSample:
+    start_time: datetime
+    end_time: datetime
+    cpu_time: float
+
+    def dump(self) -> Dict[str, Union[float, int]]:
+        return {
+            "start_time": int(self.start_time.timestamp() * 1000),
+            "end_time": int(self.end_time.timestamp() * 1000),
+            "cpu_time": self.cpu_time,
+        }
+
+
+class Sampler:
+    def __init__(self):
+        self.cpu_last_sample_value: Optional[float] = None
+        self.cpu_last_sample_time: Optional[datetime] = None
+        self.cpu_samples: List[CpuSample] = []
+
+    def start_sampling(self, interval_ms: int = DEFAULT_SAMPLING_INTERVAL):
+        self.cpu_samples = []
+        self.sample()
+        signal.signal(signal.SIGALRM, self.sample)
+        signal.setitimer(signal.ITIMER_REAL, interval_ms / 1000, interval_ms / 1000)
+
+    def stop_sampling(self):
+        signal.alarm(0)
+        signal.signal(signal.SIGALRM, signal.SIG_DFL)
+        self.sample()
+
+    def get_samples(self) -> List[CpuSample]:
+        return self.cpu_samples
+
+    def sample(self, *args):
+        now = datetime.now()
+        current_cpu = get_current_cpu_time()
+        if self.cpu_last_sample_time and self.cpu_last_sample_value and current_cpu:
+            self.cpu_samples.append(
+                CpuSample(
+                    start_time=self.cpu_last_sample_time,
+                    end_time=now,
+                    cpu_time=current_cpu - self.cpu_last_sample_value,
+                )
+            )
+        self.cpu_last_sample_time = now
+        self.cpu_last_sample_value = current_cpu

src/lumigo_tracer/extension/sampler.py

@@ -38,9 +38,7 @@
     "requestParameters",
 ]
 
-S3_BUCKET_KEYS_ORDER = str_to_list(os.environ.get("LUMIGO_S3_BUCKET_KEYS_ORDER", "")) or [
-    "arn",
-]
+S3_BUCKET_KEYS_ORDER = str_to_list(os.environ.get("LUMIGO_S3_BUCKET_KEYS_ORDER", "")) or ["arn"]
 
 S3_OBJECT_KEYS_ORDER = str_to_list(os.environ.get("LUMIGO_S3_OBJECT_KEYS_ORDER", "")) or [
     "key",

src/lumigo_tracer/parsers/event_parser.py

@@ -3,11 +3,9 @@
 import logging
 import http.client
 from io import BytesIO
-import os
 import builtins
 from functools import wraps
 import importlib.util
-import botocore.awsrequest  # noqa: F401
 
 from lumigo_tracer.auto_tag.auto_tag_event import AutoTagEvent
 from lumigo_tracer.libs.wrapt import wrap_function_wrapper
@@ -19,14 +17,14 @@
     lumigo_safe_execute,
     is_aws_environment,
     ensure_str,
+    is_kill_switch_on,
 )
 from lumigo_tracer.spans_container import SpansContainer, TimeoutMechanism
 from lumigo_tracer.parsers.http_data_classes import HttpRequest
 from collections import namedtuple
 
 _BODY_HEADER_SPLITTER = b"\r\n\r\n"
 _FLAGS_HEADER_SPLITTER = b"\r\n"
-_KILL_SWITCH = "LUMIGO_SWITCH_OFF"
 CONTEXT_WRAPPED_BY_LUMIGO_KEY = "_wrapped_by_lumigo"
 MAX_READ_SIZE = 1024
 already_wrapped = False
@@ -198,7 +196,7 @@ def _add_wrap_flag_to_context(*args):
 def _lumigo_tracer(func):
     @wraps(func)
     def lambda_wrapper(*args, **kwargs):
-        if str(os.environ.get(_KILL_SWITCH, "")).lower() == "true":
+        if is_kill_switch_on():
             return func(*args, **kwargs)
 
         if _is_context_already_wrapped(*args):
@@ -330,7 +328,10 @@ def wrap_http_calls():
             wrap_function_wrapper(
                 "http.client", "HTTPConnection.request", _headers_reminder_wrapper
             )
-            wrap_function_wrapper("botocore.awsrequest", "AWSRequest.__init__", _putheader_wrapper)
+            if importlib.util.find_spec("botocore"):
+                wrap_function_wrapper(
+                    "botocore.awsrequest", "AWSRequest.__init__", _putheader_wrapper
+                )
             wrap_function_wrapper("http.client", "HTTPConnection.getresponse", _response_wrapper)
             wrap_function_wrapper("http.client", "HTTPResponse.read", _read_wrapper)
             if importlib.util.find_spec("urllib3"):