From ef5acde710d9ecaf7d79b35e64fddc1ac53af9f4 Mon Sep 17 00:00:00 2001 From: Yuxiang Zhu Date: Fri, 15 Nov 2024 16:11:16 +0800 Subject: [PATCH] Fix env var propagation with telemetry enabled When OpenTelemetry is enabled, environment variables `TRACEPARENT` needs to be set for child process for trace context propagation. There is a bug in `cmd_gather_async`: If parameter `env` is not set or set to `None`, all environment variables of the current process will be passed to the child process. However, when injecting `TRACEPARENT`, we pass `{'TRACEPARENT': 'something'}` into child process without other environment variables of the current process, causing misbehavior in the child process. This PR will fix this by copying all environment variables of the current process if `env` is not set. --- artcommon/artcommonlib/exectools.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/artcommon/artcommonlib/exectools.py b/artcommon/artcommonlib/exectools.py index d977f5b3c..72194d66a 100644 --- a/artcommon/artcommonlib/exectools.py +++ b/artcommon/artcommonlib/exectools.py @@ -464,7 +464,6 @@ async def cmd_gather_async(cmd: Union[List[str], str], check: bool = True, **kwa span = trace.get_current_span() span.set_attribute("pyartcd.param.cmd", cmd_list) - logger.info("Executing:cmd_gather_async %s", cmd_list) # capture stdout and stderr if they are not set in kwargs if "stdout" not in kwargs: kwargs["stdout"] = asyncio.subprocess.PIPE @@ -472,13 +471,17 @@ async def cmd_gather_async(cmd: Union[List[str], str], check: bool = True, **kwa kwargs["stderr"] = asyncio.subprocess.PIPE # Propagate trace context to subprocess - env = kwargs.get("env", {}) carrier = {} TraceContextTextMapPropagator().inject(carrier) if "traceparent" in carrier: + env = kwargs.get("env") + if env is None: + # Per Popen doc, a None env means "inheriting the current process’ environment". + # To inject the trace context, we need to copy the current environment. + env = kwargs["env"] = os.environ.copy() env["TRACEPARENT"] = carrier["traceparent"] - kwargs["env"] = env + logger.info("Executing:cmd_gather_async %s", cmd_list) proc = await asyncio.subprocess.create_subprocess_exec(cmd_list[0], *cmd_list[1:], **kwargs) stdout, stderr = await proc.communicate() stdout = stdout.decode() if stdout else "" @@ -511,13 +514,15 @@ async def cmd_assert_async(cmd: Union[List[str], str], check: bool = True, **kwa span.set_attribute("pyartcd.param.cmd", cmd_list) # Propagate trace context to subprocess - env = kwargs.get("env", {}) carrier = {} TraceContextTextMapPropagator().inject(carrier) if "traceparent" in carrier: + env = kwargs.get("env") + if env is None: + # Per Popen doc, a None env means "inheriting the current process’ environment". + # To inject the trace context, we need to copy the current environment. + env = kwargs["env"] = os.environ.copy() env["TRACEPARENT"] = carrier["traceparent"] - logger.warning("Pass TRACEPARENT %s", env["TRACEPARENT"]) - kwargs["env"] = env logger.info("Executing:cmd_assert_async %s", cmd_list) proc = await asyncio.subprocess.create_subprocess_exec(cmd_list[0], *cmd_list[1:], **kwargs)