From c926c7ce6b26ba4003c2769cde9de5edfc6570f6 Mon Sep 17 00:00:00 2001
From: Joel Scheuner <joel.scheuner.dev@gmail.com>
Date: Tue, 26 Nov 2024 17:00:59 +0100
Subject: [PATCH] Add UsageMultiSetCounter

---
 .../localstack/utils/analytics/usage.py       | 58 +++++++++++++++++++
 tests/unit/utils/analytics/test_usage.py      | 21 +++++++
 2 files changed, 79 insertions(+)
 create mode 100644 tests/unit/utils/analytics/test_usage.py

diff --git a/localstack-core/localstack/utils/analytics/usage.py b/localstack-core/localstack/utils/analytics/usage.py
index f4b067df8867f..defb2db6b2c86 100644
--- a/localstack-core/localstack/utils/analytics/usage.py
+++ b/localstack-core/localstack/utils/analytics/usage.py
@@ -1,5 +1,6 @@
 import datetime
 import math
+import threading
 from collections import defaultdict
 from itertools import count
 from typing import Any
@@ -49,6 +50,63 @@ def aggregate(self) -> dict:
         return self.state
 
 
+class UsageMultiSetCounter:
+    """
+    Use this counter to count occurrences of unique values for multiple dimensions.
+    This dynamically creates UsageSetCounters and should be used with care (i.e., with limited keys).
+    Example:
+    my_feature_counter = UsageMultiSetCounter("pipes:invocation")
+    my_feature_counter.record("aws:sqs", "aws:lambda")
+    my_feature_counter.record("aws:sqs", "aws:lambda")
+    my_feature_counter.record("aws:sqs", "aws:stepfunctions")
+    my_feature_counter.record("aws:kinesis", "aws:lambda")
+    aggregate is implemented for each counter individually
+    my_feature_counter.aggregate() is available for testing purposes:
+    {
+       "aws:sqs": {
+         "aws:lambda": 2,
+         "aws:stepfunctions": 1,
+       },
+       "aws:kinesis": {
+         "aws:lambda": 1
+       }
+    }
+    """
+
+    namespace: str
+    _counters: dict[str, UsageSetCounter]
+    lock = threading.Lock()
+
+    def __init__(self, namespace: str):
+        self._counters = {}
+        self.namespace = namespace
+
+    def record(self, key: str, value: str):
+        namespace = f"{self.namespace}:{key}"
+        if namespace in self._counters:
+            set_counter = self._counters[namespace]
+        else:
+            with self.lock:
+                if namespace in self._counters:
+                    set_counter = self._counters[namespace]
+                else:
+                    # We cannot use setdefault here because Python always instantiates a new UsageSetCounter,
+                    # which overwrites the collector_registry
+                    set_counter = UsageSetCounter(namespace)
+                    self._counters[namespace] = set_counter
+
+        self._counters[namespace] = set_counter
+        set_counter.record(value)
+
+    def aggregate(self) -> dict:
+        """aggregate is invoked on a per UsageSetCounter level because each counter is registered individually.
+        This utility is only for testing!"""
+        merged_dict = {}
+        for namespace, counter in self._counters.items():
+            merged_dict[namespace] = counter.aggregate()
+        return merged_dict
+
+
 class UsageCounter:
     """
     Use this counter to count numeric values
diff --git a/tests/unit/utils/analytics/test_usage.py b/tests/unit/utils/analytics/test_usage.py
new file mode 100644
index 0000000000000..f58cc24d78077
--- /dev/null
+++ b/tests/unit/utils/analytics/test_usage.py
@@ -0,0 +1,21 @@
+from localstack.utils.analytics.usage import UsageMultiSetCounter
+
+
+def test_multi_set_counter():
+    my_feature_counter = UsageMultiSetCounter("pipes:invocation")
+    my_feature_counter.record("aws:sqs", "aws:lambda")
+    my_feature_counter.record("aws:sqs", "aws:lambda")
+    my_feature_counter.record("aws:sqs", "aws:stepfunctions")
+    my_feature_counter.record("aws:kinesis", "aws:lambda")
+    assert my_feature_counter.aggregate() == {
+        "pipes:invocation:aws:sqs": {
+            "aws:lambda": 2,
+            "aws:stepfunctions": 1,
+        },
+        "pipes:invocation:aws:kinesis": {"aws:lambda": 1},
+    }
+    assert my_feature_counter._counters["pipes:invocation:aws:sqs"].state == {
+        "aws:lambda": 2,
+        "aws:stepfunctions": 1,
+    }
+    assert my_feature_counter._counters["pipes:invocation:aws:kinesis"].state == {"aws:lambda": 1}