Pymongo trace tcp connections

[LetzNico]

cf. #1442

[Kyle-Verhoog]

Looks great after a first pass! Just need a few tests to confirm.

To address your questions from the issue thread:

What could be a good approach to test the new trace?

What we typically do is we replace the default tracer instance with our own dummy tracer which will capture all the spans so that we can access them for testing. Looking at the pymongo test suite it's a bit of a mess... I'd recommend to create a separate test case class in tests/contrib/pymongo/test.py starting with something like:

class TestPymongoPool(BaseTracerTestCase):
    def setUp(self):
        patch()
        self.tracer = get_dummy_tracer()
        Pin.override(pymongo.pool.Pool, tracer=self.tracer)

    def tearDown(self):
        unpatch()

    def test_pool(self):
        # do a pool operation which will generate traces
        spans = self.tracer.writer.pop()
        # assert on the spans

Are the tags correct?

They look good to me after a first pass. Will look closer on the next pass :)

Is the way I instanciate the pin correct to?

Yep this looks good to me!

To make the old tests pass, I'm considering dropping from the spans retrieved by the writer all "tcp connect" spans. That would mean, for example, adding a list comprehension like spans = [span for span in spans if span.name in ("mongo_op", "pymongo.cmd")] after this line. Do you think it is a good idea? If not, can you recommend me another approach?

Hmm yeah I was thinking about that as I was going through the test suite. Yes, I think this is a good approach as any to get around the design of the current tests. (As an aside we should certainly rewrite the test suite for pymongo as part of the rewrite 😛)

[LetzNico]

Hey!

Sorry I have not taken the time lately to continue working on that PR. Thanks a lot for your review @Kyle-Verhoog, it has been very helpful.

The PR is pretty much ready, I've just been stuck on a problem. I'm leaving a comment on the PR. If you could help me once again, that would be greatly appreciated!

[LetzNico]

I have discovered that pymongo maintains a background monitoring pool.

The thing is that since we patch a function in the pool class, we get traces from both the monitoring pool (created here) and the active pool (created here) while we only want to retrieve traces from the active pool.

I think we can workaround the problem by performing a check on the instance. For the sake of the example, I've used an implementation detail on the timeout of the pool sockets to detect which pool emitted the span but that is definitely not satisfying.
Do you have an idea to detect which pool is emitting the span? Would patching pymongo's _create_pool_for_server and _create_pool_for_monitor to add an hidden attribute to the pool be an acceptable solution?

[Kyle-Verhoog]

Hey! Sorry for the long delay. Hmmm interesting. After looking at the examples (thanks for the code references!) it does seem like patching _create_pool_for_server would make sense. IMO it'd be better if we're able to avoid patching the undesired path entirely. Is it relatively easy to accomplish?

[LetzNico]

No worries :) Once again, thanks for your time @Kyle-Verhoog

IMO it'd be better if we're able to avoid patching the undesired path entirely. Is it relatively easy to accomplish?

Agree with you but tbh I have not found a way to do it :/ Basically in pseudo-python, the flow is like this:

class PyMongoClient:
    def __init__(self, ip, port):
        self.ip = ip
        self.port = port
        # This runs in a separate thread
        self.start_monitoring()

    def start_monitoring(self):
        while True:
            for monitor in self.monitors:
                monitor.get_socket(self.ip, self.port) as sock_info:
            	    sock_info.health_check()
                time.sleep(x)

    def run_cmd(self, cmd):
        server = self.servers[(self.ip, self.port)]
        with server.get_socket(self.ip, self.port) as sock_info:
            sock_info.send(cmd)

class Server:
    def get_socket(self):
        return self.pool.get_socket()

class Monitor:
    def get_socket(self);
        return self.pool.get_socket()

class Pool:
    def get_socket(self):
        # (...)
        try:
            # Lookup of an already existing, ready socket. This call does not take time
            socket = self.sockets.popleft()
        except IndexError:
            # Creation of a new socket+handshake. This call does take time
            socket = self.connect()
        return socket

So basically:

• If we patch pymongo.server.Server's get_socket, we actually trace both Pool's self.sockets.popleft() and Pool's self.connect(). The second one is effectively what we are trying to trace while the first one is not a network call but a lookup in memory which we do not want to trace (I think?)
• If we patch pymongo.pool.Pool's self.connect(), we only trace actual network calls but we trace network calls of both the Monitor object and the Server object. That means that we generate traces for both the pymongo thread that handle incoming requests (good) but also from the pymongo thread that performs monitoring health_checks on the remote address (bad)

I do not see a way to only patch the critical path, unfortunately :( That leaves us with 2 viable options imo:

• Trace pymongo.pool.Pool's get_socket: it will generate tiny traces when the socket already exists and classic traces when the socket needs to be created
• Do the hack which adds an hidden attribute pool instances that belong to Servers: it will generate traces only for actual network calls linked to pymongo direct requests but well, it's a bit of a hack

[Kyle-Verhoog]

Thanks for that write-up! It's actually really helpful as I'd forgotten some of the context 🙂.

So I'm okay with patching get_socket in general even if there's the chance that it's a quick memory look up. It's meaningful sometimes which is enough reason IMO and our overhead shouldn't be too bad (unless this is done a crazy amount of times). I think we should avoid instrumenting the health checks from Monitor.

Trace pymongo.pool.Pool's get_socket: it will generate tiny traces when the socket already exists and classic traces when the socket needs to be created

Did you mean Server.get_socket here? Else we'd still be getting the health checks? If you did mean Server.get_socket I think that sounds like the way to go 👍. What do you think?

[LetzNico]

Did you mean Server.get_socket here?

Yep my bad, messed things up! Alright I'm changing the code to patch Server.get_socket :). I'll update you when it's done!

[Kyle-Verhoog]

Haha no problem. And great! I'm excited to get this merged 😄

[LetzNico]

Hmm I've ran into a problem implementing that. Actually Pool's get_socket's code creates&connects the socket and then yields it, making Server's get_socket's a generator. Thus, we cannot trace it effectively with wrap_function_wrapper since generator functions instantly return a generator and the actual code is executed when the parent function decides to consume it.

I believe a workaround could be to consume the socket from the generator in our patch and then yield the socket but I don't think it is a good idea. Are there some wrappers to work with generators in ddtrace?

[Kyle-Verhoog]

Agh, I see. Hmm we don't have any utils for patching generators unfortunately. I think we're okay to intercept the generator and yield the socket ourselves as you described. We just unfortunately can't make use of yield from to maintain compatibility with Python 2.

What are your thoughts?

[Kyle-Verhoog]

Nice 👍! I think once the monitor pool issue is figured out I think we're good to go 🙂

[LetzNico]

Hey there @Kyle-Verhoog ! I've finally taken the time to finish this PR, sorry for the delay!

Do you mind having a final look? Also, it looks like profile tests are flaky: https://app.circleci.com/pipelines/github/DataDog/dd-trace-py/1818/workflows/5bf2b1c6-52f2-45b1-a39c-5bc0adb8bb46/jobs/291047/steps. Unfortunately it seems like I lack the rights to relaunch those, so the CI won't pass. Do you mind doing it for me 🙏 ? Also, do you want me to create an issue regarding that matter?

[Kyle-Verhoog]

Sweet! I think it's ready in terms of the integration. I think it'd be good to have a few tests (we can use the ones you've already written) checking that the functionality of pymongo is maintained.

Also feel free to format the files with black since we're starting to do that across the project 😄

[LetzNico]

Regarding black, I guess I'm lucky that nothing has to be done (I think? I'm not 100% sure about the cmds I used)

# tox -e black tests/contrib/pymongo/
black installed: appdirs==1.4.4,attrs==19.3.0,black==19.10b0,click==7.1.2,pathspec==0.8.0,regex==2020.7.14,toml==0.10.1,typed-ast==1.4.1
black run-test-pre: PYTHONHASHSEED='1888914511'
black run-test: commands[0] | black --check .
All done! ✨ 🍰 ✨
219 files would be left unchanged.
__________________________________________ summary ___________________________________________
  black: commands succeeded
  congratulations :)
# tox -e black ddtrace/contrib/pymongo/
black installed: appdirs==1.4.4,attrs==19.3.0,black==19.10b0,click==7.1.2,pathspec==0.8.0,regex==2020.7.14,toml==0.10.1,typed-ast==1.4.1
black run-test-pre: PYTHONHASHSEED='1678242913'
black run-test: commands[0] | black --check .
All done! ✨ 🍰 ✨
219 files would be left unchanged.
__________________________________________ summary ___________________________________________
  black: commands succeeded
  congratulations :)

[Kyle-Verhoog]

@LetzNico oh right. pymongo is probably on the ignore list in pyproject.toml. The change might be better suited for the entire integration refactor, up to you if you still want to do it 😄

[LetzNico]

Hey there @Kyle-Verhoog ! I finally took some time to wrap up this PR. Excited to hear back from you and hope it will get merged! 🙌

[Kyle-Verhoog]

hey @LetzNico! Thanks for getting back around to it! I was out on vacation last week but definitely gonna take a look at this, this week once I get caught back up 😄

[Kyle-Verhoog]

🎉 @LetzNico thanks for seeing this to the finish line!! 😄

[LetzNico]

Yeeey! Thanks for your time! Please hit me up in the future for support regarding this feature or if I can be of any help on related topics :)