From eb5958dc8fe95db7b8f323c0f941f109c639ded0 Mon Sep 17 00:00:00 2001
From: Alex Boten <aboten@lightstep.com>
Date: Thu, 10 Mar 2022 09:00:46 -0800
Subject: [PATCH] [instrumentation/wsgi] fix NonRecordingSpan bug

There was a bug caused by accessing `.kind` on a NonRecordingSpan. Added a test to validate the fix.

Fix #956
---
 .../instrumentation/wsgi/__init__.py          |  5 +++--
 .../tests/test_wsgi_middleware.py             | 21 +++++++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/instrumentation/opentelemetry-instrumentation-wsgi/src/opentelemetry/instrumentation/wsgi/__init__.py b/instrumentation/opentelemetry-instrumentation-wsgi/src/opentelemetry/instrumentation/wsgi/__init__.py
index ad4d425b25..8531d02c56 100644
--- a/instrumentation/opentelemetry-instrumentation-wsgi/src/opentelemetry/instrumentation/wsgi/__init__.py
+++ b/instrumentation/opentelemetry-instrumentation-wsgi/src/opentelemetry/instrumentation/wsgi/__init__.py
@@ -116,6 +116,7 @@ def response_hook(span: Span, environ: WSGIEnvironment, status: str, response_he
 from opentelemetry.instrumentation.wsgi.version import __version__
 from opentelemetry.propagators.textmap import Getter
 from opentelemetry.semconv.trace import SpanAttributes
+from opentelemetry.trace import NonRecordingSpan
 from opentelemetry.trace.status import Status, StatusCode
 from opentelemetry.util.http import (
     OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST,
@@ -313,7 +314,7 @@ def _create_start_response(span, start_response, response_hook):
         @functools.wraps(start_response)
         def _start_response(status, response_headers, *args, **kwargs):
             add_response_attributes(span, status, response_headers)
-            if span.kind == trace.SpanKind.SERVER:
+            if not isinstance(span, NonRecordingSpan) and span.kind == trace.SpanKind.SERVER:
                 add_custom_response_headers(span, response_headers)
             if response_hook:
                 response_hook(status, response_headers)
@@ -336,7 +337,7 @@ def __call__(self, environ, start_response):
             context_getter=wsgi_getter,
             attributes=collect_request_attributes(environ),
         )
-        if span.kind == trace.SpanKind.SERVER:
+        if not isinstance(span, NonRecordingSpan) and span.kind == trace.SpanKind.SERVER:
             add_custom_request_headers(span, environ)
 
         if self.request_hook:
diff --git a/instrumentation/opentelemetry-instrumentation-wsgi/tests/test_wsgi_middleware.py b/instrumentation/opentelemetry-instrumentation-wsgi/tests/test_wsgi_middleware.py
index 35f8e0577a..945e6f823f 100644
--- a/instrumentation/opentelemetry-instrumentation-wsgi/tests/test_wsgi_middleware.py
+++ b/instrumentation/opentelemetry-instrumentation-wsgi/tests/test_wsgi_middleware.py
@@ -475,6 +475,27 @@ def iterate_response(self, response):
             except StopIteration:
                 break
 
+    @mock.patch.dict(
+        "os.environ",
+        {
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST: "Custom-Test-Header-1,Custom-Test-Header-2,Custom-Test-Header-3",
+        },
+    )
+    def test_custom_request_headers_non_recording_span(self):
+        try:
+            tracer_provider = trace_api.NoOpTracerProvider()
+            self.environ.update(
+                {
+                    "HTTP_CUSTOM_TEST_HEADER_1": "Test Value 2",
+                    "HTTP_CUSTOM_TEST_HEADER_2": "TestValue2,TestValue3",
+                }
+            )
+            app = otel_wsgi.OpenTelemetryMiddleware(simple_wsgi, tracer_provider=tracer_provider)
+            response = app(self.environ, self.start_response)
+            self.iterate_response(response)
+        except Exception as e:
+            self.fail(f"Exception raised with NonRecordingSpan {e}")
+
     @mock.patch.dict(
         "os.environ",
         {