Refactor forwarding and logs (#755)

Author: ge0Aja

aws/logs_monitoring/customized_log_group.py

@@ -1,3 +1,9 @@
+# Unless explicitly stated otherwise all files in this repository are licensed
+# under the Apache License Version 2.0.
+# This product includes software developed at Datadog (https://www.datadoghq.com/).
+# Copyright 2021 Datadog, Inc.
+
+
 import re
 
 """

aws/logs_monitoring/forwarders.py renamed to aws/logs_monitoring/forwarder.py

@@ -1,28 +1,33 @@
+# Unless explicitly stated otherwise all files in this repository are licensed
+# under the Apache License Version 2.0.
+# This product includes software developed at Datadog (https://www.datadoghq.com/).
+# Copyright 2021 Datadog, Inc.
+
+
 import logging
 import json
 import os
 
 from telemetry import send_event_metric, send_log_metric
 from trace_forwarder.connection import TraceConnection
-from logs.logs import (
-    DatadogScrubber,
-    DatadogBatcher,
-    DatadogClient,
-    DatadogHTTPClient,
-    DatadogTCPClient,
-)
-from logs.logs_helpers import filter_logs
+from logs.datadog_http_client import DatadogHTTPClient
+from logs.datadog_batcher import DatadogBatcher
+from logs.datadog_client import DatadogClient
+from logs.datadog_tcp_client import DatadogTCPClient
+from logs.datadog_scrubber import DatadogScrubber
+from logs.helpers import filter_logs
 from settings import (
     DD_API_KEY,
     DD_USE_TCP,
     DD_NO_SSL,
     DD_SKIP_SSL_VALIDATION,
     DD_URL,
     DD_PORT,
+    DD_TRACE_INTAKE_URL,
+    DD_FORWARD_LOG,
     SCRUBBING_RULE_CONFIGS,
     INCLUDE_AT_MATCH,
     EXCLUDE_AT_MATCH,
-    DD_TRACE_INTAKE_URL,
 )
 
 logger = logging.getLogger()
@@ -32,7 +37,20 @@
 )
 
 
-def forward_logs(logs):
+def forward(logs, metrics, traces):
+    """
+    Forward logs, metrics, and traces to Datadog in a background thread.
+    """
+    if DD_FORWARD_LOG:
+        _forward_logs(logs)
+
+    _forward_metrics(metrics)
+
+    if len(traces) > 0:
+        _forward_traces(traces)
+
+
+def _forward_logs(logs):
     """Forward logs to Datadog"""
     if logger.isEnabledFor(logging.DEBUG):
         logger.debug(f"Forwarding {len(logs)} logs")
@@ -56,15 +74,15 @@ def forward_logs(logs):
             try:
                 client.send(batch)
             except Exception:
-                logger.exception(f"Exception while forwarding log batch {batch}")
-            else:
+                logger.exception("Exception while forwarding log batch")
+            finally:
                 if logger.isEnabledFor(logging.DEBUG):
                     logger.debug(f"Forwarded log batch: {json.dumps(batch)}")
 
     send_event_metric("logs_forwarded", len(logs_to_forward))
 
 
-def forward_metrics(metrics):
+def _forward_metrics(metrics):
     """
     Forward custom metrics submitted via logs to Datadog in a background thread
     using `lambda_stats` that is provided by the Datadog Python Lambda Layer.
@@ -76,25 +94,23 @@ def forward_metrics(metrics):
         try:
             send_log_metric(metric)
         except Exception:
-            logger.exception(f"Exception while forwarding metric {json.dumps(metric)}")
-        else:
+            logger.exception("Exception while forwarding metrics")
+        finally:
             if logger.isEnabledFor(logging.DEBUG):
                 logger.debug(f"Forwarded metric: {json.dumps(metric)}")
 
     send_event_metric("metrics_forwarded", len(metrics))
 
 
-def forward_traces(trace_payloads):
+def _forward_traces(trace_payloads):
     if logger.isEnabledFor(logging.DEBUG):
         logger.debug(f"Forwarding {len(trace_payloads)} traces")
 
     try:
         trace_connection.send_traces(trace_payloads)
     except Exception:
-        logger.exception(
-            f"Exception while forwarding traces {json.dumps(trace_payloads)}"
-        )
-    else:
+        logger.exception("Exception while forwarding traces")
+    finally:
         if logger.isEnabledFor(logging.DEBUG):
             logger.debug(f"Forwarded traces: {json.dumps(trace_payloads)}")
 

aws/logs_monitoring/lambda_function.py

@@ -17,14 +17,9 @@
 from steps.transformation import transform
 from steps.splitting import split
 from caching.cache_layer import CacheLayer
-from forwarders import (
-    forward_metrics,
-    forward_traces,
-    forward_logs,
-)
+from forwarder import forward
 from settings import (
     DD_API_KEY,
-    DD_FORWARD_LOG,
     DD_SKIP_SSL_VALIDATION,
     DD_API_URL,
     DD_FORWARDER_VERSION,
@@ -78,14 +73,7 @@ def datadog_forwarder(event, context):
     transformed = transform(enriched)
     metrics, logs, trace_payloads = split(transformed)
 
-    if DD_FORWARD_LOG:
-        forward_logs(logs)
-
-    forward_metrics(metrics)
-
-    if len(trace_payloads) > 0:
-        forward_traces(trace_payloads)
-
+    forward(logs, metrics, trace_payloads)
     parse_and_submit_enhanced_metrics(logs, cache_layer)
 
 

aws/logs_monitoring/logs/datadog_batcher.py

@@ -0,0 +1,44 @@
+# Unless explicitly stated otherwise all files in this repository are licensed
+# under the Apache License Version 2.0.
+# This product includes software developed at Datadog (https://www.datadoghq.com/).
+# Copyright 2021 Datadog, Inc.
+
+
+class DatadogBatcher(object):
+    def __init__(self, max_item_size_bytes, max_batch_size_bytes, max_items_count):
+        self._max_item_size_bytes = max_item_size_bytes
+        self._max_batch_size_bytes = max_batch_size_bytes
+        self._max_items_count = max_items_count
+
+    def _sizeof_bytes(self, item):
+        return len(str(item).encode("UTF-8"))
+
+    def batch(self, items):
+        """
+        Returns an array of batches.
+        Each batch contains at most max_items_count items and
+        is not strictly greater than max_batch_size_bytes.
+        All items strictly greater than max_item_size_bytes are dropped.
+        """
+        batches = []
+        batch = []
+        size_bytes = 0
+        size_count = 0
+        for item in items:
+            item_size_bytes = self._sizeof_bytes(item)
+            if size_count > 0 and (
+                size_count >= self._max_items_count
+                or size_bytes + item_size_bytes > self._max_batch_size_bytes
+            ):
+                batches.append(batch)
+                batch = []
+                size_bytes = 0
+                size_count = 0
+            # all items exceeding max_item_size_bytes are dropped here
+            if item_size_bytes <= self._max_item_size_bytes:
+                batch.append(item)
+                size_bytes += item_size_bytes
+                size_count += 1
+        if size_count > 0:
+            batches.append(batch)
+        return batches

aws/logs_monitoring/logs/datadog_client.py

@@ -0,0 +1,37 @@
+# Unless explicitly stated otherwise all files in this repository are licensed
+# under the Apache License Version 2.0.
+# This product includes software developed at Datadog (https://www.datadoghq.com/).
+# Copyright 2021 Datadog, Inc.
+
+
+import time
+from logs.exceptions import RetriableException
+
+
+class DatadogClient(object):
+    """
+    Client that implements a exponential retrying logic to send a batch of logs.
+    """
+
+    def __init__(self, client, max_backoff=30):
+        self._client = client
+        self._max_backoff = max_backoff
+
+    def send(self, logs):
+        backoff = 1
+        while True:
+            try:
+                self._client.send(logs)
+                return
+            except RetriableException:
+                time.sleep(backoff)
+                if backoff < self._max_backoff:
+                    backoff *= 2
+                continue
+
+    def __enter__(self):
+        self._client.__enter__()
+        return self
+
+    def __exit__(self, ex_type, ex_value, traceback):
+        self._client.__exit__(ex_type, ex_value, traceback)

aws/logs_monitoring/logs/datadog_http_client.py

@@ -0,0 +1,94 @@
+# Unless explicitly stated otherwise all files in this repository are licensed
+# under the Apache License Version 2.0.
+# This product includes software developed at Datadog (https://www.datadoghq.com/).
+# Copyright 2021 Datadog, Inc.
+
+
+import os
+import logging
+
+from concurrent.futures import as_completed
+from requests_futures.sessions import FuturesSession
+from logs.helpers import compress_logs
+from logs.exceptions import ScrubbingException
+
+from settings import (
+    DD_USE_COMPRESSION,
+    DD_COMPRESSION_LEVEL,
+    DD_MAX_WORKERS,
+    DD_FORWARDER_VERSION,
+)
+
+logger = logging.getLogger()
+logger.setLevel(logging.getLevelName(os.environ.get("DD_LOG_LEVEL", "INFO").upper()))
+
+
+class DatadogHTTPClient(object):
+    """
+    Client that sends a batch of logs over HTTP.
+    """
+
+    _POST = "POST"
+    if DD_USE_COMPRESSION:
+        _HEADERS = {"Content-type": "application/json", "Content-Encoding": "gzip"}
+    else:
+        _HEADERS = {"Content-type": "application/json"}
+
+    _HEADERS["DD-EVP-ORIGIN"] = "aws_forwarder"
+    _HEADERS["DD-EVP-ORIGIN-VERSION"] = DD_FORWARDER_VERSION
+
+    def __init__(
+        self, host, port, no_ssl, skip_ssl_validation, api_key, scrubber, timeout=10
+    ):
+        self._HEADERS.update({"DD-API-KEY": api_key})
+        protocol = "http" if no_ssl else "https"
+        self._url = "{}://{}:{}/api/v2/logs".format(protocol, host, port)
+        self._scrubber = scrubber
+        self._timeout = timeout
+        self._session = None
+        self._ssl_validation = not skip_ssl_validation
+        self._futures = []
+        if logger.isEnabledFor(logging.DEBUG):
+            logger.debug(
+                f"Initialized http client for logs intake: "
+                f"<host: {host}, port: {port}, url: {self._url}, no_ssl: {no_ssl}, "
+                f"skip_ssl_validation: {skip_ssl_validation}, timeout: {timeout}>"
+            )
+
+    def _connect(self):
+        self._session = FuturesSession(max_workers=DD_MAX_WORKERS)
+        self._session.headers.update(self._HEADERS)
+
+    def _close(self):
+        # Resolve all the futures and log exceptions if any
+        for future in as_completed(self._futures):
+            try:
+                future.result()
+            except Exception:
+                logger.exception("Exception while forwarding logs")
+
+        self._session.close()
+
+    def send(self, logs):
+        """
+        Sends a batch of log, only retry on server and network errors.
+        """
+        try:
+            data = self._scrubber.scrub("[{}]".format(",".join(logs)))
+        except ScrubbingException:
+            raise Exception("could not scrub the payload")
+        if DD_USE_COMPRESSION:
+            data = compress_logs(data, DD_COMPRESSION_LEVEL)
+
+        # FuturesSession returns immediately with a future object
+        future = self._session.post(
+            self._url, data, timeout=self._timeout, verify=self._ssl_validation
+        )
+        self._futures.append(future)
+
+    def __enter__(self):
+        self._connect()
+        return self
+
+    def __exit__(self, ex_type, ex_value, traceback):
+        self._close()

aws/logs_monitoring/logs/datadog_scrubber.py

@@ -0,0 +1,36 @@
+# Unless explicitly stated otherwise all files in this repository are licensed
+# under the Apache License Version 2.0.
+# This product includes software developed at Datadog (https://www.datadoghq.com/).
+# Copyright 2021 Datadog, Inc.
+
+
+import os
+from logs.exceptions import ScrubbingException
+from logs.helpers import compileRegex
+
+
+class DatadogScrubber(object):
+    def __init__(self, configs):
+        rules = []
+        for config in configs:
+            if config.name in os.environ:
+                rules.append(
+                    ScrubbingRule(
+                        compileRegex(config.name, config.pattern), config.placeholder
+                    )
+                )
+        self._rules = rules
+
+    def scrub(self, payload):
+        for rule in self._rules:
+            try:
+                payload = rule.regex.sub(rule.placeholder, payload)
+            except Exception:
+                raise ScrubbingException()
+        return payload
+
+
+class ScrubbingRule(object):
+    def __init__(self, regex, placeholder):
+        self.regex = regex
+        self.placeholder = placeholder