Add option --x-trace-all for forcing distribute tracing (#1973)

CHANGELOG.D/1973.feature

@@ -0,0 +1 @@
+Added private option `--x-trace-all` which forces distribute tracing.

neuro-cli/src/neuro_cli/main.py

@@ -181,6 +181,7 @@ def make_context(
             network_timeout=kwargs["network_timeout"],
             config_path=Path(kwargs["neuromation_config"]),
             trace=kwargs["trace"],
+            force_trace_all=kwargs["x_trace_all"],
             trace_hide_token=hide_token_bool,
             command_path="",
             command_params=[],
@@ -383,6 +384,12 @@ def print_options(
     is_flag=True,
     help="Trace sent HTTP requests and received replies to stderr.",
 )
+@option(
+    "--x-trace-all",
+    hidden=True,
+    is_flag=True,
+    help="Force distribute tracing in all HTTP requests.",
+)
 @option(
     "--hide-token/--no-hide-token",
     is_flag=True,
@@ -415,6 +422,7 @@ def cli(
     disable_pypi_version_check: bool,
     network_timeout: float,
     trace: bool,
+    x_trace_all: bool,
     hide_token: Optional[bool],
     skip_stats: bool,
 ) -> None:

neuro-cli/src/neuro_cli/root.py

@@ -68,6 +68,7 @@ class Root:
     network_timeout: float
     config_path: Path
     trace: bool
+    force_trace_all: bool
     verbosity: int
     trace_hide_token: bool
     command_path: str
@@ -174,6 +175,7 @@ def factory(self) -> Factory:
                 path=self.config_path,
                 trace_configs=trace_configs,
                 trace_id=gen_trace_id(),
+                trace_sampled=True if self.force_trace_all else None,
             )
         return self._factory
 

neuro-cli/tests/unit/conftest.py

@@ -69,6 +69,7 @@ def create_root(config_path: Path) -> Root:
         verbosity=0,
         trace=False,
         trace_hide_token=True,
+        force_trace_all=False,
         command_path="",
         command_params=[],
         skip_gmp_stats=True,

neuro-cli/tests/unit/test_config.py

@@ -52,6 +52,7 @@ def test_prompt_cluster(make_client: Callable[..., Client]) -> None:
         verbosity=0,
         trace=False,
         trace_hide_token=True,
+        force_trace_all=False,
         command_path="",
         command_params=[],
         skip_gmp_stats=True,
@@ -112,6 +113,7 @@ def test_prompt_cluster_default(make_client: Callable[..., Client]) -> None:
         verbosity=0,
         trace=False,
         trace_hide_token=True,
+        force_trace_all=False,
         command_path="",
         command_params=[],
         skip_gmp_stats=True,

neuro-cli/tests/unit/test_root.py

@@ -18,6 +18,7 @@ def root_uninitialized() -> Iterator[Root]:
         verbosity=0,
         trace=False,
         trace_hide_token=True,
+        force_trace_all=False,
         command_path="",
         command_params=[],
         skip_gmp_stats=True,

neuro-cli/tests/unit/test_storage_progress.py

@@ -76,6 +76,7 @@ def make(color: bool, tty: bool, verbose: bool) -> Root:
             verbosity=int(verbose),
             trace=False,
             trace_hide_token=True,
+            force_trace_all=False,
             command_path="",
             command_params=[],
             skip_gmp_stats=True,

neuro-sdk/src/neuro_sdk/client.py

@@ -26,11 +26,11 @@ def __init__(
         session: aiohttp.ClientSession,
         path: Path,
         trace_id: Optional[str],
+        trace_sampled: Optional[bool] = None,
     ) -> None:
         self._closed = False
-        self._trace_id = trace_id
         self._session = session
-        self._core = _Core(session, trace_id)
+        self._core = _Core(session, trace_id, trace_sampled)
         self._config = Config._create(self._core, path)
 
         # Order does matter, need to check the main config before loading

neuro-sdk/src/neuro_sdk/config_factory.py

@@ -56,6 +56,7 @@ def __init__(
         path: Optional[Path] = None,
         trace_configs: Optional[List[aiohttp.TraceConfig]] = None,
         trace_id: Optional[str] = None,
+        trace_sampled: Optional[bool] = None,
     ) -> None:
         if path is None:
             path = Path(os.environ.get(CONFIG_ENV_NAME, DEFAULT_CONFIG_PATH))
@@ -64,6 +65,7 @@ def __init__(
         if trace_configs:
             self._trace_configs += trace_configs
         self._trace_id = trace_id
+        self._trace_sampled = trace_sampled
 
     @property
     def path(self) -> Path:
@@ -78,7 +80,9 @@ async def get(self, *, timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT) -> Clie
             await self.login_with_passed_config(timeout=timeout)
         session = await _make_session(timeout, self._trace_configs)
         try:
-            client = Client._create(session, self._path, self._trace_id)
+            client = Client._create(
+                session, self._path, self._trace_id, self._trace_sampled
+            )
             await client.config.check_server()
         except (asyncio.CancelledError, Exception):
             await session.close()

neuro-sdk/src/neuro_sdk/core.py

@@ -63,9 +63,11 @@ def __init__(
         self,
         session: aiohttp.ClientSession,
         trace_id: Optional[str],
+        trace_sampled: Optional[bool] = None,
     ) -> None:
         self._session = session
         self._trace_id = trace_id
+        self._trace_sampled = trace_sampled
         self._exception_map = {
             400: IllegalArgumentError,
             401: AuthenticationError,
@@ -131,6 +133,7 @@ async def request(
         if trace_id is None:
             trace_id = gen_trace_id()
         trace_request_ctx.trace_id = trace_id
+        trace_request_ctx.trace_sampled = self._trace_sampled
         if params:
             url = url.with_query(params)
         async with self._session.request(

neuro-sdk/src/neuro_sdk/tracing.py

@@ -4,6 +4,7 @@
 import random
 import time
 import types
+from typing import Optional
 
 import aiohttp
 from multidict import CIMultiDict
@@ -32,10 +33,22 @@ def _gen_span_id() -> str:
     return os.urandom(8).hex()
 
 
-def _update_headers(headers: "CIMultiDict[str]", trace_id: str, span_id: str) -> None:
+def _update_headers(
+    headers: "CIMultiDict[str]",
+    trace_id: str,
+    span_id: str,
+    sampled: Optional[bool] = None,
+) -> None:
     """Creates dict with zipkin single header format."""
     # b3={TraceId}-{SpanId}-{SamplingState}-{ParentSpanId}
     headers["b3"] = f"{trace_id}-{span_id}"
+    if sampled is True:
+        sampled_str = "1"
+    elif sampled is False:
+        sampled_str = "0"
+    else:
+        sampled_str = ""
+    headers["sentry-trace"] = f"{trace_id}-{span_id}-{sampled_str}"
 
 
 async def _on_request_start(
@@ -47,8 +60,9 @@ async def _on_request_start(
     trace_id = getattr(trace_ctx, "trace_id", None)
     if trace_id is None:
         return
+    sampled = getattr(trace_ctx, "trace_sampled", None)
     span_id = _gen_span_id()
-    _update_headers(params.headers, trace_id, span_id)
+    _update_headers(params.headers, trace_id, span_id, sampled)
 
 
 def _make_trace_config() -> aiohttp.TraceConfig: