Repository: DataDog/dd-trace-py

Issue:
pymongo: Incomplete trace
### Which version of dd-trace-py are you using?

```
$pip3 freeze | grep ddtrace
ddtrace==0.37.1
```

### Which version of the libraries are you using?

```
$pip3 freeze | grep pymongo
pymongo==3.10.1
```

### How can we reproduce your problem?

```
import pymongo
from ddtrace import patch_all; patch_all()

ip = "xxx.xxx.xxx.xxx" # With a local server the problem will be difficult to notice as establishing a TCP connection on localhost is nowhere as slow as establishing it across the Internet
client = pymongo.MongoClient(ip, 27017)
with ddtrace.tracer.trace('QueryMongoMetadata', service='xxx', resource='GetMongoMetadatas'):
    server_info = client.server_info()                                                                     
```

produces

![image](https://user-images.githubusercontent.com/6052976/82158395-e3518980-9887-11ea-9ef1-0848a0f08d98.png)

### What is the result that you get?

In this example, when calling `client.server_info()`, pymongo silently establishes a new TCP connection. It then uses the newly created socket to perform the actual `buildinfo` request. This can be verified by capturing network while running that python snippet.
Only the second request is traced by ddtrace. The establishment of the new TCP connection is not. Still, it can take some time, as in this example. This behavior can lead to misleading traces like this one: in this example, the network calls traced by ddtrace only represent 1/5 of the total time taken by pymongo's public function,`server_info()`.

### What is the result that you expected?

I do not know if this is a known behavior. It could also be a regression (cf. 4yo comment [here](https://github.com/DataDog/dd-trace-py/blob/master/ddtrace/contrib/pymongo/client.py#L56), the internal API may have changed).
In any case, I was wondering if you would accept pull requests for tracing new connections or at least for new connections established while calling a traced, public pymongo function.

It looks like pymongo is handling its sockets with `Pool` instances, attached to `Server` instances. Since ddtrace already traces `Server` instances (https://github.com/DataDog/dd-trace-py/blob/master/ddtrace/contrib/pymongo/client.py#L86), I believe it could be done easily by tracing some functions of `Server`'s pool instances.
Another option could be to trace the [already overridden `get_socket`](https://github.com/DataDog/dd-trace-py/blob/master/ddtrace/contrib/pymongo/client.py#L169) function of the `Server` instance. It looks like this function's role is to give a working socket to pymongo (using the appropriate `Pool`) so it can query the server. It would be faster to trace it. The downside would be that ddtrace would generate a trace both in the case where `get_socket` creates a new socket (good) and in the case where `get_socket` picks an already existing one (bad?).

Would you accept a contribution on that issue?


Discussion:
Hey @LetzNico!

Thanks for the write-up! I'm not sure if this is a regression or not either, pymongo's one of our oldest integrations.

Without looking into pymongo, both of your approaches seem good. The `Pool` approach may give more visibility not just into this particular case, which would be great (also more work). We might want to avoid doing the `get_socket` approach if `get_socket` is called a bunch, which I feel like it could be, otherwise it might be a quick and easy solution.

We'd love to accept a contribution if you're up for it! 😄
Hey @Kyle-Verhoog !

I have started working on that issue. I'm considering the Pool patching approach. Still, I am not familiar with ddtrace internals, particularly `ObjectProxy`. I have been hitting my head on a problem with context managers used by pymongo.
Below is a minimal example.

```
from ddtrace.vendor.wrapt import ObjectProxy
import contextlib

class Pool():
    @contextlib.contextmanager
    def parent_function(self):
        # do some stuff
        yield self.target_function()                                                                                                                                         
        # do some stuff

    def target_function(self):
        print("original target")
        return 1

class TracedPool(ObjectProxy):
    def __init__(self, pool):
        super(TracedPool, self).__init__(pool)

    def target_function(self):
        print("overriden target")
        return 2

pool = Pool()
traced_pool = TracedPool(pool)
# stuff
with traced_pool.parent_function() as blablabla:
    pass
# stuff
```

The function I want to trace is `target_function`, called by `parent_function`. However, since `parent_function` is decorated by a `contextmanager`, its return value is a `contextlib._GeneratorContextManager`.
The annoying part is that the function and the arguments referenced by the `contextlib._GeneratorContextManager` are the original `target_function` and `Pool` objects. Thus, even though `target_function` is overriden in `TracedPool`, the original `target_function` will be executed.

Is it a known issue? I will continue working but if you have some insights, that would be greatly appreciated!
Thanks for starting work on it!

Ahhh I see. Surprisingly, I don't think I've personally come across where we've had to instrument a `@contextlib.contextmanager` and don't think we do it anywhere in the tracer.

I think the advice I'd give is to not go the `ObjectProxy` approach. With our newer integrations we're instead wrapping the methods individually. See the [flask integration](https://github.com/DataDog/dd-trace-py/blob/a6df2696bc906758c027be3257deabce344fb5a9/ddtrace/contrib/flask/patch.py#L56-L77) for an example of this. This might make things a bit more straightforward. Let me know if you have any questions about this approach.

That is strange that it has the original references 🤔. Darn, the context manager decorator might be the reason for that... Might have to think about that one.

Do you have a develop branch you could share? Feel free to open up a draft PR I'd love to take a look. 🙂 

Resolving patch:
diff --git a/ddtrace/contrib/pymongo/__init__.py b/ddtrace/contrib/pymongo/__init__.py
--- a/ddtrace/contrib/pymongo/__init__.py
+++ b/ddtrace/contrib/pymongo/__init__.py
@@ -31,6 +31,5 @@
 
 with require_modules(required_modules) as missing_modules:
     if not missing_modules:
-        from .client import trace_mongo_client
-        from .patch import patch
+        from .patch import trace_mongo_client, patch
         __all__ = ['trace_mongo_client', 'patch']
diff --git a/ddtrace/contrib/pymongo/client.py b/ddtrace/contrib/pymongo/client.py
--- a/ddtrace/contrib/pymongo/client.py
+++ b/ddtrace/contrib/pymongo/client.py
@@ -13,7 +13,6 @@
 from ...ext import SpanTypes, mongo as mongox, net as netx
 from ...internal.logger import get_logger
 from ...settings import config
-from ...utils.deprecation import deprecated
 from .parse import parse_spec, parse_query, parse_msg
 
 # Original Client class
@@ -22,13 +21,6 @@
 log = get_logger(__name__)
 
 
-@deprecated(message='Use patching instead (see the docs).', version='1.0.0')
-def trace_mongo_client(client, tracer, service=mongox.SERVICE):
-    traced_client = TracedMongoClient(client)
-    ddtrace.Pin(service=service, tracer=tracer).onto(traced_client)
-    return traced_client
-
-
 class TracedMongoClient(ObjectProxy):
 
     def __init__(self, client=None, *args, **kwargs):
@@ -137,7 +129,7 @@ def run_operation_with_response(self, sock_info, operation, *args, **kwargs):
             )
 
             if result and result.address:
-                _set_address_tags(span, result.address)
+                set_address_tags(span, result.address)
             return result
         finally:
             span.finish()
@@ -160,7 +152,7 @@ def send_message_with_response(self, operation, *args, **kwargs):
             )
 
             if result and result.address:
-                _set_address_tags(span, result.address)
+                set_address_tags(span, result.address)
             return result
         finally:
             span.finish()
@@ -243,7 +235,7 @@ def __trace(self, cmd):
         )
 
         if self.address:
-            _set_address_tags(s, self.address)
+            set_address_tags(s, self.address)
         return s
 
 
@@ -275,7 +267,7 @@ def normalize_filter(f=None):
         return {}
 
 
-def _set_address_tags(span, address):
+def set_address_tags(span, address):
     # the address is only set after the cursor is done.
     if address:
         span.set_tag(netx.TARGET_HOST, address[0])
diff --git a/ddtrace/contrib/pymongo/patch.py b/ddtrace/contrib/pymongo/patch.py
--- a/ddtrace/contrib/pymongo/patch.py
+++ b/ddtrace/contrib/pymongo/patch.py
@@ -1,14 +1,70 @@
+import contextlib
 import pymongo
 
-from .client import TracedMongoClient
+from ddtrace import Pin
+from ddtrace.vendor.wrapt import wrap_function_wrapper as _w
+from ...utils.wrappers import unwrap as _u
+
+from ...ext import mongo as mongox, SpanTypes
+from ...utils.deprecation import deprecated
+from ...constants import SPAN_MEASURED_KEY
+from .client import TracedMongoClient, set_address_tags
 
 # Original Client class
 _MongoClient = pymongo.MongoClient
 
 
 def patch():
+    patch_pymongo_module()
+    # We should progressively get rid of TracedMongoClient. We now try to
+    # wrap methods individually. cf #1501
     setattr(pymongo, 'MongoClient', TracedMongoClient)
 
 
 def unpatch():
+    unpatch_pymongo_module()
     setattr(pymongo, 'MongoClient', _MongoClient)
+
+
+@deprecated(message='Use patching instead (see the docs).', version='1.0.0')
+def trace_mongo_client(client, tracer, service=mongox.SERVICE):
+    traced_client = TracedMongoClient(client)
+    Pin(service=service, tracer=tracer).onto(traced_client)
+    return traced_client
+
+
+def patch_pymongo_module():
+    if getattr(pymongo, '_datadog_patch', False):
+        return
+    setattr(pymongo, '_datadog_patch', True)
+    Pin(service=mongox.SERVICE, app=mongox.SERVICE).onto(pymongo.server.Server)
+
+    # Whenever a pymongo command is invoked, the lib either:
+    # - Creates a new socket & performs a TCP handshake
+    # - Grabs a socket already initialized before
+    _w("pymongo.server", "Server.get_socket", traced_get_socket)
+
+
+def unpatch_pymongo_module():
+    if not getattr(pymongo, '_datadog_patch', False):
+        return
+    setattr(pymongo, '_datadog_patch', False)
+
+    _u(pymongo.server.Server, "get_socket")
+
+
+@contextlib.contextmanager
+def traced_get_socket(wrapped, instance, args, kwargs):
+    pin = Pin._find(wrapped, instance)
+    if not pin or not pin.enabled():
+        with wrapped(*args, **kwargs) as sock_info:
+            yield sock_info
+            return
+
+    with pin.tracer.trace('pymongo.get_socket',
+                          service=pin.service,
+                          span_type=SpanTypes.MONGODB) as span:
+        with wrapped(*args, **kwargs) as sock_info:
+            set_address_tags(span, sock_info.address)
+            span.set_tag(SPAN_MEASURED_KEY)
+            yield sock_info


Test patch:
diff --git a/tests/contrib/pymongo/test.py b/tests/contrib/pymongo/test.py
--- a/tests/contrib/pymongo/test.py
+++ b/tests/contrib/pymongo/test.py
@@ -7,9 +7,9 @@
 # project
 from ddtrace import Pin
 from ddtrace.constants import ANALYTICS_SAMPLE_RATE_KEY
-from ddtrace.ext import mongo as mongox
-from ddtrace.contrib.pymongo.client import trace_mongo_client, normalize_filter
-from ddtrace.contrib.pymongo.patch import patch, unpatch
+from ddtrace.ext import mongo as mongox, SpanTypes
+from ddtrace.contrib.pymongo.client import normalize_filter
+from ddtrace.contrib.pymongo.patch import patch, unpatch, trace_mongo_client
 
 # testing
 from tests.opentracer.utils import init_tracer
@@ -350,6 +350,8 @@ def get_tracer_and_client(self):
         tracer = get_dummy_tracer()
         original_client = pymongo.MongoClient(port=MONGO_CONFIG['port'])
         client = trace_mongo_client(original_client, tracer, service=self.TEST_SERVICE)
+        # No need to disable tcp spans tracer here as trace_mongo_client does not call
+        # patch()
         return tracer, client
 
 
@@ -368,6 +370,8 @@ def get_tracer_and_client(self):
         tracer = get_dummy_tracer()
         client = pymongo.MongoClient(port=MONGO_CONFIG['port'])
         Pin.get_from(client).clone(tracer=tracer).onto(client)
+        # We do not wish to trace tcp spans here
+        Pin.get_from(pymongo.server.Server).remove_from(pymongo.server.Server)
         return tracer, client
 
     def test_service(self):
@@ -409,6 +413,8 @@ def get_tracer_and_client(self):
         tracer = get_dummy_tracer()
         client = pymongo.MongoClient(port=MONGO_CONFIG['port'])
         Pin(service=self.TEST_SERVICE, tracer=tracer).onto(client)
+        # We do not wish to trace tcp spans here
+        Pin.get_from(pymongo.server.Server).remove_from(pymongo.server.Server)
         return tracer, client
 
     def test_patch_unpatch(self):
@@ -421,6 +427,8 @@ def test_patch_unpatch(self):
 
         client = pymongo.MongoClient(port=MONGO_CONFIG['port'])
         Pin.get_from(client).clone(tracer=tracer).onto(client)
+        # We do not wish to trace tcp spans here
+        Pin.get_from(pymongo.server.Server).remove_from(pymongo.server.Server)
         client['testdb'].drop_collection('whatever')
 
         spans = writer.pop()
@@ -441,6 +449,8 @@ def test_patch_unpatch(self):
 
         client = pymongo.MongoClient(port=MONGO_CONFIG['port'])
         Pin.get_from(client).clone(tracer=tracer).onto(client)
+        # We do not wish to trace tcp spans here
+        Pin.get_from(pymongo.server.Server).remove_from(pymongo.server.Server)
         client['testdb'].drop_collection('whatever')
 
         spans = writer.pop()
@@ -460,7 +470,103 @@ def test_user_specified_service(self):
         tracer = get_dummy_tracer()
         client = pymongo.MongoClient(port=MONGO_CONFIG["port"])
         Pin.get_from(client).clone(tracer=tracer).onto(client)
+        # We do not wish to trace tcp spans here
+        Pin.get_from(pymongo.server.Server).remove_from(pymongo.server.Server)
         client["testdb"].drop_collection("whatever")
         spans = tracer.writer.pop()
         assert len(spans) == 1
         assert spans[0].service != "mysvc"
+
+
+class TestPymongoSocketTracing(TracerTestCase):
+    """
+    Test suite which checks that tcp socket creation/retrieval is correctly traced
+    """
+    def setUp(self):
+        patch()
+        self.tracer = get_dummy_tracer()
+        # Override server pin's tracer with our dummy tracer
+        Pin.override(pymongo.server.Server, tracer=self.tracer)
+        # maxPoolSize controls the number of sockets that the client can instanciate
+        # and choose from to perform classic operations. For the sake of our tests,
+        # let's limit this number to 1
+        self.client = pymongo.MongoClient(port=MONGO_CONFIG['port'], maxPoolSize=1)
+        # Override TracedMongoClient's pin's tracer with our dummy tracer
+        Pin.override(self.client, tracer=self.tracer)
+
+    def tearDown(self):
+        unpatch()
+        self.client.close()
+
+    @staticmethod
+    def check_socket_metadata(span):
+        assert span.name == "pymongo.get_socket"
+        assert span.service == mongox.SERVICE
+        assert span.span_type == SpanTypes.MONGODB.value
+        assert span.meta.get('out.host') == "localhost"
+        assert span.metrics.get('out.port') == MONGO_CONFIG["port"]
+
+    def test_single_op(self):
+        self.client["some_db"].drop_collection("some_collection")
+        spans = self.tracer.writer.pop()
+
+        assert len(spans) == 2
+        self.check_socket_metadata(spans[0])
+        assert spans[1].name == "pymongo.cmd"
+
+    def test_multiple_ops(self):
+        db = self.client["medias"]
+        input_movies = [{'name': 'Kool movie', 'filmmaker': 'John Mc Johnson'}]
+
+        # Perform a few pymongo operations
+        db.drop_collection('movies')
+        db.movies.insert_many(input_movies)
+        docs_in_collection = list(db.movies.find())
+        db.drop_collection('movies')
+
+        # Ensure patched pymongo's insert/find behave normally
+        assert input_movies == docs_in_collection
+
+        spans = self.tracer.writer.pop()
+        # Check that we generated one get_socket span and one cmd span per cmd
+        # drop(), insert_many(), find() drop_collection() -> four cmds
+        count_commands_executed = 4
+        # one get_socket() per cmd executed, four cmds executed -> height spans total
+        count_expected_spans_emitted = 2 * count_commands_executed
+        assert len(spans) == count_expected_spans_emitted
+
+        for i in range(0, count_expected_spans_emitted, 2):
+            # Each cmd span should be associated with exactly one get_socket span
+            first_span, second_span = spans[i], spans[i + 1]
+            # We have no guarantee on which span comes first (= which is the parent of which).
+            # When in theory, get_socket() should always be first, pymongo has multiple flows
+            # depending on the operation it will perform against the server.
+            #
+            # Write operations (e.g. `insert_many()`) follow a flow which calls get_socket()
+            # first and then sock_info.command() which we trace correctly.
+            # Then, some flows, namely the flows which necessit a response from the server,
+            # we do not patch correctly. For example, for pymongo 3.0 through 3.8, we trace
+            # `find()`-like operations by patching Server's send_message_with_response(). This
+            # function actually calls get_socket(), then sock_info.send_message(data) and
+            # finally sock_info.receive_message(request_id). The correct way to patch those
+            # three functions and generate three traces, one per network call. Still, that
+            # would be lots of additional work to retrieve the span tags in send_message() and
+            # receive_message().
+            #
+            # Now, with pymongo 3.9+, send_message_with_response() has been replaced with
+            # run_operation_with_response() which takes as an argument a sock_info. The
+            # lib now first calls get_socket() and then run_operation_with_response(sock_info),
+            # which makes more sense and also allows us to trace the function correctly.
+            assert {first_span.name, second_span.name} == {"pymongo.cmd", "pymongo.get_socket"}
+            if first_span.name == "pymongo.get_socket":
+                self.check_socket_metadata(first_span)
+            else:
+                self.check_socket_metadata(second_span)
+
+    def test_server_info(self):
+        self.client.server_info()
+        spans = self.tracer.writer.pop()
+
+        assert len(spans) == 2
+        self.check_socket_metadata(spans[0])
+        assert spans[1].name == "pymongo.cmd"