refactor(ipc): more general shared file

We refactor the file queue into a more generic shared string file for
IPC that can be used by more components.

Author: P403n1x87

ddtrace/internal/_file_queue.py

@@ -0,0 +1,210 @@
+from contextlib import contextmanager
+import os
+import secrets
+import tempfile
+import typing
+
+from ddtrace.internal._unpatched import unpatched_open
+from ddtrace.internal.compat import Path
+from ddtrace.internal.logger import get_logger
+
+
+log = get_logger(__name__)
+
+
+MAX_FILE_SIZE = 8192
+
+
+class BaseLock:
+    def __init__(self, file: typing.IO[typing.Any]):
+        self.file = file
+
+    def acquire(self):
+        ...
+
+    def release(self):
+        ...
+
+    def __enter__(self):
+        self.acquire()
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        self.release()
+
+
+try:
+    # Unix based file locking
+    # Availability: Unix, not Emscripten, not WASI.
+    import fcntl
+
+    class BaseUnixLock(BaseLock):
+        __acquire_mode__ = None
+
+        def acquire(self):
+            if self.__acquire_mode__ is None:
+                msg = f"Cannot use lock of type {type(self)} directly"
+                raise ValueError(msg)
+
+            fcntl.lockf(self.file, self.__acquire_mode__)
+
+        def release(self):
+            fcntl.lockf(self.file, fcntl.LOCK_UN)
+
+    class ReadLock(BaseUnixLock):
+        __acquire_mode__ = fcntl.LOCK_SH
+
+    class WriteLock(BaseUnixLock):
+        __acquire_mode__ = fcntl.LOCK_EX
+
+    @contextmanager
+    def open_file(path, mode):
+        yield unpatched_open(path, mode)
+
+except ModuleNotFoundError:
+    # Availability: Windows
+    import _winapi
+    import ctypes
+    from ctypes import wintypes
+    import msvcrt
+    import os
+
+    kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
+
+    LOCKFILE_EXCLUSIVE_LOCK = 0x00000002
+    INVALID_HANDLE_VALUE = -1
+
+    class OVERLAPPED(ctypes.Structure):
+        _fields_ = [
+            ("Internal", wintypes.LPVOID),
+            ("InternalHigh", wintypes.LPVOID),
+            ("Offset", wintypes.DWORD),
+            ("OffsetHigh", wintypes.DWORD),
+            ("hEvent", wintypes.HANDLE),
+        ]
+
+    kernel32.LockFileEx.argtypes = (
+        wintypes.HANDLE,
+        wintypes.DWORD,
+        wintypes.DWORD,
+        wintypes.DWORD,
+        wintypes.DWORD,
+        ctypes.POINTER(OVERLAPPED),
+    )
+    kernel32.LockFileEx.restype = wintypes.BOOL
+
+    kernel32.UnlockFileEx.argtypes = (
+        wintypes.HANDLE,
+        wintypes.DWORD,
+        wintypes.DWORD,
+        wintypes.DWORD,
+        ctypes.POINTER(OVERLAPPED),
+    )
+    kernel32.UnlockFileEx.restype = wintypes.BOOL
+
+    SHARED_READ_WRITE = 0x7
+    OPEN_ALWAYS = 4
+    RANDOM_ACCESS = 0x10000000
+
+    class BaseWinLock(BaseLock):
+        __acquire_mode__ = None
+
+        def __init__(self, file: typing.IO[typing.Any]):
+            super().__init__(file)
+            self.handle = None
+            self.overlapped: typing.Optional[OVERLAPPED] = None
+
+        def acquire(self):
+            self.handle = _winapi.CreateFile(
+                self.file.name,
+                _winapi.GENERIC_READ | _winapi.GENERIC_WRITE,
+                SHARED_READ_WRITE,
+                0,
+                OPEN_ALWAYS,
+                RANDOM_ACCESS,
+                0,
+            )
+
+            self.overlapped = OVERLAPPED()
+
+            kernel32.LockFileEx(self.handle, self.__acquire_mode__, 0, 0, 0, ctypes.byref(self.overlapped))
+
+        def release(self):
+            if self.handle is None or self.overlapped is None:
+                msg = "Lock was not acquired"
+                raise RuntimeError(msg)
+
+            try:
+                kernel32.UnlockFileEx(self.handle, 0, 0, 0, ctypes.byref(self.overlapped))
+            finally:
+                _winapi.CloseHandle(self.handle)
+
+    class ReadLock(BaseWinLock):
+        __acquire_mode__ = 0  # Shared by default
+
+    class WriteLock(BaseWinLock):
+        __acquire_mode__ = LOCKFILE_EXCLUSIVE_LOCK
+
+    @contextmanager
+    def open_file(path, mode):
+        # force all modes to be read/write binary
+        mode = "r+b"
+        flag = _winapi.GENERIC_READ | _winapi.GENERIC_WRITE
+        fd_flag = os.O_RDWR | os.O_CREAT | os.O_BINARY | os.O_RANDOM
+
+        handle = _winapi.CreateFile(path, flag, SHARED_READ_WRITE, 0, OPEN_ALWAYS, RANDOM_ACCESS, 0)
+        fd = msvcrt.open_osfhandle(handle, fd_flag | os.O_NOINHERIT)
+
+        yield unpatched_open(fd, mode)
+
+        _winapi.CloseHandle(handle)
+
+
+TMPDIR = Path(tempfile.gettempdir())
+
+
+class SharedStringFile:
+    """A simple shared-file implementation for multiprocess communication."""
+
+    def __init__(self) -> None:
+        self.filename: typing.Optional[str] = str(TMPDIR / secrets.token_hex(8))
+
+    def put(self, data: str) -> None:
+        """Put a string into the file."""
+        if self.filename is None:
+            return
+
+        try:
+            with open_file(self.filename, "ab") as f, WriteLock(f):
+                f.seek(0, os.SEEK_END)
+                dt = (data + "\x00").encode()
+                if f.tell() + len(dt) <= MAX_FILE_SIZE:
+                    f.write(dt)
+        except Exception:  # nosec
+            pass
+
+    def peekall(self) -> typing.List[str]:
+        """Peek at all strings from the file."""
+        if self.filename is None:
+            return []
+
+        try:
+            with open_file(self.filename, "r+b") as f, ReadLock(f):
+                return f.read().strip(b"\x00").decode().split("\x00")
+        except Exception:  # nosec
+            return []
+
+    def snatchall(self) -> typing.List[str]:
+        """Retrieve and remove all strings from the file."""
+        if self.filename is None:
+            return []
+
+        try:
+            with open_file(self.filename, "r+b") as f, WriteLock(f):
+                try:
+                    return f.read().strip(b"\x00").decode().split("\x00")
+                finally:
+                    f.seek(0)
+                    f.truncate()
+        except Exception:  # nosec
+            return []

ddtrace/internal/ipc.py

@@ -510,9 +510,9 @@ def __init__(self):
         self._extra_services_queue = None
         if self._remote_config_enabled and not in_aws_lambda():
             # lazy load slow import
-            from ddtrace.internal._file_queue import File_Queue
+            from ddtrace.internal.ipc import SharedStringFile
 
-            self._extra_services_queue = File_Queue()
+            self._extra_services_queue = SharedStringFile()
 
         self._unparsed_service_mapping = _get_config("DD_SERVICE_MAPPING", "")
         self.service_mapping = parse_tags_str(self._unparsed_service_mapping)
@@ -703,7 +703,7 @@ def _get_extra_services(self):
         # type: () -> set[str]
         if self._extra_services_queue is None:
             return set()
-        self._extra_services.update(self._extra_services_queue.get_all())
+        self._extra_services.update(set(self._extra_services_queue.snatchall()))
         while len(self._extra_services) > 64:
             self._extra_services.pop()
         return self._extra_services