Store flags on the scope

Author: cmanallen

sentry_sdk/flag_utils.py

@@ -1,23 +1,6 @@
-import itertools
-
-
-class FlagManager:
-    """
-    Right now this is just an interface for the buffer but it might contain
-    thread-local state handling in the future.
-    """
-
-    def __init__(self, capacity):
-        # type: (int) -> None
-        self.buffer = FlagBuffer(capacity)
-
-    def get_flags(self):
-        # type: () -> list[dict]
-        return self.buffer.serialize()
+from copy import copy
 
-    def set_flag(self, flag, result):
-        # type: (str, bool) -> None
-        self.buffer.insert(flag, result)
+import itertools
 
 
 class FlagBuffer:
@@ -32,18 +15,18 @@ def __init__(self, capacity):
     def index(self):
         return self.ip % self.capacity
 
-    def insert(self, flag, result):
-        # type: (str, bool) -> None
-        flag_ = Flag(flag, result)
-
-        if self.ip >= self.capacity:
-            self.buffer[self.index] = flag_
-        else:
-            self.buffer.append(flag_)
+    def clear(self):
+        self.buffer = []
+        self.ip = 0
 
-        self.ip += 1
+    def copy(self):
+        # type: () -> FlagBuffer
+        buffer = FlagBuffer(capacity=self.capacity)
+        buffer.buffer = copy(self.buffer)
+        buffer.ip = self.ip
+        return buffer
 
