Adds test for multithreaded writes to files.

elistevens · elistevens · commit dca6eb690f19 · 2024-05-16T18:37:49.000-07:00
Multithreaded writes to files (including standard calls to print() from threads) can result in assertion failures and potentially missed output due to race conditions in _io_TextIOWrapper_write_impl and _textiowrapper_writeflush. See: python#118138
diff --git a/Lib/test/test_io_threading.py b/Lib/test/test_io_threading.py
@@ -0,0 +1,36 @@
+import unittest
+from concurrent.futures import ThreadPoolExecutor
+
+
+NUM_THREADS = 100
+NUM_LOOPS = 10
+
+def looping_print_to_file(f, i):
+    for j in range(NUM_LOOPS):
+        # p = 'x' * 90 + '123456789'
+        p = f"{i:2}i,{j}j\n" + '0123456789' * 10
+        print(p, file=f)
+
+class IoThreadingTestCase(unittest.TestCase):
+    def test_io_threading(self):
+        with ThreadPoolExecutor(max_workers=NUM_THREADS) as executor:
+            with open('test_io_threading.out', 'w') as f:
+                for i in range(NUM_THREADS):
+                    executor.submit(looping_print_to_file, f, i)
+
+            executor.shutdown(wait=True)
+
+        with open('test_io_threading.out', 'r') as f:
+            lines = set(x.rstrip() for x in f.readlines() if ',' in x)
+
+        assert len(lines) == NUM_THREADS * NUM_LOOPS, repr(len(lines))
+
+        for i in range(NUM_THREADS):
+            for j in range(NUM_LOOPS):
+                p = f"{i:2}i,{j}j"
+
+                assert p in lines, repr(p)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/Modules/_io/textio.c b/Modules/_io/textio.c
@@ -1729,6 +1729,10 @@ _io_TextIOWrapper_write_impl(textio *self, PyObject *text)
             Py_DECREF(b);
             return NULL;
         }
+        // The assumption that self won't have been modified from another
+        // thread between the flush above and the assignment below is
+        // incorrect.
+        assert(self->pending_bytes == NULL);
         self->pending_bytes = b;
     }
     else if (!PyList_CheckExact(self->pending_bytes)) {

Original file line number	Diff line number	Diff line change
`@@ -1729,6 +1729,10 @@ _io_TextIOWrapper_write_impl(textio self, PyObject text)`
`1729`	`1729`	`Py_DECREF(b);`
`1730`	`1730`	`return NULL;`
`1731`	`1731`	`}`
	`1732`	`+ // The assumption that self won't have been modified from another`
	`1733`	`+ // thread between the flush above and the assignment below is`
	`1734`	`+ // incorrect.`
	`1735`	`+ assert(self->pending_bytes == NULL);`
`1732`	`1736`	`self->pending_bytes = b;`
`1733`	`1737`	`}`
`1734`	`1738`	`else if (!PyList_CheckExact(self->pending_bytes)) {`