support trace propagation in httplib integration (#237)

Author: reyang

opencensus/trace/ext/httplib/trace.py

@@ -57,7 +57,7 @@ def wrap_httplib_request(request_func):
     """Wrap the httplib request function to trace. Create a new span and update
     and close the span in the response later.
     """
-    def call(self, method, url, *args, **kwargs):
+    def call(self, method, url, body, headers, *args, **kwargs):
         _tracer = execution_context.get_opencensus_tracer()
         _span = _tracer.start_span()
         _span.name = '[httplib]{}'.format(request_func.__name__)
@@ -71,8 +71,13 @@ def call(self, method, url, *args, **kwargs):
         # Store the current span id to thread local.
         execution_context.set_opencensus_attr(
             'httplib/current_span_id', _span.span_id)
-
-        return request_func(self, method, url, *args, **kwargs)
+        try:
+            headers = headers.copy()
+            headers.update(_tracer.propagator.to_headers(
+                _span.context_tracer.span_context))
+        except Exception:  # pragma: NO COVER
+            pass
+        return request_func(self, method, url, body, headers, *args, **kwargs)
 
     return call
 

tests/unit/trace/ext/httplib/test_httplib_trace.py

@@ -17,6 +17,7 @@
 import mock
 
 from opencensus.trace.ext.httplib import trace
+from opencensus.trace.propagation import trace_context_http_header_format
 
 
 class Test_httplib_trace(unittest.TestCase):
@@ -61,7 +62,10 @@ def test_trace_integration(self):
                          wrap_response_result)
 
     def test_wrap_httplib_request(self):
-        mock_tracer = MockTracer()
+        mock_span = mock.Mock()
+        span_id = '1234'
+        mock_span.span_id = span_id
+        mock_tracer = MockTracer(mock_span)
         mock_request_func = mock.Mock()
         mock_request_func.__name__ = 'request'
 
@@ -72,17 +76,25 @@ def test_wrap_httplib_request(self):
 
         wrapped = trace.wrap_httplib_request(mock_request_func)
 
-        url = 'http://localhost:8080'
+        mock_self = mock.Mock()
         method = 'GET'
+        url = 'http://localhost:8080'
+        body = None
+        headers = {}
 
         with patch:
-            wrapped(mock.Mock(), method, url)
+            wrapped(mock_self, method, url, body, headers)
 
         expected_attributes = {
             '/http/url': url,
             '/http/method': method}
         expected_name = '[httplib]request'
 
+        mock_request_func.assert_called_with(
+            mock_self, method, url, body, {
+                'traceparent': '00-123-456-01',
+            }
+        )
         self.assertEqual(expected_attributes,
                          mock_tracer.span.attributes)
         self.assertEqual(expected_name, mock_tracer.span.name)
@@ -151,13 +163,20 @@ def test_wrap_httplib_response_no_open_span(self):
 class MockTracer(object):
     def __init__(self, span=None):
         self.span = span
+        self.propagator = (
+            trace_context_http_header_format.TraceContextPropagator()
+        )
 
     def current_span(self):
         return self.span
 
     def start_span(self):
         span = mock.Mock()
         span.attributes = {}
+        span.context_tracer = mock.Mock()
+        span.context_tracer.span_context = mock.Mock()
+        span.context_tracer.span_context.trace_id = '123'
+        span.context_tracer.span_context.span_id = '456'
         self.span = span
         return span