-    def serialize(self):
+    def get(self):
         # type: () -> list[dict]
         if self.ip >= self.capacity:
             iterator = itertools.chain(
@@ -53,6 +36,17 @@ def serialize(self):
         else:
             return [flag.asdict for flag in self.buffer]
 
+    def set(self, flag, result):
+        # type: (str, bool) -> None
+        flag_ = Flag(flag, result)
+
+        if self.ip >= self.capacity:
+            self.buffer[self.index] = flag_
+        else:
+            self.buffer.append(flag_)
+
+        self.ip += 1
+
 
 class Flag:
     __slots__ = ("flag", "result")

sentry_sdk/integrations/openfeature.py

@@ -1,7 +1,6 @@
 from typing import TYPE_CHECKING
 import sentry_sdk
 
-from sentry_sdk.flag_utils import FlagManager
 from sentry_sdk.integrations import DidNotEnable, Integration
 
 try:
@@ -22,41 +21,28 @@ class OpenFeatureIntegration(Integration):
     thread-local state to be serialized and sent off in the error payload.
     """
 
-    def __init__(self, capacity):
+    def __init__(self):
         # type: (int) -> None
-        self.flag_manager = FlagManager(capacity=capacity)
-
-        # Get or create a new isolation scope and register the integration's
-        # error processing hook on it.
-        scope = sentry_sdk.get_isolation_scope()
+        # Store the error processor on the current scope. If its forked
+        # (i.e. threads are spawned) the callback will be copied to the
+        # new Scope.
+        scope = sentry_sdk.get_current_scope()
         scope.add_error_processor(self.error_processor)
 
-        # This is a globally registered hook (its a list singleton). FlagManager
-        # expects itself to be in a THREAD-LOCAL context. Whatever hooks are
-        # triggered will not be THREAD-LOCAL unless we seed the open feature hook
-        # class with thread-local context.
-        api.add_hooks(hooks=[OpenFeatureHook(self.flag_manager)])
+        # Register the hook within the global openfeature hooks list.
+        api.add_hooks(hooks=[OpenFeatureHook()])
 
     def error_processor(self, event, exc_info):
-        """
-        On error Sentry will call this hook. This needs to serialize the flags
-        from the THREAD-LOCAL context and put the result into the error event.
-        """
-        event["contexts"]["flags"] = {"values": self.flag_manager.get_flags()}
+        event["contexts"]["flags"] = {
+            "values": sentry_sdk.get_current_scope().flags.get()
+        }
         return event
 
 
 class OpenFeatureHook(Hook):
-    """
-    OpenFeature will call the `after` method after each flag evaluation. We need to
-    accept the method call and push the result into our THREAD-LOCAL buffer.
-    """
-
-    def __init__(self, flag_manager):
-        # type: (FlagManager) -> None
-        self.flag_manager = flag_manager
 
     def after(self, hook_context, details, hints) -> None:
         # type: (HookContext, FlagEvaluationDetails, HookHints) -> None
         if isinstance(details.value, bool):
-            self.flag_manager.set_flag(details.flag_key, details.value)
+            flags = sentry_sdk.get_current_scope().flags
+            flags.set(details.flag_key, details.value)

sentry_sdk/scope.py

@@ -11,6 +11,7 @@
 
 from sentry_sdk.attachments import Attachment
 from sentry_sdk.consts import DEFAULT_MAX_BREADCRUMBS, FALSE_VALUES, INSTRUMENTER
+from sentry_sdk.flag_utils import FlagBuffer
 from sentry_sdk.profiler.continuous_profiler import try_autostart_continuous_profiler
 from sentry_sdk.profiler.transaction_profiler import Profile
 from sentry_sdk.session import Session
@@ -192,6 +193,7 @@ class Scope:
         "client",
         "_type",
         "_last_event_id",
+        "_flags",
     )
 
     def __init__(self, ty=None, client=None):
@@ -249,6 +251,10 @@ def __copy__(self):
 
         rv._last_event_id = self._last_event_id
 
+        # TODO: Do I need to target `self.flags`? I don't want to create an instance
+        # before I copy but I guess theres no harm.
+        rv._flags = self._flags.copy()
+
         return rv
 
     @classmethod
@@ -685,6 +691,7 @@ def clear(self):
 
         # self._last_event_id is only applicable to isolation scopes
         self._last_event_id = None  # type: Optional[str]
+        self._flags = None
 
     @_attr_setter
     def level(self, value):
@@ -1546,6 +1553,13 @@ def __repr__(self):
             self._type,
         )
 
+    @property
+    def flags(self):
+        # type: () -> FlagBuffer
+        if self._flags is None:
+            self._flags = FlagBuffer(capacity=50)
+        self._flags
+
 
 @contextmanager
 def new_scope():

tests/test_flag_utils.py

@@ -1,88 +1,43 @@
-# import asyncio
-# import pytest
-# import concurrent.futures as cf
-
-from sentry_sdk.flag_utils import get_flags, set_flag
+from sentry_sdk.flag_utils import FlagBuffer
 
 
 def test_flag_tracking():
     """Assert the ring buffer works."""
-    set_flag("a", True)
-    flags = get_flags()
+    buffer = FlagBuffer(capacity=3)
+    buffer.set("a", True)
+    flags = buffer.get()
     assert len(flags) == 1
     assert flags == [{"flag": "a", "result": True}]
 
-    set_flag("b", True)
-    flags = get_flags()
+    buffer.set("b", True)
+    flags = buffer.get()
     assert len(flags) == 2
     assert flags == [{"flag": "a", "result": True}, {"flag": "b", "result": True}]
 
-    set_flag("c", True)
-    flags = get_flags()
+    buffer.set("c", True)
+    flags = buffer.get()
     assert len(flags) == 3
     assert flags == [
         {"flag": "a", "result": True},
         {"flag": "b", "result": True},
         {"flag": "c", "result": True},
     ]
 
-    set_flag("d", False)
-    flags = get_flags()
+    buffer.set("d", False)
+    flags = buffer.get()
     assert len(flags) == 3
     assert flags == [
         {"flag": "b", "result": True},
         {"flag": "c", "result": True},
         {"flag": "d", "result": False},
     ]
 
-    set_flag("e", False)
-    set_flag("f", False)
-    flags = get_flags()
+    buffer.set("e", False)
+    buffer.set("f", False)
+    flags = buffer.get()
     assert len(flags) == 3
     assert flags == [
         {"flag": "d", "result": False},
         {"flag": "e", "result": False},
         {"flag": "f", "result": False},
     ]
-
-
-# Not applicable right now. Thread-specific testing might be moved to another
-# module depending on who eventually managees it.
-
-
-# def test_flag_manager_asyncio_isolation(i):
-#     """Assert concurrently evaluated flags do not pollute one another."""
-
-#     async def task(chars: str):
-#         for char in chars:
-#             set_flag(char, True)
-#         return [f["flag"] for f in get_flags()]
-
-#     async def runner():
-#         return asyncio.gather(
-#             task("abc"),
-#             task("de"),
-#             task("fghijk"),
-#         )
-
-#     results = asyncio.run(runner()).result()
-
-#     assert results[0] == ["a", "b", "c"]
-#     assert results[1] == ["d", "e"]
-#     assert results[2] == ["i", "j", "k"]
-
-
-# def test_flag_manager_thread_isolation(i):
-#     """Assert concurrently evaluated flags do not pollute one another."""
-
-#     def task(chars: str):
-#         for char in chars:
-#             set_flag(char, True)
-#         return [f["flag"] for f in get_flags()]
-
-#     with cf.ThreadPoolExecutor(max_workers=3) as pool:
-#         results = list(pool.map(task, ["abc", "de", "fghijk"]))
-
-#     assert results[0] == ["a", "b", "c"]
-#     assert results[1] == ["d", "e"]
-#     assert results[2] == ["i", "j", "k"]