Correlate cli to service (#3137)

* Start event retention policy * . * Correlate telemetry from cli to service and out * Traces end to end * Linting * . * Fix build failures * Trying to fix python dependency error * . * Lets let pip figure it out * . * Modified the wrong file * . * . * . * . * . * . * This is the one * fix lints? * I _love_ python * ... * Undo some unnecessary changes * Works again * PR comments
microsoft · Jun 28, 2023 · 85c3e10 · 85c3e10
1 parent b3fa826
commit 85c3e10
Show file tree

Hide file tree

Showing 14 changed files with 131 additions and 36 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -148,7 +148,7 @@ jobs:
           python setup.py sdist bdist_wheel
           pip install -r ./requirements.txt ../../artifacts/sdk/*.whl
           pip install six
-          pyinstaller onefuzz/__main__.py --onefile --name onefuzz --additional-hooks-dir extra/pyinstaller --hidden-import='pkg_resources.py2_warn' --exclude-module tkinter --exclude-module PySide2 --exclude-module PIL.ImageDraw --exclude-module Pillow --clean --add-data "onefuzz/data/privacy.txt;onefuzz/data" --add-data "onefuzz/data/licenses.json;onefuzz/data"
+          pyinstaller onefuzz/__main__.py --onefile --name onefuzz --additional-hooks-dir extra/pyinstaller --hidden-import='pkg_resources.py2_warn' --hidden-import='opentelemetry.baggage' --hidden-import='opentelemetry.baggage.propagation' --hidden-import='opentelemetry.context.contextvars_context' --copy-metadata opentelemetry-sdk --copy-metadata opentelemetry-api --exclude-module tkinter --exclude-module PySide2 --exclude-module PIL.ImageDraw --exclude-module Pillow --clean --add-data "onefuzz/data/privacy.txt;onefuzz/data" --add-data "onefuzz/data/licenses.json;onefuzz/data"
           ./dist/onefuzz.exe --version
           ./dist/onefuzz.exe privacy_statement
           mkdir -p ${GITHUB_WORKSPACE}/artifacts/windows-cli/
@@ -172,7 +172,7 @@ jobs:
           isort --profile black ./onefuzz ./examples/ ./tests/ --check
           pytest -v tests
           ../ci/disable-py-cache.sh
-          mypy --ignore-missing-imports ./onefuzz ./examples ./tests
+          mypy --ignore-missing-imports --implicit-reexport --namespace-packages ./onefuzz ./examples ./tests
 
           # set a minimum confidence to ignore known false positives
           vulture --min-confidence 61 onefuzz

diff --git a/src/ApiService/ApiService/ApiService.csproj b/src/ApiService/ApiService/ApiService.csproj
@@ -51,6 +51,8 @@
     <PackageReference Include="Octokit" Version="2.0.1" />
     <PackageReference Include="Microsoft.TeamFoundationServer.Client" Version="19.219.0-preview" />
     <PackageReference Include="SmartAnalyzers.CSharpExtensions.Annotations" Version="4.2.7" />
+    <PackageReference Include="OpenTelemetry.Api" Version="1.5.0-rc.1" />
+    <PackageReference Include="System.Diagnostics.DiagnosticSource" Version="8.0.0-preview.4.23259.5" />
   </ItemGroup>
   <ItemGroup>
     <None Update="host.json">

diff --git a/src/ApiService/ApiService/Functions/Node.cs b/src/ApiService/ApiService/Functions/Node.cs
@@ -174,6 +174,7 @@ private async Async.Task<HttpResponseData> Delete(HttpRequestData req) {
             }
         }
 
+
         return await RequestHandling.Ok(req, true);
     }
 }
diff --git a/src/ApiService/ApiService/Log.cs b/src/ApiService/ApiService/Log.cs
@@ -26,6 +26,8 @@ public record TelemetryConfig(TelemetryClient TelemetryClient, ISet<Telemetry>?
 public class OneFuzzLogger : ILogger {
 
     public const string CorrelationId = "CorrelationId";
+    public const string TraceId = "TraceId";
+    public const string SpanId = "SpanId";
 
     private readonly string categoryName;
 

diff --git a/src/ApiService/ApiService/Program.cs b/src/ApiService/ApiService/Program.cs
@@ -9,7 +9,6 @@
 using Azure.Identity;
 using Microsoft.ApplicationInsights.DependencyCollector;
 using Microsoft.Azure.Functions.Worker;
-using Microsoft.Azure.Functions.Worker.Http;
 using Microsoft.Azure.Functions.Worker.Middleware;
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.DependencyInjection;
@@ -36,27 +35,26 @@ public class LoggingMiddleware : IFunctionsWorkerMiddleware {
         public async Async.Task Invoke(FunctionContext context, FunctionExecutionDelegate next) {
             //https://learn.microsoft.com/en-us/azure/azure-monitor/app/custom-operations-tracking#applicationinsights-operations-vs-systemdiagnosticsactivity
             using var activity = OneFuzzLogger.Activity;
-            _ = activity.Start();
-            string correlationId = Guid.NewGuid().ToString();
-
-            if (await context.GetHttpRequestDataAsync() is HttpRequestData requestData) {
-                //if header has 1f-CorrelationId then use that
-                //otherwise check if message can be deserialized to {"correlationId": "SOME-GUID"}, then use that
 
-                if (requestData.Headers.TryGetValues("Correlation-ID", out var values1f)) {
-                    correlationId = values1f.First();
-                } else if (requestData.Headers.TryGetValues("X-Correlation-ID", out var values)) {
-                    correlationId = values.First();
-                }
+            // let azure functions identify the headers for us
+            if (context.TraceContext is not null && !string.IsNullOrEmpty(context.TraceContext.TraceParent)) {
+                activity.TraceStateString = context.TraceContext.TraceState;
+                _ = activity.SetParentId(context.TraceContext.TraceParent);
             }
 
-            _ = activity.AddTag(OneFuzzLogger.CorrelationId, correlationId);
+            _ = activity.Start();
+
+            _ = activity.AddTag(OneFuzzLogger.CorrelationId, activity.TraceId);
+            _ = activity.AddTag(OneFuzzLogger.TraceId, activity.TraceId);
+            _ = activity.AddTag(OneFuzzLogger.SpanId, activity.SpanId);
             _ = activity.AddTag("FunctionId", context.FunctionId);
             _ = activity.AddTag("InvocationId", context.InvocationId);
 
             await next(context);
 
-            context.GetHttpResponseData()?.Headers.Add("X-Correlation-ID", correlationId);
+            var response = context.GetHttpResponseData();
+
+            response?.Headers.Add("traceparent", activity.Id);
         }
     }
 

diff --git a/src/ApiService/ApiService/packages.lock.json b/src/ApiService/ApiService/packages.lock.json
@@ -341,6 +341,15 @@
         "resolved": "2.0.1",
         "contentHash": "JVlfUY+sfItl6RSyVKDJTutuy28cDydUwKKfzcelwNMor2Sa18pYVKna6phO8lug1b+ep+pcuFh/FPayuImsQw=="
       },
+      "OpenTelemetry.Api": {
+        "type": "Direct",
+        "requested": "[1.5.0-rc.1, )",
+        "resolved": "1.5.0-rc.1",
+        "contentHash": "KQYeO/UgV1WoaxOh9l4qV15T1FVow7Aflz1yplSgdnJPqji62hD9+hCe2tQ001iJ9l/mv1/x0xfVEZlRfwDaUA==",
+        "dependencies": {
+          "System.Diagnostics.DiagnosticSource": "7.0.0"
+        }
+      },
       "Scriban": {
         "type": "Direct",
         "requested": "[5.5.0, )",
@@ -359,6 +368,12 @@
         "resolved": "4.2.7",
         "contentHash": "9fRFxTUwPmH7lukckwEvvKawMcP8ObwnOngN8kx5Bx773WHSku1EGa5BIteV07th5553il76fPX7U1xz2bFmuQ=="
       },
+      "System.Diagnostics.DiagnosticSource": {
+        "type": "Direct",
+        "requested": "[8.0.0-preview.4.23259.5, )",
+        "resolved": "8.0.0-preview.4.23259.5",
+        "contentHash": "p7bkTixVHLOALaNlqX+8uBMT68Fbo0yoEigLQ6ya9wNRPThGx5V14et7kf9pVQbR9gN1H2zqSPbX0y8nrXTq+g=="
+      },
       "System.IdentityModel.Tokens.Jwt": {
         "type": "Direct",
         "requested": "[6.29.0, )",
@@ -1423,11 +1438,6 @@
           "System.Runtime": "4.3.0"
         }
       },
-      "System.Diagnostics.DiagnosticSource": {
-        "type": "Transitive",
-        "resolved": "7.0.0",
-        "contentHash": "9W0ewWDuAyDqS2PigdTxk6jDKonfgscY/hP8hm7VpxYhNHZHKvZTdRckberlFk3VnCmr3xBUyMBut12Q+T2aOw=="
-      },
       "System.Diagnostics.EventLog": {
         "type": "Transitive",
         "resolved": "5.0.0",

diff --git a/src/ApiService/IntegrationTests/packages.lock.json b/src/ApiService/IntegrationTests/packages.lock.json
@@ -1253,6 +1253,14 @@
         "resolved": "2.0.1",
         "contentHash": "JVlfUY+sfItl6RSyVKDJTutuy28cDydUwKKfzcelwNMor2Sa18pYVKna6phO8lug1b+ep+pcuFh/FPayuImsQw=="
       },
+      "OpenTelemetry.Api": {
+        "type": "Transitive",
+        "resolved": "1.5.0-rc.1",
+        "contentHash": "KQYeO/UgV1WoaxOh9l4qV15T1FVow7Aflz1yplSgdnJPqji62hD9+hCe2tQ001iJ9l/mv1/x0xfVEZlRfwDaUA==",
+        "dependencies": {
+          "System.Diagnostics.DiagnosticSource": "7.0.0"
+        }
+      },
       "runtime.debian.8-x64.runtime.native.System.Security.Cryptography.OpenSsl": {
         "type": "Transitive",
         "resolved": "4.3.2",
@@ -1567,8 +1575,8 @@
       },
       "System.Diagnostics.DiagnosticSource": {
         "type": "Transitive",
-        "resolved": "7.0.0",
-        "contentHash": "9W0ewWDuAyDqS2PigdTxk6jDKonfgscY/hP8hm7VpxYhNHZHKvZTdRckberlFk3VnCmr3xBUyMBut12Q+T2aOw=="
+        "resolved": "8.0.0-preview.4.23259.5",
+        "contentHash": "p7bkTixVHLOALaNlqX+8uBMT68Fbo0yoEigLQ6ya9wNRPThGx5V14et7kf9pVQbR9gN1H2zqSPbX0y8nrXTq+g=="
       },
       "System.Diagnostics.EventLog": {
         "type": "Transitive",
@@ -2607,9 +2615,11 @@
           "Microsoft.Identity.Web.TokenCache": "[2.7.0, )",
           "Microsoft.TeamFoundationServer.Client": "[19.219.0-preview, )",
           "Octokit": "[2.0.1, )",
+          "OpenTelemetry.Api": "[1.5.0-rc.1, )",
           "Scriban": "[5.5.0, )",
           "Semver": "[2.1.0, )",
           "SmartAnalyzers.CSharpExtensions.Annotations": "[4.2.7, )",
+          "System.Diagnostics.DiagnosticSource": "[8.0.0-preview.4.23259.5, )",
           "System.IdentityModel.Tokens.Jwt": "[6.29.0, )",
           "System.Linq.Async": "[6.0.1, )",
           "System.Security.Cryptography.Pkcs": "[7.0.2, )",

diff --git a/src/ApiService/Tests/packages.lock.json b/src/ApiService/Tests/packages.lock.json
@@ -1255,6 +1255,14 @@
         "resolved": "2.0.1",
         "contentHash": "JVlfUY+sfItl6RSyVKDJTutuy28cDydUwKKfzcelwNMor2Sa18pYVKna6phO8lug1b+ep+pcuFh/FPayuImsQw=="
       },
+      "OpenTelemetry.Api": {
+        "type": "Transitive",
+        "resolved": "1.5.0-rc.1",
+        "contentHash": "KQYeO/UgV1WoaxOh9l4qV15T1FVow7Aflz1yplSgdnJPqji62hD9+hCe2tQ001iJ9l/mv1/x0xfVEZlRfwDaUA==",
+        "dependencies": {
+          "System.Diagnostics.DiagnosticSource": "7.0.0"
+        }
+      },
       "runtime.debian.8-x64.runtime.native.System.Security.Cryptography.OpenSsl": {
         "type": "Transitive",
         "resolved": "4.3.2",
@@ -1569,8 +1577,8 @@
       },
       "System.Diagnostics.DiagnosticSource": {
         "type": "Transitive",
-        "resolved": "7.0.0",
-        "contentHash": "9W0ewWDuAyDqS2PigdTxk6jDKonfgscY/hP8hm7VpxYhNHZHKvZTdRckberlFk3VnCmr3xBUyMBut12Q+T2aOw=="
+        "resolved": "8.0.0-preview.4.23259.5",
+        "contentHash": "p7bkTixVHLOALaNlqX+8uBMT68Fbo0yoEigLQ6ya9wNRPThGx5V14et7kf9pVQbR9gN1H2zqSPbX0y8nrXTq+g=="
       },
       "System.Diagnostics.EventLog": {
         "type": "Transitive",
@@ -2609,9 +2617,11 @@
           "Microsoft.Identity.Web.TokenCache": "[2.7.0, )",
           "Microsoft.TeamFoundationServer.Client": "[19.219.0-preview, )",
           "Octokit": "[2.0.1, )",
+          "OpenTelemetry.Api": "[1.5.0-rc.1, )",
           "Scriban": "[5.5.0, )",
           "Semver": "[2.1.0, )",
           "SmartAnalyzers.CSharpExtensions.Annotations": "[4.2.7, )",
+          "System.Diagnostics.DiagnosticSource": "[8.0.0-preview.4.23259.5, )",
           "System.IdentityModel.Tokens.Jwt": "[6.29.0, )",
           "System.Linq.Async": "[6.0.1, )",
           "System.Security.Cryptography.Pkcs": "[7.0.2, )",

diff --git a/src/ci/build_cli.ps1 b/src/ci/build_cli.ps1
@@ -31,7 +31,7 @@ try {
     pip install -r .\requirements.txt -r .\requirements-dev.txt
 
     # Build exe
-    pyinstaller onefuzz/__main__.py --onefile --name "onefuzz" --additional-hooks-dir extra/pyinstaller --hidden-import='pkg_resources.py2_warn' --exclude-module tkinter --exclude-module PySide2 --exclude-module PIL.ImageDraw --exclude-module Pillow --clean
+    pyinstaller onefuzz/__main__.py --onefile --name "onefuzz" --additional-hooks-dir extra/pyinstaller --hidden-import='pkg_resources.py2_warn' --hidden-import='opentelemetry.baggage' --hidden-import='opentelemetry.baggage.propagation' --hidden-import='opentelemetry.context.contextvars_context' --copy-metadata opentelemetry-sdk --copy-metadata opentelemetry-api --exclude-module tkinter --exclude-module PySide2 --exclude-module PIL.ImageDraw --exclude-module Pillow --clean
 
     # Cleanup
     (Get-Content -path "requirements.txt") -replace "./onefuzztypes-$_version-py3-none-any.whl", "onefuzztypes==$version" | Out-File -FilePath "requirements.txt" -Encoding "ascii"

diff --git a/src/cli/onefuzz/__main__.py b/src/cli/onefuzz/__main__.py
@@ -7,13 +7,27 @@
 
 import sys
 
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+
 from onefuzz.__version__ import __version__
 from onefuzz.api import Command, Endpoint, Onefuzz
 from onefuzz.cli import execute_api
+from onefuzz.opentelemetry_exporter import OneFuzzSpanExporter
+
+provider = TracerProvider()
+processor = SimpleSpanProcessor(OneFuzzSpanExporter())
+provider.add_span_processor(processor)
+trace.set_tracer_provider(provider)
+
+
+tracer = trace.get_tracer(__name__)
 
 
 def main() -> int:
-    return execute_api(Onefuzz(), [Endpoint, Command], __version__)
+    with tracer.start_as_current_span("cli"):
+        return execute_api(Onefuzz(), [Endpoint, Command], __version__)
 
 
 if __name__ == "__main__":

diff --git a/src/cli/onefuzz/backend.py b/src/cli/onefuzz/backend.py
@@ -32,6 +32,8 @@
 import requests
 from azure.storage.blob import ContainerClient
 from onefuzztypes import responses
+from opentelemetry import context, propagate
+from opentelemetry.instrumentation.requests import RequestsInstrumentor
 from pydantic import BaseModel
 from requests import Response
 from tenacity import RetryCallState, retry
@@ -53,6 +55,7 @@
 REQUEST_READ_TIMEOUT = 120.0
 
 LOGGER = logging.getLogger("backend")
+PROPAGATOR = propagate.get_global_textmap()
 
 
 @contextlib.contextmanager
@@ -120,6 +123,7 @@ def __init__(
         token_path: Optional[str] = None,
         client_secret: Optional[str] = None,
     ):
+        RequestsInstrumentor().instrument()
         self.config_path = os.path.expanduser(config_path or DEFAULT_CONFIG_PATH)
         self.token_path = os.path.expanduser(token_path or DEFAULT_TOKEN_PATH)
         self.client_secret = client_secret
@@ -382,14 +386,28 @@ def request(
         response = None
         for backoff in range(1, 10):
             try:
+                current_context = context.get_current()
                 LOGGER.debug("request %s %s %s", method, url, repr(json_data))
-                response = self.session.request(
-                    method,
-                    url,
-                    headers=headers,
-                    json=json_data,
-                    params=params,
-                    timeout=(REQUEST_CONNECT_TIMEOUT, REQUEST_READ_TIMEOUT),
+
+                request_to_downstream = requests.Request(
+                    method, url, headers=headers, json=json_data, params=params
+                )
+
+                PROPAGATOR.inject(
+                    carrier=request_to_downstream.headers,
+                    context=current_context,
+                )
+
+                correlation_id = ""
+                for span in current_context.values():
+                    correlation_id = str(hex(span.__dict__["_context"].trace_id))[2:]
+                    break
+
+                LOGGER.debug("OneFuzz CorrelationId: %s", correlation_id)
+                prep_req = self.session.prepare_request(request_to_downstream)
+
+                response = self.session.send(
+                    prep_req, timeout=(REQUEST_CONNECT_TIMEOUT, REQUEST_READ_TIMEOUT)
                 )
 
                 if response.status_code not in retry_codes:

diff --git a/src/cli/onefuzz/opentelemetry_exporter.py b/src/cli/onefuzz/opentelemetry_exporter.py
@@ -0,0 +1,27 @@
+import logging
+import os
+import typing
+
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+
+LOGGER = logging.getLogger("opentelemetry")
+
+
+class OneFuzzSpanExporter(SpanExporter):
+    def __init__(
+        self,
+        formatter: typing.Callable[[ReadableSpan], str] = lambda span: str(
+            span.to_json()
+        )
+        + os.linesep,
+    ):
+        self.formatter = formatter
+
+    def export(self, spans: typing.Sequence[ReadableSpan]) -> SpanExportResult:
+        for span in spans:
+            LOGGER.debug(self.formatter(span))
+        return SpanExportResult.SUCCESS
+
+    def force_flush(self, _timeout_millis: int = 30000) -> bool:
+        return True
diff --git a/src/cli/requirements-lint.txt b/src/cli/requirements-lint.txt
@@ -3,7 +3,7 @@ mypy==0.910
 pytest
 isort
 vulture
-click==8.0.4
+click>=8.0.4
 black
 bandit
 types-requests

diff --git a/src/cli/requirements.txt b/src/cli/requirements.txt
@@ -13,8 +13,11 @@ azure-applicationinsights==0.1.0
 tenacity==8.0.1
 docstring_parser==0.8.1
 azure-identity==1.10.0
-azure-cli-core==2.47.0
+azure-cli-core==2.49.0
 PyJWT>=2.4.0
+opentelemetry-api==1.16.0
+opentelemetry-sdk==1.16.0
+opentelemetry-instrumentation-requests==0.37b0
 # install rsa version >=4.7 to fix CVE-2020-25658
 rsa>=4.7
 # onefuzztypes version is set during build
-Original file line number
+Diff line change
@@ Expand Up @@
                 }
             }
             return await RequestHandling.Ok(req, true);
         }
     }