From fcba751969a4fecdcff51fc6e0f4b74f85f6390a Mon Sep 17 00:00:00 2001
From: Alex Boten <aboten@lightstep.com>
Date: Thu, 10 Mar 2022 10:02:49 -0800
Subject: [PATCH] [instrumentation/wsgi] fix NonRecordingSpan bug (#957)

* [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

* fix lint

* use is_recording

* fix lint

* fix lint

* fix lint
---
 .../instrumentation/wsgi/__init__.py          |  4 ++--
 .../tests/test_wsgi_middleware.py             | 23 +++++++++++++++++++
 2 files changed, 25 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 ad4d425b2..57a0c8a4c 100644
--- a/instrumentation/opentelemetry-instrumentation-wsgi/src/opentelemetry/instrumentation/wsgi/__init__.py
+++ b/instrumentation/opentelemetry-instrumentation-wsgi/src/opentelemetry/instrumentation/wsgi/__init__.py
@@ -313,7 +313,7 @@ class OpenTelemetryMiddleware:
         @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 span.is_recording() and span.kind == trace.SpanKind.SERVER:
                 add_custom_response_headers(span, response_headers)
             if response_hook:
                 response_hook(status, response_headers)
@@ -336,7 +336,7 @@ class OpenTelemetryMiddleware:
             context_getter=wsgi_getter,
             attributes=collect_request_attributes(environ),
         )
-        if span.kind == trace.SpanKind.SERVER:
+        if span.is_recording() 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 35f8e0577..323893079 100644
--- a/instrumentation/opentelemetry-instrumentation-wsgi/tests/test_wsgi_middleware.py
+++ b/instrumentation/opentelemetry-instrumentation-wsgi/tests/test_wsgi_middleware.py
@@ -475,6 +475,29 @@ class TestAdditionOfCustomRequestResponseHeaders(WsgiTestBase, TestBase):
             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 exc:  # pylint: disable=W0703
+            self.fail(f"Exception raised with NonRecordingSpan {exc}")
+
     @mock.patch.dict(
         "os.environ",
         {