Skip to content

Commit

Permalink
[3.11] pythongh-119506: fix _io.TextIOWrapper.write() write during fl…
Browse files Browse the repository at this point in the history
…ush (python#119507) (python#120314)

Co-authored-by: Hugo van Kemenade <[email protected]>
fix _io.TextIOWrapper.write() write during flush (python#119507)
  • Loading branch information
methane authored Aug 9, 2024
1 parent b396360 commit 8a978a7
Show file tree
Hide file tree
Showing 3 changed files with 67 additions and 22 deletions.
22 changes: 22 additions & 0 deletions Lib/test/test_io.py
Original file line number Diff line number Diff line change
Expand Up @@ -3960,6 +3960,28 @@ def write(self, data):
t.write("x"*chunk_size)
self.assertEqual([b"abcdef", b"ghi", b"x"*chunk_size], buf._write_stack)

def test_issue119506(self):
chunk_size = 8192

class MockIO(self.MockRawIO):
written = False
def write(self, data):
if not self.written:
self.written = True
t.write("middle")
return super().write(data)

buf = MockIO()
t = self.TextIOWrapper(buf)
t.write("abc")
t.write("def")
# writing data which size >= chunk_size cause flushing buffer before write.
t.write("g" * chunk_size)
t.flush()

self.assertEqual([b"abcdef", b"middle", b"g"*chunk_size],
buf._write_stack)


class PyTextIOWrapperTest(TextIOWrapperTest):
io = pyio
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Fix :meth:`!io.TextIOWrapper.write` method breaks internal buffer when the method is called again during flushing internal buffer.
66 changes: 44 additions & 22 deletions Modules/_io/textio.c
Original file line number Diff line number Diff line change
Expand Up @@ -1701,34 +1701,56 @@ _io_TextIOWrapper_write_impl(textio *self, PyObject *text)
bytes_len = PyBytes_GET_SIZE(b);
}

if (self->pending_bytes == NULL) {
self->pending_bytes_count = 0;
self->pending_bytes = b;
}
else if (self->pending_bytes_count + bytes_len > self->chunk_size) {
// Prevent to concatenate more than chunk_size data.
if (_textiowrapper_writeflush(self) < 0) {
Py_DECREF(b);
return NULL;
// We should avoid concatinating huge data.
// Flush the buffer before adding b to the buffer if b is not small.
// https://github.com/python/cpython/issues/87426
if (bytes_len >= self->chunk_size) {
// _textiowrapper_writeflush() calls buffer.write().
// self->pending_bytes can be appended during buffer->write()
// or other thread.
// We need to loop until buffer becomes empty.
// https://github.com/python/cpython/issues/118138
// https://github.com/python/cpython/issues/119506
while (self->pending_bytes != NULL) {
if (_textiowrapper_writeflush(self) < 0) {
Py_DECREF(b);
return NULL;
}
}
self->pending_bytes = b;
}
else if (!PyList_CheckExact(self->pending_bytes)) {
PyObject *list = PyList_New(2);
if (list == NULL) {
Py_DECREF(b);
return NULL;
}
PyList_SET_ITEM(list, 0, self->pending_bytes);
PyList_SET_ITEM(list, 1, b);
self->pending_bytes = list;

if (self->pending_bytes == NULL) {
assert(self->pending_bytes_count == 0);
self->pending_bytes = b;
}
else {
if (PyList_Append(self->pending_bytes, b) < 0) {
Py_DECREF(b);
return NULL;
if (!PyList_CheckExact(self->pending_bytes)) {
PyObject *list = PyList_New(0);
if (list == NULL) {
Py_DECREF(b);
return NULL;
}
// PyList_New() may trigger GC and other thread may call write().
// So, we need to check the self->pending_bytes is a list again.
if (PyList_CheckExact(self->pending_bytes)) {
// Releasing empty list won't trigger GC and/or __del__.
Py_DECREF(list);
}
else {
if (PyList_Append(list, self->pending_bytes) < 0) {
Py_DECREF(list);
Py_DECREF(b);
return NULL;
}
Py_SETREF(self->pending_bytes, list);
}
}

int ret = PyList_Append(self->pending_bytes, b);
Py_DECREF(b);
if (ret < 0) {
return NULL;
}
}

self->pending_bytes_count += bytes_len;
Expand Down

0 comments on commit 8a978a7

Please sign in to comment.