Skip to content

Commit 4172644

Browse files
authored
gh-142206: multiprocessing.resource_tracker: Decode messages using older protocol (GH-142215)
1 parent 88cd5d9 commit 4172644

File tree

3 files changed

+73
-20
lines changed

3 files changed

+73
-20
lines changed

Lib/multiprocessing/resource_tracker.py

Lines changed: 45 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,13 @@ def __init__(self):
6868
self._exitcode = None
6969
self._reentrant_messages = deque()
7070

71+
# True to use colon-separated lines, rather than JSON lines,
72+
# for internal communication. (Mainly for testing).
73+
# Filenames not supported by the simple format will always be sent
74+
# using JSON.
75+
# The reader should understand all formats.
76+
self._use_simple_format = False
77+
7178
def _reentrant_call_error(self):
7279
# gh-109629: this happens if an explicit call to the ResourceTracker
7380
# gets interrupted by a garbage collection, invoking a finalizer (*)
@@ -200,7 +207,9 @@ def _launch(self):
200207
os.close(r)
201208

202209
def _make_probe_message(self):
203-
"""Return a JSON-encoded probe message."""
210+
"""Return a probe message."""
211+
if self._use_simple_format:
212+
return b'PROBE:0:noop\n'
204213
return (
205214
json.dumps(
206215
{"cmd": "PROBE", "rtype": "noop"},
@@ -267,6 +276,15 @@ def _write(self, msg):
267276
assert nbytes == len(msg), f"{nbytes=} != {len(msg)=}"
268277

269278
def _send(self, cmd, name, rtype):
279+
if self._use_simple_format and '\n' not in name:
280+
msg = f"{cmd}:{name}:{rtype}\n".encode("ascii")
281+
if len(msg) > 512:
282+
# posix guarantees that writes to a pipe of less than PIPE_BUF
283+
# bytes are atomic, and that PIPE_BUF >= 512
284+
raise ValueError('msg too long')
285+
self._ensure_running_and_write(msg)
286+
return
287+
270288
# POSIX guarantees that writes to a pipe of less than PIPE_BUF (512 on Linux)
271289
# bytes are atomic. Therefore, we want the message to be shorter than 512 bytes.
272290
# POSIX shm_open() and sem_open() require the name, including its leading slash,
@@ -286,6 +304,7 @@ def _send(self, cmd, name, rtype):
286304

287305
# The entire JSON message is guaranteed < PIPE_BUF (512 bytes) by construction.
288306
assert len(msg) <= 512, f"internal error: message too long ({len(msg)} bytes)"
307+
assert msg.startswith(b'{')
289308

290309
self._ensure_running_and_write(msg)
291310

@@ -296,6 +315,30 @@ def _send(self, cmd, name, rtype):
296315
getfd = _resource_tracker.getfd
297316

298317

318+
def _decode_message(line):
319+
if line.startswith(b'{'):
320+
try:
321+
obj = json.loads(line.decode('ascii'))
322+
except Exception as e:
323+
raise ValueError("malformed resource_tracker message: %r" % (line,)) from e
324+
325+
cmd = obj["cmd"]
326+
rtype = obj["rtype"]
327+
b64 = obj.get("base64_name", "")
328+
329+
if not isinstance(cmd, str) or not isinstance(rtype, str) or not isinstance(b64, str):
330+
raise ValueError("malformed resource_tracker fields: %r" % (obj,))
331+
332+
try:
333+
name = base64.urlsafe_b64decode(b64).decode('utf-8', 'surrogateescape')
334+
except ValueError as e:
335+
raise ValueError("malformed resource_tracker base64_name: %r" % (b64,)) from e
336+
else:
337+
cmd, rest = line.strip().decode('ascii').split(':', maxsplit=1)
338+
name, rtype = rest.rsplit(':', maxsplit=1)
339+
return cmd, rtype, name
340+
341+
299342
def main(fd):
300343
'''Run resource tracker.'''
301344
# protect the process from ^C and "killall python" etc
@@ -318,23 +361,7 @@ def main(fd):
318361
with open(fd, 'rb') as f:
319362
for line in f:
320363
try:
321-
try:
322-
obj = json.loads(line.decode('ascii'))
323-
except Exception as e:
324-
raise ValueError("malformed resource_tracker message: %r" % (line,)) from e
325-
326-
cmd = obj["cmd"]
327-
rtype = obj["rtype"]
328-
b64 = obj.get("base64_name", "")
329-
330-
if not isinstance(cmd, str) or not isinstance(rtype, str) or not isinstance(b64, str):
331-
raise ValueError("malformed resource_tracker fields: %r" % (obj,))
332-
333-
try:
334-
name = base64.urlsafe_b64decode(b64).decode('utf-8', 'surrogateescape')
335-
except ValueError as e:
336-
raise ValueError("malformed resource_tracker base64_name: %r" % (b64,)) from e
337-
364+
cmd, rtype, name = _decode_message(line)
338365
cleanup_func = _CLEANUP_FUNCS.get(rtype, None)
339366
if cleanup_func is None:
340367
raise ValueError(

Lib/test/_test_multiprocessing.py

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
from test.support import socket_helper
4040
from test.support import threading_helper
4141
from test.support import warnings_helper
42+
from test.support import subTests
4243
from test.support.script_helper import assert_python_failure, assert_python_ok
4344

4445
# Skip tests if _multiprocessing wasn't built.
@@ -4383,6 +4384,19 @@ def test_copy(self):
43834384
self.assertEqual(bar.z, 2 ** 33)
43844385

43854386

4387+
def resource_tracker_format_subtests(func):
4388+
"""Run given test using both resource tracker communication formats"""
4389+
def _inner(self, *args, **kwargs):
4390+
tracker = resource_tracker._resource_tracker
4391+
for use_simple_format in False, True:
4392+
with (
4393+
self.subTest(use_simple_format=use_simple_format),
4394+
unittest.mock.patch.object(
4395+
tracker, '_use_simple_format', use_simple_format)
4396+
):
4397+
func(self, *args, **kwargs)
4398+
return _inner
4399+
43864400
@unittest.skipUnless(HAS_SHMEM, "requires multiprocessing.shared_memory")
43874401
@hashlib_helper.requires_hashdigest('sha256')
43884402
class _TestSharedMemory(BaseTestCase):
@@ -4662,6 +4676,7 @@ def test_shared_memory_SharedMemoryServer_ignores_sigint(self):
46624676
smm.shutdown()
46634677

46644678
@unittest.skipIf(os.name != "posix", "resource_tracker is posix only")
4679+
@resource_tracker_format_subtests
46654680
def test_shared_memory_SharedMemoryManager_reuses_resource_tracker(self):
46664681
# bpo-36867: test that a SharedMemoryManager uses the
46674682
# same resource_tracker process as its parent.
@@ -4913,6 +4928,7 @@ def test_shared_memory_cleaned_after_process_termination(self):
49134928
"shared_memory objects to clean up at shutdown", err)
49144929

49154930
@unittest.skipIf(os.name != "posix", "resource_tracker is posix only")
4931+
@resource_tracker_format_subtests
49164932
def test_shared_memory_untracking(self):
49174933
# gh-82300: When a separate Python process accesses shared memory
49184934
# with track=False, it must not cause the memory to be deleted
@@ -4940,6 +4956,7 @@ def test_shared_memory_untracking(self):
49404956
mem.close()
49414957

49424958
@unittest.skipIf(os.name != "posix", "resource_tracker is posix only")
4959+
@resource_tracker_format_subtests
49434960
def test_shared_memory_tracking(self):
49444961
# gh-82300: When a separate Python process accesses shared memory
49454962
# with track=True, it must cause the memory to be deleted when
@@ -7353,13 +7370,18 @@ def test_forkpty(self):
73537370

73547371
@unittest.skipUnless(HAS_SHMEM, "requires multiprocessing.shared_memory")
73557372
class TestSharedMemoryNames(unittest.TestCase):
7356-
def test_that_shared_memory_name_with_colons_has_no_resource_tracker_errors(self):
7373+
@subTests('use_simple_format', (True, False))
7374+
def test_that_shared_memory_name_with_colons_has_no_resource_tracker_errors(
7375+
self, use_simple_format):
73577376
# Test script that creates and cleans up shared memory with colon in name
73587377
test_script = textwrap.dedent("""
73597378
import sys
73607379
from multiprocessing import shared_memory
7380+
from multiprocessing import resource_tracker
73617381
import time
73627382
7383+
resource_tracker._resource_tracker._use_simple_format = %s
7384+
73637385
# Test various patterns of colons in names
73647386
test_names = [
73657387
"a:b",
@@ -7387,7 +7409,7 @@ def test_that_shared_memory_name_with_colons_has_no_resource_tracker_errors(self
73877409
sys.exit(1)
73887410
73897411
print("SUCCESS")
7390-
""")
7412+
""" % use_simple_format)
73917413

73927414
rc, out, err = assert_python_ok("-c", test_script)
73937415
self.assertIn(b"SUCCESS", out)
Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
The resource tracker in the :mod:`multiprocessing` module can now understand
2+
messages from older versions of itself. This avoids issues with upgrading
3+
Python while it is running. (Note that such 'in-place' upgrades are not
4+
tested.)

0 commit comments

Comments
 (0)