Fixes

Author: jy-tan

drift/core/adaptive_sampling.py

@@ -5,6 +5,7 @@
 import logging
 import math
 import random
+import threading
 import time
 from dataclasses import dataclass
 from typing import Literal
@@ -75,6 +76,7 @@ def __init__(
         self._config = config
         self._random_fn = random_fn
         self._now_fn = now_fn
+        self._lock = threading.RLock()
 
         self._admission_multiplier = 1.0
         self._state: AdaptiveSamplingState = "fixed" if config.mode == "fixed" else "healthy"
@@ -90,135 +92,140 @@ def __init__(
         self._recent_failure_signal = 0.0
 
     def update(self, snapshot: AdaptiveSamplingHealthSnapshot) -> None:
-        if self._config.mode != "adaptive":
-            self._state = "fixed"
-            self._admission_multiplier = 1.0
-            return
-
-        now_s = self._now_fn()
-        elapsed_s = 2.0 if self._last_updated_at_s == 0 else max(0.001, now_s - self._last_updated_at_s)
-        self._last_updated_at_s = now_s
-
-        decay = math.exp(-(elapsed_s * 1000.0) / 30000.0)
-        self._recent_drop_signal *= decay
-        self._recent_failure_signal *= decay
-
-        dropped_delta = max(0, snapshot.dropped_span_count - self._prev_dropped_span_count)
-        export_failure_delta = max(0, snapshot.export_failure_count - self._prev_export_failure_count)
-
-        self._prev_dropped_span_count = snapshot.dropped_span_count
-        self._prev_export_failure_count = snapshot.export_failure_count
-
-        self._recent_drop_signal += dropped_delta
-        self._recent_failure_signal += export_failure_delta
-
-        if snapshot.queue_fill_ratio is not None:
-            queue_fill_ratio = _clamp01(snapshot.queue_fill_ratio)
-            self._queue_fill_ewma = (
-                queue_fill_ratio
-                if self._queue_fill_ewma is None
-                else (0.25 * queue_fill_ratio) + (0.75 * self._queue_fill_ewma)
+        with self._lock:
+            if self._config.mode != "adaptive":
+                self._state = "fixed"
+                self._admission_multiplier = 1.0
+                return
+
+            now_s = self._now_fn()
+            elapsed_s = 2.0 if self._last_updated_at_s == 0 else max(0.001, now_s - self._last_updated_at_s)
+            self._last_updated_at_s = now_s
+
+            decay = math.exp(-(elapsed_s * 1000.0) / 30000.0)
+            self._recent_drop_signal *= decay
+            self._recent_failure_signal *= decay
+
+            dropped_delta = max(0, snapshot.dropped_span_count - self._prev_dropped_span_count)
+            export_failure_delta = max(0, snapshot.export_failure_count - self._prev_export_failure_count)
+
+            self._prev_dropped_span_count = snapshot.dropped_span_count
+            self._prev_export_failure_count = snapshot.export_failure_count
+
+            self._recent_drop_signal += dropped_delta
+            self._recent_failure_signal += export_failure_delta
+
+            if snapshot.queue_fill_ratio is not None:
+                queue_fill_ratio = _clamp01(snapshot.queue_fill_ratio)
+                self._queue_fill_ewma = (
+                    queue_fill_ratio
+                    if self._queue_fill_ewma is None
+                    else (0.25 * queue_fill_ratio) + (0.75 * self._queue_fill_ewma)
+                )
+
+            queue_pressure = _normalize_between(self._queue_fill_ewma, 0.20, 0.85)
+            memory_pressure = _normalize_between(snapshot.memory_pressure_ratio, 0.80, 0.92)
+            export_failure_pressure = _clamp01(self._recent_failure_signal / 5.0)
+            pressure = max(queue_pressure, memory_pressure, export_failure_pressure)
+
+            hard_brake = (
+                dropped_delta > 0 or snapshot.export_circuit_open or (snapshot.memory_pressure_ratio or 0.0) >= 0.92
             )
 
-        queue_pressure = _normalize_between(self._queue_fill_ewma, 0.20, 0.85)
-        memory_pressure = _normalize_between(snapshot.memory_pressure_ratio, 0.80, 0.92)
-        export_failure_pressure = _clamp01(self._recent_failure_signal / 5.0)
-        pressure = max(queue_pressure, memory_pressure, export_failure_pressure)
+            previous_state = self._state
+            previous_multiplier = self._admission_multiplier
+
+            if hard_brake:
+                self._paused_until_s = now_s + 15.0
+                self._admission_multiplier = 0.0
+                self._state = "critical_pause"
+                self._last_decrease_at_s = now_s
+                self._log_transition(previous_state, previous_multiplier, pressure, snapshot)
+                return
+
+            if now_s < self._paused_until_s:
+                self._state = "critical_pause"
+                self._log_transition(previous_state, previous_multiplier, pressure, snapshot)
+                return
+
+            min_multiplier = self._get_min_multiplier()
+            if pressure >= 0.70:
+                self._admission_multiplier = max(min_multiplier, self._admission_multiplier * 0.4)
+                self._state = "hot"
+                self._last_decrease_at_s = now_s
+            elif pressure >= 0.45:
+                self._admission_multiplier = max(min_multiplier, self._admission_multiplier * 0.7)
+                self._state = "warm"
+                self._last_decrease_at_s = now_s
+            else:
+                if pressure <= 0.20 and (now_s - self._last_decrease_at_s) >= 10.0:
+                    self._admission_multiplier = min(1.0, self._admission_multiplier + 0.05)
+                self._state = "healthy"
 
-        hard_brake = (
-            dropped_delta > 0 or snapshot.export_circuit_open or (snapshot.memory_pressure_ratio or 0.0) >= 0.92
-        )
-
-        previous_state = self._state
-        previous_multiplier = self._admission_multiplier
-
-        if hard_brake:
-            self._paused_until_s = now_s + 15.0
-            self._admission_multiplier = 0.0
-            self._state = "critical_pause"
-            self._last_decrease_at_s = now_s
             self._log_transition(previous_state, previous_multiplier, pressure, snapshot)
-            return
-
-        if now_s < self._paused_until_s:
-            self._state = "critical_pause"
-            self._log_transition(previous_state, previous_multiplier, pressure, snapshot)
-            return
-
-        min_multiplier = self._get_min_multiplier()
-        if pressure >= 0.70:
-            self._admission_multiplier = max(min_multiplier, self._admission_multiplier * 0.4)
-            self._state = "hot"
-            self._last_decrease_at_s = now_s
-        elif pressure >= 0.45:
-            self._admission_multiplier = max(min_multiplier, self._admission_multiplier * 0.7)
-            self._state = "warm"
-            self._last_decrease_at_s = now_s
-        else:
-            if pressure <= 0.20 and (now_s - self._last_decrease_at_s) >= 10.0:
-                self._admission_multiplier = min(1.0, self._admission_multiplier + 0.05)
-            self._state = "healthy"
-
-        self._log_transition(previous_state, previous_multiplier, pressure, snapshot)
 
     def get_decision(self, *, is_pre_app_start: bool) -> RootSamplingDecision:
-        if is_pre_app_start:
-            return RootSamplingDecision(
-                should_record=True,
-                reason="pre_app_start",
-                mode=self._config.mode,
-                state=self._state,
-                base_rate=self._config.base_rate,
-                min_rate=self._config.min_rate,
-                effective_rate=1.0,
-                admission_multiplier=1.0,
+        with self._lock:
+            if is_pre_app_start:
+                return RootSamplingDecision(
+                    should_record=True,
+                    reason="pre_app_start",
+                    mode=self._config.mode,
+                    state=self._state,
+                    base_rate=self._config.base_rate,
+                    min_rate=self._config.min_rate,
+                    effective_rate=1.0,
+                    admission_multiplier=1.0,
+                )
+
+            effective_rate = (
+                self.get_effective_sampling_rate()
+                if self._config.mode == "adaptive"
+                else _clamp01(self._config.base_rate)
             )
 
-        effective_rate = (
-            self.get_effective_sampling_rate() if self._config.mode == "adaptive" else _clamp01(self._config.base_rate)
-        )
-
-        if effective_rate <= 0.0:
+            if effective_rate <= 0.0:
+                return RootSamplingDecision(
+                    should_record=False,
+                    reason="critical_pause" if self._state == "critical_pause" else "not_sampled",
+                    mode=self._config.mode,
+                    state=self._state,
+                    base_rate=self._config.base_rate,
+                    min_rate=self._config.min_rate,
+                    effective_rate=effective_rate,
+                    admission_multiplier=self._admission_multiplier,
+                )
+
+            should_record = self._random_fn() < effective_rate
             return RootSamplingDecision(
-                should_record=False,
-                reason="critical_pause" if self._state == "critical_pause" else "not_sampled",
+                should_record=should_record,
+                reason=(
+                    "sampled"
+                    if should_record
+                    else "load_shed"
+                    if self._config.mode == "adaptive" and effective_rate < self._config.base_rate
+                    else "not_sampled"
+                ),
                 mode=self._config.mode,
                 state=self._state,
                 base_rate=self._config.base_rate,
                 min_rate=self._config.min_rate,
                 effective_rate=effective_rate,
-                admission_multiplier=self._admission_multiplier,
+                admission_multiplier=self._admission_multiplier if self._config.mode == "adaptive" else 1.0,
             )
 
-        should_record = self._random_fn() < effective_rate
-        return RootSamplingDecision(
-            should_record=should_record,
-            reason=(
-                "sampled"
-                if should_record
-                else "load_shed"
-                if self._config.mode == "adaptive" and effective_rate < self._config.base_rate
-                else "not_sampled"
-            ),
-            mode=self._config.mode,
-            state=self._state,
-            base_rate=self._config.base_rate,
-            min_rate=self._config.min_rate,
-            effective_rate=effective_rate,
-            admission_multiplier=self._admission_multiplier if self._config.mode == "adaptive" else 1.0,
-        )
-
     def get_effective_sampling_rate(self) -> float:
-        if self._config.mode != "adaptive":
-            return _clamp01(self._config.base_rate)
-        if self._state == "critical_pause" and self._now_fn() < self._paused_until_s:
-            return 0.0
-        effective_rate = self._config.base_rate * self._admission_multiplier
-        return _clamp(
-            effective_rate,
-            min(self._config.base_rate, self._config.min_rate),
-            self._config.base_rate,
-        )
+        with self._lock:
+            if self._config.mode != "adaptive":
+                return _clamp01(self._config.base_rate)
+            if self._state == "critical_pause" and self._now_fn() < self._paused_until_s:
+                return 0.0
+            effective_rate = self._config.base_rate * self._admission_multiplier
+            return _clamp(
+                effective_rate,
+                min(self._config.base_rate, self._config.min_rate),
+                self._config.base_rate,
+            )
 
     def _get_min_multiplier(self) -> float:
         if self._config.base_rate <= 0.0 or self._config.min_rate <= 0.0:

drift/core/drift_sdk.py

@@ -352,13 +352,14 @@ def _determine_sampling_config(self, init_param: float | None) -> ResolvedSampli
                 config_sampling.mode,
             )
 
-        base_rate = 1.0
+        base_rate: float | None = None
         if init_param is not None:
             validated = validate_sampling_rate(init_param, "init params")
             if validated is not None:
                 logger.debug(f"Using sampling rate from init params: {validated}")
                 base_rate = validated
-        else:
+
+        if base_rate is None:
             env_rate = os.environ.get("TUSK_SAMPLING_RATE")
             if env_rate is not None:
                 try:
@@ -369,20 +370,22 @@ def _determine_sampling_config(self, init_param: float | None) -> ResolvedSampli
                         base_rate = validated
                 except ValueError:
                     logger.warning(f"Invalid TUSK_SAMPLING_RATE env var: {env_rate}")
-            elif config_sampling and config_sampling.base_rate is not None:
-                validated = validate_sampling_rate(
-                    config_sampling.base_rate, "config file recording.sampling.base_rate"
-                )
-                if validated is not None:
-                    base_rate = validated
-            elif recording_config and recording_config.sampling_rate is not None:
-                validated = validate_sampling_rate(
-                    recording_config.sampling_rate, "config file recording.sampling_rate"
-                )
-                if validated is not None:
-                    base_rate = validated
-            else:
-                logger.debug("Using default sampling rate: 1.0")
+
+        if base_rate is None and config_sampling and config_sampling.base_rate is not None:
+            validated = validate_sampling_rate(config_sampling.base_rate, "config file recording.sampling.base_rate")
+            if validated is not None:
+                logger.debug(f"Using sampling rate from config file recording.sampling.base_rate: {validated}")
+                base_rate = validated
+
+        if base_rate is None and recording_config and recording_config.sampling_rate is not None:
+            validated = validate_sampling_rate(recording_config.sampling_rate, "config file recording.sampling_rate")
+            if validated is not None:
+                logger.debug(f"Using sampling rate from config file recording.sampling_rate: {validated}")
+                base_rate = validated
+
+        if base_rate is None:
+            logger.debug("Using default sampling rate: 1.0")
+            base_rate = 1.0
 
         min_rate = 0.0
         if mode == "adaptive":
@@ -423,11 +426,17 @@ def _start_adaptive_sampling_control_loop(self) -> None:
             name="drift-adaptive-sampling",
         )
         self._adaptive_sampling_thread.start()
-        self._update_adaptive_sampling_health()
+        self._safe_update_adaptive_sampling_health()
 
     def _adaptive_sampling_loop(self) -> None:
         while not self._adaptive_sampling_stop_event.wait(timeout=2.0):
+            self._safe_update_adaptive_sampling_health()
+
+    def _safe_update_adaptive_sampling_health(self) -> None:
+        try:
             self._update_adaptive_sampling_health()
+        except Exception:
+            logger.error("Adaptive sampling health update failed; keeping previous controller state.", exc_info=True)
 
     def _update_adaptive_sampling_health(self) -> None:
         if self._adaptive_sampling_controller is None:
@@ -484,14 +493,52 @@ def _get_memory_pressure_ratio(self) -> float | None:
         if cgroup_v1_current is not None:
             return cgroup_v1_current / self._effective_memory_limit_bytes
 
+        current_rss_bytes = self._read_current_rss_bytes()
+        if current_rss_bytes is not None:
+            return current_rss_bytes / self._effective_memory_limit_bytes
+
+        return None
+
+    @staticmethod
+    def _parse_proc_status_rss_bytes(raw_status: str) -> int | None:
+        for line in raw_status.splitlines():
+            if not line.startswith("VmRSS:"):
+                continue
+
+            parts = line.split()
+            if len(parts) < 3 or parts[2].lower() != "kb":
+                return None
+
+            return int(parts[1]) * 1024
+
+        return None
+
+    @staticmethod
+    def _parse_proc_statm_rss_bytes(raw_statm: str, page_size: int) -> int | None:
+        fields = raw_statm.split()
+        if len(fields) < 2:
+            return None
+
+        return int(fields[1]) * page_size
+
+    def _read_current_rss_bytes(self) -> int | None:
         try:
-            import resource
+            proc_status_path = Path("/proc/self/status")
+            if proc_status_path.exists():
+                parsed = self._parse_proc_status_rss_bytes(proc_status_path.read_text())
+                if parsed is not None:
+                    return parsed
+        except Exception:
+            pass
 
-            rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
-            rss_bytes = rss if platform.system() == "Darwin" else rss * 1024
-            return rss_bytes / self._effective_memory_limit_bytes
+        try:
+            proc_statm_path = Path("/proc/self/statm")
+            if proc_statm_path.exists():
+                return self._parse_proc_statm_rss_bytes(proc_statm_path.read_text(), int(os.sysconf("SC_PAGE_SIZE")))
         except Exception:
-            return None
+            pass
+
+        return None
 
     def _read_numeric_control_file(self, path: str) -> int | None:
         try: