gh-157242: Add byteswriter_check_consistency()

Doc/c-api/bytes.rst

@@ -305,8 +305,8 @@ object.
 
    A bytes writer object.
 
-   The API is **not thread safe**. A :c:type:`PyBytesWriter` object must only
-   be used by a single thread, it must not be shared between threads.
+   The API is **not thread safe**. To share a writer with multiple threads, a
+   critical section or a lock is needed.
 
    The instance must be destroyed by :c:func:`PyBytesWriter_Finish` on
    success, or :c:func:`PyBytesWriter_Discard` on error.

Include/internal/pycore_unicodeobject.h

@@ -116,7 +116,8 @@ _PyUnicodeWriter_CanWrite(_PyUnicodeWriter *writer)
     PyObject *buffer = writer->buffer;
     assert(buffer != NULL);
     // Do not use _PyObject_IsUniquelyReferenced(): the caller can have its own
-    // lock to prevent a writer being used by two theads at the same time.
+    // lock to prevent a writer from being used by two threads at the same
+    // time.
     assert(Py_REFCNT(buffer) == 1);
     assert(PyUnstable_Unicode_GET_CACHED_HASH(buffer) == -1);
     assert(!PyUnicode_CHECK_INTERNED(buffer));

Lib/test/test_capi/test_bytes.py

@@ -1,8 +1,10 @@
 import sys
 import textwrap
+import threading
 import unittest
 from test import support
 from test.support import import_helper
+from test.support import threading_helper
 from test.support.script_helper import assert_python_failure
 
 _testlimitedcapi = import_helper.import_module('_testlimitedcapi')
@@ -461,6 +463,14 @@ def test_resize(self):
         self.assertEqual(writer.finish(),
                          b's' * small + b'L' * (large - small))
 
+        # Make sure that it's possible to write after a resize to zero
+        # when a bytes/bytearray object is allocated.
+        writer = self.create_writer()
+        writer.resize(self.LARGE_BUFFER)
+        writer.resize(0)
+        writer.write_bytes(b'abc', 3)
+        self.assertEqual(writer.finish(), b'abc')
+
         # invalid size
         for size in (self.SMALL_BUFFER, self.LARGE_BUFFER):
             with self.subTest(size=size):
@@ -665,6 +675,36 @@ def test_get_data_canary(self):
         self.assertEqual(get_data_canary(writer),
                          b'abc123' + CANARY_BYTE)
 
+    @threading_helper.requires_working_threading()
+    def test_thread(self):
+        # PyBytesWriter can be used by multiple threads: it's up to the caller
+        # to implement a lock to prevent concurrent accesses.
+        writer = self.create_writer(0)
+        size = None
+        data = None
+
+        def thread_func(writer, LARGE_BUFFER):
+            nonlocal size, data
+
+            # create a bytes object for the buffer
+            writer.write_bytes(b'x' * LARGE_BUFFER, LARGE_BUFFER)
+
+            # so we can check a write with a bytes object
+            writer.write_bytes(b'yz', 2)
+            writer.format_i(b'i=%i', 5)
+            writer.resize(10)
+            data = writer.get_data()
+            size = writer.get_size()
+
+        thread = threading.Thread(target=thread_func,
+                                  args=(writer, self.LARGE_BUFFER))
+        thread.start()
+        threading_helper.join_thread(thread)
+
+        self.assertEqual(size, 10)
+        self.assertEqual(data, b'x' * 10)
+        self.assertEqual(writer.finish(), b'x' * 10)
+
 
 class BytesWriterTest(BaseWriterTest, unittest.TestCase):
     RESULT_TYPE = bytes

Objects/bytesobject.c

@@ -3329,18 +3329,22 @@ PyBytes_ConcatAndDel(PyObject **pv, PyObject *w)
 //
 // Usage: assert(_PyBytes_IsMutable(obj)).
 int
-_PyBytes_IsMutable(PyObject *v)
+_PyBytes_IsMutable(PyObject *self)
 {
-    // Singleton objects must never be modified
-    assert(!_Py_IsImmortal(v));
+    assert(PyBytes_Check(self));
+    // Do not use _PyObject_IsUniquelyReferenced(): this function is called
+    // by bytearray and PyBytesWriter which can be used by multiple threads.
+    assert(Py_REFCNT(self) == 1);
+    assert(!_Py_IsImmortal(self));
 
-    Py_ssize_t size = PyBytes_GET_SIZE(v);
+    // Check that the object is not a singleton
+    Py_ssize_t size = PyBytes_GET_SIZE(self);
     if (size == 0) {
-        assert(v != bytes_get_empty());
+        assert(self != bytes_get_empty());
     }
     else if (size == 1) {
-        unsigned char ch = PyBytes_AS_STRING(v)[0];
-        assert(v != (PyObject*)CHARACTER(ch));
+        unsigned char ch = PyBytes_AS_STRING(self)[0];
+        assert(self != (PyObject*)CHARACTER(ch));
     }
     return 1;
 }
@@ -3675,20 +3679,6 @@ byteswriter_allocated(PyBytesWriter *writer)
 
 
 #ifdef Py_DEBUG
-static void
-byteswriter_check_canary_byte(PyBytesWriter *writer)
-{
-    const unsigned char *data = (const unsigned char*)byteswriter_data(writer);
-    unsigned char canary = data[writer->size];
-    if (canary != PyBytesWriter_CANARY_BYTE) {
-        _Py_FatalErrorFormat(__func__,
-                             "Buffer overflow detected in PyBytesWriter %p "
-                             "at position %zd",
-                             writer, writer->size);
-    }
-}
-
-
 static void
 byteswriter_write_canary_byte(PyBytesWriter *writer)
 {
@@ -3710,6 +3700,45 @@ byteswriter_reset_trailing_byte(PyBytesWriter *writer)
 #endif
 
 
+#ifndef NDEBUG
+static int
+byteswriter_check_consistency(PyBytesWriter *writer)
+{
+    PyObject *obj = writer->obj;
+    if (obj != NULL) {
+        if (writer->use_bytearray) {
+            assert(PyByteArray_CheckExact(obj));
+            // Do not use _PyObject_IsUniquelyReferenced(): the caller can have
+            // its own lock to prevent a writer from being used by two threads
+            // at the same time.
+            assert(Py_REFCNT(obj) == 1);
+            PyByteArrayObject *bytearray = (PyByteArrayObject*)obj;
+            obj = bytearray->ob_bytes_object;
+            assert(obj != NULL);
+        }
+
+        // Code adapted from _PyBytes_IsMutable()
+        assert(PyBytes_CheckExact(obj));
+        assert(_PyBytes_IsMutable(obj));
+        // -1 since the last small buffer byte is used as the canary byte
+        assert((size_t)PyBytes_GET_SIZE(obj) > (sizeof(writer->small_buffer) - 1));
+    }
+
+#ifdef Py_DEBUG
+    const unsigned char *data = (const unsigned char*)byteswriter_data(writer);
+    unsigned char canary = data[writer->size];
+    if (canary != PyBytesWriter_CANARY_BYTE) {
+        _Py_FatalErrorFormat(__func__,
+                             "Buffer overflow detected in PyBytesWriter %p "
+                             "at position %zd",
+                             writer, writer->size);
+    }
+#endif
+    return 1;
+}
+#endif
+
+
 #ifdef MS_WINDOWS
    /* On Windows, overallocate by 50% is the best factor */
 #  define OVERALLOCATE_FACTOR 2
@@ -3743,13 +3772,15 @@ byteswriter_resize(PyBytesWriter *writer, Py_ssize_t new_size, int resize)
                 // bytearray can override the canary byte on error
                 byteswriter_write_canary_byte(writer);
 #endif
+                assert(byteswriter_check_consistency(writer));
                 return -1;
             }
         }
         else {
             // Can raise MemoryError or OverflowError
             if (_PyBytes_ResizeKeepOnError(&writer->obj, alloc)) {
                 assert(writer->obj != NULL);
+                assert(byteswriter_check_consistency(writer));
                 return -1;
             }
             assert(_PyBytes_IsMutable(writer->obj));
@@ -3821,7 +3852,7 @@ byteswriter_create(Py_ssize_t size, int use_bytearray)
     if (size >= 1) {
         if (byteswriter_resize(writer, size, 0) < 0) {
 #ifdef Py_DEBUG
-            // Write the canary byte so byteswriter_check_canary_byte()
+            // Write the canary byte so byteswriter_check_consistency()
             // doesn't fail in PyBytesWriter_Discard()
             byteswriter_write_canary_byte(writer);
 #endif
@@ -3835,6 +3866,7 @@ byteswriter_create(Py_ssize_t size, int use_bytearray)
            byteswriter_allocated(writer));
     byteswriter_write_canary_byte(writer);
 #endif
+    assert(byteswriter_check_consistency(writer));
     return writer;
 }
 
@@ -3858,8 +3890,8 @@ PyBytesWriter_Discard(PyBytesWriter *writer)
         return;
     }
 
+    assert(byteswriter_check_consistency(writer));
 #ifdef Py_DEBUG
-    byteswriter_check_canary_byte(writer);
     if (writer->obj != NULL) {
         byteswriter_reset_trailing_byte(writer);
     }
