chore: resolve review comment

Author: fatelei

Lib/test/test_array.py

@@ -1367,15 +1367,15 @@ def test_frombytearray(self):
         b = array.array(self.typecode, a)
         self.assertEqual(a, b)
 
-    def test_tofile_use_after_free(self):
-        CHUNK = 64 * 1024
-        victim = array.array('B', b'\0' * (CHUNK * 2))
+    def test_tofile_concurrent_mutation(self):
+        BLOCKSIZE = 64 * 1024
+        victim = array.array('B', b'\0' * (BLOCKSIZE * 2))
 
         class Writer:
-            armed = True
+            cleared = False
             def write(self, data):
-                if Writer.armed:
-                    Writer.armed = False
+                if not self.cleared:
+                    self.cleared = True
                     victim.clear()
                 return 0
 

Modules/arraymodule.c

@@ -1589,29 +1589,36 @@ array_array_tofile_impl(arrayobject *self, PyTypeObject *cls, PyObject *f)
 {
     /* Write 64K blocks at a time */
     /* XXX Make the block size settable */
-    int BLOCKSIZE = 64*1024;
-    Py_ssize_t offset = 0;
+    const Py_ssize_t BLOCKSIZE = 64*1024;
+    const Py_ssize_t itemsize = self->ob_descr->itemsize;
+
+    Py_ssize_t nbytes = Py_SIZE(self) * itemsize;
+    Py_ssize_t nblocks = (nbytes + BLOCKSIZE - 1) / BLOCKSIZE;
 
     if (Py_SIZE(self) == 0)
         goto done;
 
     array_state *state = get_array_state_by_class(cls);
     assert(state != NULL);
 
-    while (1) {
+    for (Py_ssize_t i = 0; i < nblocks; i++) {
         if (self->ob_item == NULL || Py_SIZE(self) == 0) {
             break;
         }
 
-        Py_ssize_t current_nbytes = Py_SIZE(self) * self->ob_descr->itemsize;
+        if (Py_SIZE(self) > PY_SSIZE_T_MAX / itemsize) {
+            return PyErr_NoMemory();
+        }
 
+        Py_ssize_t current_nbytes = Py_SIZE(self) * itemsize;
+        const Py_ssize_t offset = i * BLOCKSIZE;
         if (offset >= current_nbytes) {
             break;
         }
 
-        Py_ssize_t size = BLOCKSIZE;
-        if (offset + size > current_nbytes) {
-            size = current_nbytes - offset;
+        Py_ssize_t size = current_nbytes - offset;
+        if (size > BLOCKSIZE) {
+            size = BLOCKSIZE;
         }
 
         char* ptr = self->ob_item + offset;
@@ -1625,9 +1632,7 @@ array_array_tofile_impl(arrayobject *self, PyTypeObject *cls, PyObject *f)
         Py_DECREF(bytes);
         if (res == NULL)
             return NULL;
-        Py_DECREF(res);
-
-        offset += size;
+        Py_DECREF(res); /* drop write result */
     }
 
   done: