Skip to content

Commit 60aa8cd

Browse files
NoahStappNoahStapp
authored
PYTHON-5963 Only attach a shared operation id to retry attempts when APM/logging is enabled (#2960)
1 parent 63955d6 commit 60aa8cd

5 files changed

Lines changed: 137 additions & 7 deletions

File tree

‎pymongo/_op_id.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414

1515
"""Internal helpers for the APM operation id.
1616
17-
The retryable read/write logic sets OP_ID for the duration of each attempt so
17+
The retryable read/write logic sets OP_ID if APM/logging is enabled for the duration of each attempt so
1818
that every attempt of one logical operation publishes the same operation_id.
1919
Commands run outside that scope (handshake, auth, killCursors, pinned-cursor
2020
getMores) read the default None and fall back to their request_id.

‎pymongo/asynchronous/mongo_client.py‎

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535

3636
import asyncio
3737
import contextlib
38+
import logging
3839
import os
3940
import time as time # noqa: PLC0414 # needed in sync version
4041
import warnings
@@ -91,6 +92,8 @@
9192
)
9293
from pymongo.logger import (
9394
_CLIENT_LOGGER,
95+
_COMMAND_LOGGER,
96+
_SERVER_SELECTION_LOGGER,
9497
_log_client_error,
9598
_log_or_warn,
9699
)
@@ -2790,7 +2793,17 @@ def __init__(
27902793
self._server: Server = None # type: ignore
27912794
self._deprioritized_servers: list[Server] = []
27922795
self._operation = operation
2793-
self._operation_id = operation_id if operation_id is not None else _randint()
2796+
# Only generate an operation id when APM/logging is enabled
2797+
if operation_id is None and (
2798+
(
2799+
self._client._event_listeners is not None
2800+
and self._client._event_listeners.enabled_for_commands
2801+
)
2802+
or _COMMAND_LOGGER.isEnabledFor(logging.DEBUG)
2803+
or _SERVER_SELECTION_LOGGER.isEnabledFor(logging.DEBUG)
2804+
):
2805+
operation_id = _randint()
2806+
self._operation_id = operation_id
27942807
self._attempt_number = 0
27952808
self._is_run_command = is_run_command
27962809
self._is_aggregate_write = is_aggregate_write
@@ -3019,7 +3032,9 @@ async def _write(self) -> T:
30193032
self._retryable = False
30203033
if self._retrying:
30213034
self._log_retry(is_write=True)
3022-
# One operation id across all attempts of this operation.
3035+
# One operation id across all attempts of this operation if APM/logging is enabled
3036+
if self._operation_id is None:
3037+
return await self._func(self._session, conn, self._retryable) # type: ignore
30233038
with _op_id._OpIdContext(self._operation_id):
30243039
return await self._func(self._session, conn, self._retryable) # type: ignore
30253040
except PyMongoError as exc:
@@ -3044,7 +3059,9 @@ async def _read(self) -> T:
30443059
self._check_last_error()
30453060
if self._retrying:
30463061
self._log_retry(is_write=False)
3047-
# One operation id across all attempts of this operation.
3062+
# One operation id across all attempts of this operation if APM/logging is enabled
3063+
if self._operation_id is None:
3064+
return await self._func(self._session, self._server, conn, read_pref) # type: ignore
30483065
with _op_id._OpIdContext(self._operation_id):
30493066
return await self._func(self._session, self._server, conn, read_pref) # type: ignore
30503067

‎pymongo/synchronous/mongo_client.py‎

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535

3636
import asyncio
3737
import contextlib
38+
import logging
3839
import os
3940
import time as time # noqa: PLC0414 # needed in sync version
4041
import warnings
@@ -81,6 +82,8 @@
8182
)
8283
from pymongo.logger import (
8384
_CLIENT_LOGGER,
85+
_COMMAND_LOGGER,
86+
_SERVER_SELECTION_LOGGER,
8487
_log_client_error,
8588
_log_or_warn,
8689
)
@@ -2781,7 +2784,17 @@ def __init__(
27812784
self._server: Server = None # type: ignore
27822785
self._deprioritized_servers: list[Server] = []
27832786
self._operation = operation
2784-
self._operation_id = operation_id if operation_id is not None else _randint()
2787+
# Only generate an operation id when APM/logging is enabled
2788+
if operation_id is None and (
2789+
(
2790+
self._client._event_listeners is not None
2791+
and self._client._event_listeners.enabled_for_commands
2792+
)
2793+
or _COMMAND_LOGGER.isEnabledFor(logging.DEBUG)
2794+
or _SERVER_SELECTION_LOGGER.isEnabledFor(logging.DEBUG)
2795+
):
2796+
operation_id = _randint()
2797+
self._operation_id = operation_id
27852798
self._attempt_number = 0
27862799
self._is_run_command = is_run_command
27872800
self._is_aggregate_write = is_aggregate_write
@@ -3010,7 +3023,9 @@ def _write(self) -> T:
30103023
self._retryable = False
30113024
if self._retrying:
30123025
self._log_retry(is_write=True)
3013-
# One operation id across all attempts of this operation.
3026+
# One operation id across all attempts of this operation if APM/logging is enabled
3027+
if self._operation_id is None:
3028+
return self._func(self._session, conn, self._retryable) # type: ignore
30143029
with _op_id._OpIdContext(self._operation_id):
30153030
return self._func(self._session, conn, self._retryable) # type: ignore
30163031
except PyMongoError as exc:
@@ -3035,7 +3050,9 @@ def _read(self) -> T:
30353050
self._check_last_error()
30363051
if self._retrying:
30373052
self._log_retry(is_write=False)
3038-
# One operation id across all attempts of this operation.
3053+
# One operation id across all attempts of this operation if APM/logging is enabled
3054+
if self._operation_id is None:
3055+
return self._func(self._session, self._server, conn, read_pref) # type: ignore
30393056
with _op_id._OpIdContext(self._operation_id):
30403057
return self._func(self._session, self._server, conn, read_pref) # type: ignore
30413058

‎test/asynchronous/test_operation_id_retry.py‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,18 +16,23 @@
1616

1717
from __future__ import annotations
1818

19+
import logging
1920
import sys
21+
from unittest.mock import patch
2022

2123
sys.path[0:0] = [""]
2224

2325
import pymongo
2426
from bson.codec_options import DEFAULT_CODEC_OPTIONS
2527
from pymongo import _op_id
28+
from pymongo._telemetry import _CommandTelemetry
29+
from pymongo.asynchronous import mongo_client
2630
from pymongo.asynchronous.encryption import _Encrypter
2731
from pymongo.asynchronous.helpers import _handle_reauth
2832
from pymongo.asynchronous.pool import AsyncConnection
2933
from pymongo.errors import OperationFailure
3034
from pymongo.helpers_shared import _REAUTHENTICATION_REQUIRED_CODE
35+
from pymongo.logger import _COMMAND_LOGGER, _SERVER_SELECTION_LOGGER
3136
from pymongo.operations import InsertOne
3237
from test.asynchronous import AsyncIntegrationTest, async_client_context, unittest
3338
from test.utils_shared import AllowListEventListener
@@ -134,6 +139,49 @@ async def test_retryable_reads_reuse_operation_id(self):
134139
with self.subTest(command=name, index=i):
135140
await self._check_stable_operation_id(name, f, self.RETRIES)
136141

142+
async def test_retry_without_listeners_or_logging_creates_no_operation_id(self):
143+
appname = _APP_NAME + "noapm"
144+
client = await self.async_rs_or_single_client(appname=appname)
145+
146+
# Make sure APM and logging are disabled
147+
for logger in (_COMMAND_LOGGER, _SERVER_SELECTION_LOGGER):
148+
self.assertFalse(logger.isEnabledFor(logging.DEBUG))
149+
self.assertFalse(client._event_listeners.enabled_for_commands)
150+
151+
find_op_ids = []
152+
original_init = _CommandTelemetry.__init__
153+
154+
def recording_init(self, topology_id, conn, listeners, cmd, dbname, request_id, op_id):
155+
if next(iter(cmd)) == "find":
156+
find_op_ids.append(op_id)
157+
original_init(self, topology_id, conn, listeners, cmd, dbname, request_id, op_id)
158+
159+
fail_point = {
160+
"mode": {"times": 1},
161+
"data": {
162+
"failCommands": ["find"],
163+
"closeConnection": True,
164+
"appName": appname,
165+
},
166+
}
167+
async with self.fail_point(fail_point):
168+
with (
169+
patch.object(mongo_client, "_randint") as randint,
170+
patch.object(_CommandTelemetry, "__init__", recording_init),
171+
):
172+
self.assertIsNotNone(
173+
await client.pymongo_test.test_operation_id_retry.find_one({"_id": 1})
174+
)
175+
176+
self.assertEqual(
177+
randint.call_count, 0, "generated an operation id without APM/logging enabled"
178+
)
179+
self.assertEqual(
180+
find_op_ids,
181+
[None, None],
182+
"expected two attempts, neither carrying a shared operation id",
183+
)
184+
137185
async def test_reauth_does_not_reuse_operation_id(self):
138186
class FakeConnection(AsyncConnection):
139187
def __init__(self):

‎test/test_operation_id_retry.py‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,16 +16,21 @@
1616

1717
from __future__ import annotations
1818

19+
import logging
1920
import sys
21+
from unittest.mock import patch
2022

2123
sys.path[0:0] = [""]
2224

2325
import pymongo
2426
from bson.codec_options import DEFAULT_CODEC_OPTIONS
2527
from pymongo import _op_id
28+
from pymongo._telemetry import _CommandTelemetry
2629
from pymongo.errors import OperationFailure
2730
from pymongo.helpers_shared import _REAUTHENTICATION_REQUIRED_CODE
31+
from pymongo.logger import _COMMAND_LOGGER, _SERVER_SELECTION_LOGGER
2832
from pymongo.operations import InsertOne
33+
from pymongo.synchronous import mongo_client
2934
from pymongo.synchronous.encryption import _Encrypter
3035
from pymongo.synchronous.helpers import _handle_reauth
3136
from pymongo.synchronous.pool import Connection
@@ -132,6 +137,49 @@ def test_retryable_reads_reuse_operation_id(self):
132137
with self.subTest(command=name, index=i):
133138
self._check_stable_operation_id(name, f, self.RETRIES)
134139

140+
def test_retry_without_listeners_or_logging_creates_no_operation_id(self):
141+
appname = _APP_NAME + "noapm"
142+
client = self.rs_or_single_client(appname=appname)
143+
144+
# Make sure APM and logging are disabled
145+
for logger in (_COMMAND_LOGGER, _SERVER_SELECTION_LOGGER):
146+
self.assertFalse(logger.isEnabledFor(logging.DEBUG))
147+
self.assertFalse(client._event_listeners.enabled_for_commands)
148+
149+
find_op_ids = []
150+
original_init = _CommandTelemetry.__init__
151+
152+
def recording_init(self, topology_id, conn, listeners, cmd, dbname, request_id, op_id):
153+
if next(iter(cmd)) == "find":
154+
find_op_ids.append(op_id)
155+
original_init(self, topology_id, conn, listeners, cmd, dbname, request_id, op_id)
156+
157+
fail_point = {
158+
"mode": {"times": 1},
159+
"data": {
160+
"failCommands": ["find"],
161+
"closeConnection": True,
162+
"appName": appname,
163+
},
164+
}
165+
with self.fail_point(fail_point):
166+
with (
167+
patch.object(mongo_client, "_randint") as randint,
168+
patch.object(_CommandTelemetry, "__init__", recording_init),
169+
):
170+
self.assertIsNotNone(
171+
client.pymongo_test.test_operation_id_retry.find_one({"_id": 1})
172+
)
173+
174+
self.assertEqual(
175+
randint.call_count, 0, "generated an operation id without APM/logging enabled"
176+
)
177+
self.assertEqual(
178+
find_op_ids,
179+
[None, None],
180+
"expected two attempts, neither carrying a shared operation id",
181+
)
182+
135183
def test_reauth_does_not_reuse_operation_id(self):
136184
class FakeConnection(Connection):
137185
def __init__(self):

0 commit comments

Comments
 (0)