@@ -3873,6 +3905,8 @@ PyBytesWriter_Discard(PyBytesWriter *writer)
 PyObject*
 PyBytesWriter_FinishWithSize(PyBytesWriter *writer, Py_ssize_t size)
 {
+    assert(byteswriter_check_consistency(writer));
+
     // Check for negative size here to raise ValueError in all cases, rather
     // than having a different exception depending on the code path. For
     // example, _PyBytes_Resize() raises SystemError on negative size.
@@ -3886,10 +3920,6 @@ PyBytesWriter_FinishWithSize(PyBytesWriter *writer, Py_ssize_t size)
         goto error;
     }
 
-#ifdef Py_DEBUG
-    byteswriter_check_canary_byte(writer);
-#endif
-
     PyObject *result;
     if (size == 0) {
         result = bytes_get_empty();
@@ -3938,7 +3968,7 @@ PyBytesWriter_FinishWithSize(PyBytesWriter *writer, Py_ssize_t size)
     }
 
 #ifdef Py_DEBUG
-    // Reset the writer, so byteswriter_check_canary_byte() doesn't fail
+    // Reset the writer, so byteswriter_check_consistency() doesn't fail
     // in PyBytesWriter_Discard().
     writer->size = 0;
     byteswriter_write_canary_byte(writer);
@@ -3970,9 +4000,7 @@ PyBytesWriter_FinishWithPointer(PyBytesWriter *writer, void *buf)
 void*
 PyBytesWriter_GetData(PyBytesWriter *writer)
 {
-#ifdef Py_DEBUG
-    byteswriter_check_canary_byte(writer);
-#endif
+    assert(byteswriter_check_consistency(writer));
 
     return byteswriter_data(writer);
 }
@@ -3981,9 +4009,7 @@ PyBytesWriter_GetData(PyBytesWriter *writer)
 Py_ssize_t
 PyBytesWriter_GetSize(PyBytesWriter *writer)
 {
-#ifdef Py_DEBUG
-    byteswriter_check_canary_byte(writer);
-#endif
+    assert(byteswriter_check_consistency(writer));
 
     return _PyBytesWriter_GetSize(writer);
 }
@@ -3992,9 +4018,7 @@ PyBytesWriter_GetSize(PyBytesWriter *writer)
 int
 PyBytesWriter_Resize(PyBytesWriter *writer, Py_ssize_t new_size)
 {
-#ifdef Py_DEBUG
-    byteswriter_check_canary_byte(writer);
-#endif
+    assert(byteswriter_check_consistency(writer));
 
     if (new_size < 0) {
         PyErr_SetString(PyExc_ValueError, "size must be >= 0");
@@ -4012,6 +4036,7 @@ PyBytesWriter_Resize(PyBytesWriter *writer, Py_ssize_t new_size)
 #ifdef Py_DEBUG
     byteswriter_write_canary_byte(writer);
 #endif
+    assert(byteswriter_check_consistency(writer));
     return 0;
 }
 
@@ -4031,9 +4056,7 @@ _PyBytesWriter_ResizeAndUpdatePointer(PyBytesWriter *writer, Py_ssize_t size,
 int
 PyBytesWriter_Grow(PyBytesWriter *writer, Py_ssize_t grow)
 {
-#ifdef Py_DEBUG
-    byteswriter_check_canary_byte(writer);
-#endif
+    assert(byteswriter_check_consistency(writer));
 
     if (grow == 0) {
         // Nothing to do
@@ -4064,6 +4087,7 @@ PyBytesWriter_Grow(PyBytesWriter *writer, Py_ssize_t grow)
 #ifdef Py_DEBUG
     byteswriter_write_canary_byte(writer);
 #endif
+    assert(byteswriter_check_consistency(writer));
     return 0;
 }
 
@@ -4099,6 +4123,8 @@ PyBytesWriter_WriteBytes(PyBytesWriter *writer,
     }
     char *buf = byteswriter_data(writer);
     memcpy(buf + pos, bytes, size);
+
+    assert(byteswriter_check_consistency(writer));
     return 0;
 }
 
@@ -4127,14 +4153,12 @@ PyBytesWriter_Format(PyBytesWriter *writer, const char *format, ...)
 static Py_ssize_t
 _PyBytesWriter_ResizeToAllocated(PyBytesWriter *writer)
 {
-#ifdef Py_DEBUG
-    byteswriter_check_canary_byte(writer);
-#endif
-
     Py_ssize_t allocated = byteswriter_allocated(writer);
     writer->size = allocated;
 #ifdef Py_DEBUG
     byteswriter_write_canary_byte(writer);
 #endif
+
+    assert(byteswriter_check_consistency(writer));
     return allocated;
 }