diff --git a/CHANGELOG.D/1973.feature b/CHANGELOG.D/1973.feature new file mode 100644 index 000000000..4a2c00ea3 --- /dev/null +++ b/CHANGELOG.D/1973.feature @@ -0,0 +1 @@ +Added private option `--x-trace-all` which forces distribute tracing. diff --git a/neuro-cli/src/neuro_cli/main.py b/neuro-cli/src/neuro_cli/main.py index 67e1c4dcd..2d086f489 100644 --- a/neuro-cli/src/neuro_cli/main.py +++ b/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: diff --git a/neuro-cli/src/neuro_cli/root.py b/neuro-cli/src/neuro_cli/root.py index 917ed5534..fd5a5f0b5 100644 --- a/neuro-cli/src/neuro_cli/root.py +++ b/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 diff --git a/neuro-cli/tests/unit/conftest.py b/neuro-cli/tests/unit/conftest.py index 24865abc6..be0e0b14a 100644 --- a/neuro-cli/tests/unit/conftest.py +++ b/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, diff --git a/neuro-cli/tests/unit/test_config.py b/neuro-cli/tests/unit/test_config.py index ee7c8c04c..9b9aeda58 100644 --- a/neuro-cli/tests/unit/test_config.py +++ b/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, diff --git a/neuro-cli/tests/unit/test_root.py b/neuro-cli/tests/unit/test_root.py index 7467b23c8..d641196a5 100644 --- a/neuro-cli/tests/unit/test_root.py +++ b/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, diff --git a/neuro-cli/tests/unit/test_storage_progress.py b/neuro-cli/tests/unit/test_storage_progress.py index 04ddc6a84..23087a37d 100644 --- a/neuro-cli/tests/unit/test_storage_progress.py +++ b/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, diff --git a/neuro-sdk/src/neuro_sdk/client.py b/neuro-sdk/src/neuro_sdk/client.py index 064b67496..3beb65990 100644 --- a/neuro-sdk/src/neuro_sdk/client.py +++ b/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 diff --git a/neuro-sdk/src/neuro_sdk/config_factory.py b/neuro-sdk/src/neuro_sdk/config_factory.py index d7f657f2c..6280226ec 100644 --- a/neuro-sdk/src/neuro_sdk/config_factory.py +++ b/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() diff --git a/neuro-sdk/src/neuro_sdk/core.py b/neuro-sdk/src/neuro_sdk/core.py index 11c4710c6..30b47f774 100644 --- a/neuro-sdk/src/neuro_sdk/core.py +++ b/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( diff --git a/neuro-sdk/src/neuro_sdk/tracing.py b/neuro-sdk/src/neuro_sdk/tracing.py index 37ae86bb7..d3ae8376a 100644 --- a/neuro-sdk/src/neuro_sdk/tracing.py +++ b/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: