bpo-34572: change _pickle unpickling to use import rather than retrie… · python/cpython@4371c0a (original) (raw)
`@@ -3,18 +3,22 @@
`
3
3
`import dbm
`
4
4
`import io
`
5
5
`import functools
`
``
6
`+
import os
`
6
7
`import pickle
`
7
8
`import pickletools
`
``
9
`+
import shutil
`
8
10
`import struct
`
9
11
`import sys
`
``
12
`+
import threading
`
10
13
`import unittest
`
11
14
`import weakref
`
``
15
`+
from textwrap import dedent
`
12
16
`from http.cookies import SimpleCookie
`
13
17
``
14
18
`from test import support
`
15
19
`from test.support import (
`
16
20
`TestFailed, TESTFN, run_with_locale, no_tracing,
`
17
``
`-
_2G, _4G, bigmemtest,
`
``
21
`+
_2G, _4G, bigmemtest, reap_threads, forget,
`
18
22
` )
`
19
23
``
20
24
`from pickle import bytes_types
`
`@@ -1174,6 +1178,67 @@ def test_truncated_data(self):
`
1174
1178
`for p in badpickles:
`
1175
1179
`self.check_unpickling_error(self.truncated_errors, p)
`
1176
1180
``
``
1181
`+
@reap_threads
`
``
1182
`+
def test_unpickle_module_race(self):
`
``
1183
`+
https://bugs.python.org/issue34572
`
``
1184
`+
locker_module = dedent("""
`
``
1185
`+
import threading
`
``
1186
`+
barrier = threading.Barrier(2)
`
``
1187
`+
""")
`
``
1188
`+
locking_import_module = dedent("""
`
``
1189
`+
import locker
`
``
1190
`+
locker.barrier.wait()
`
``
1191
`+
class ToBeUnpickled(object):
`
``
1192
`+
pass
`
``
1193
`+
""")
`
``
1194
+
``
1195
`+
os.mkdir(TESTFN)
`
``
1196
`+
self.addCleanup(shutil.rmtree, TESTFN)
`
``
1197
`+
sys.path.insert(0, TESTFN)
`
``
1198
`+
self.addCleanup(sys.path.remove, TESTFN)
`
``
1199
`+
with open(os.path.join(TESTFN, "locker.py"), "wb") as f:
`
``
1200
`+
f.write(locker_module.encode('utf-8'))
`
``
1201
`+
with open(os.path.join(TESTFN, "locking_import.py"), "wb") as f:
`
``
1202
`+
f.write(locking_import_module.encode('utf-8'))
`
``
1203
`+
self.addCleanup(forget, "locker")
`
``
1204
`+
self.addCleanup(forget, "locking_import")
`
``
1205
+
``
1206
`+
import locker
`
``
1207
+
``
1208
`+
pickle_bytes = (
`
``
1209
`+
b'\x80\x03clocking_import\nToBeUnpickled\nq\x00)\x81q\x01.')
`
``
1210
+
``
1211
`+
Then try to unpickle two of these simultaneously
`
``
1212
`+
One of them will cause the module import, and we want it to block
`
``
1213
`+
until the other one either:
`
``
1214
`+
- fails (before the patch for this issue)
`
``
1215
`+
- blocks on the import lock for the module, as it should
`
``
1216
`+
results = []
`
``
1217
`+
barrier = threading.Barrier(3)
`
``
1218
`+
def t():
`
``
1219
`+
This ensures the threads have all started
`
``
1220
`+
presumably barrier release is faster than thread startup
`
``
1221
`+
barrier.wait()
`
``
1222
`+
results.append(pickle.loads(pickle_bytes))
`
``
1223
+
``
1224
`+
t1 = threading.Thread(target=t)
`
``
1225
`+
t2 = threading.Thread(target=t)
`
``
1226
`+
t1.start()
`
``
1227
`+
t2.start()
`
``
1228
+
``
1229
`+
barrier.wait()
`
``
1230
`+
could have delay here
`
``
1231
`+
locker.barrier.wait()
`
``
1232
+
``
1233
`+
t1.join()
`
``
1234
`+
t2.join()
`
``
1235
+
``
1236
`+
from locking_import import ToBeUnpickled
`
``
1237
`+
self.assertEqual(
`
``
1238
`+
[type(x) for x in results],
`
``
1239
`+
[ToBeUnpickled] * 2)
`
``
1240
+
``
1241
+
1177
1242
``
1178
1243
`class AbstractPickleTests(unittest.TestCase):
`
1179
1244
`# Subclass must define self.dumps, self.loads.
`