py/modio: Handle partial writes on BufferedWriter.

klondi · DvdGiessen · commit 92aaf70b8e3e · 2025-05-14T17:46:26.000+02:00
To simplify the logic create bufwriter_do_write. This function keeps
track of the data left on self-&gt;len and, if a partial write happens,
this function uses memmove to move the remaining data back to the
beginning of the buffer.

This allows simplifying bufwriter_write and bufwriter_flush significantly.

bufwriter_flush now only needs to call this function if the buffer has
some data stored and check for any error codes.
Additionally, if the buffer is only partially flushed we notify the user
too (so that they know there might be data left).

bufwriter_write now just needs to call this function whenever the buffer
gets full and copy the input into the buffer.
It will return when either no data is written at all or when all of the
input is consumed. In the case of a partial write it returns exactly the
amount of data which was written.

Additionally allow caching of errors to better handle partial writes.
Until now if an error occurred during the write, the error would be
raised and the caller had no way to know if any data was written at all
(for example in prior calls if more than one block of data was passed as
input). Now when we have written out some data and an error happens, we
reset the buffer to the state it would have if it did not contain the
data that was not written (and which was not buffered previously), and
then, we return the data that was written (if any) or raise an error if
no data from the input was written.

This allows the programmer better control of writes. In particular, the
programmer will know exactly how much of its last input data was written,
consequently allowing it to handle whatever data left to be written in
a better way.

Signed-off-by: Francisco Blas (klondike) Izquierdo Riera &lt;klondike@klondike.es&gt;
diff --git a/py/modio.c b/py/modio.c
@@ -27,6 +27,8 @@
 #include <assert.h>
 #include <string.h>
 
+#include "py/mpconfig.h"
+#include "py/misc.h"
 #include "py/runtime.h"
 #include "py/builtin.h"
 #include "py/stream.h"
@@ -113,6 +115,7 @@ typedef struct _mp_obj_bufwriter_t {
     mp_obj_t stream;
     size_t alloc;
     size_t len;
+    int error;
     byte buf[0];
 } mp_obj_bufwriter_t;
 
@@ -123,61 +126,97 @@ static mp_obj_t bufwriter_make_new(const mp_obj_type_t *type, size_t n_args, siz
     o->stream = args[0];
     o->alloc = alloc;
     o->len = 0;
+    o->error = 0;
     return o;
 }
 
+// Writes out the data stored in the buffer so far
+static int bufwriter_do_write(mp_obj_bufwriter_t *self) {
+    int rv = 0;
+    // This cannot return 0 without an error
+    mp_uint_t out_sz = mp_stream_write_exactly(self->stream, self->buf, self->len, &rv);
+    self->len -= out_sz;
+    // Copy the non written characters back to the beginning
+    if (self->len != 0) {
+        // Use memmove since there might be overlaps
+        memmove(self->buf, self->buf + out_sz, self->len);
+    }
+    return rv;
+}
+
 static mp_uint_t bufwriter_write(mp_obj_t self_in, const void *buf, mp_uint_t size, int *errcode) {
     mp_obj_bufwriter_t *self = MP_OBJ_TO_PTR(self_in);
 
     mp_uint_t org_size = size;
+    // Alloc should always remain the same so cache it.
+    size_t alloc = self->alloc;
+    mp_uint_t rem = 0; // No data has been copied from the buffer so far
 
-    while (size > 0) {
-        mp_uint_t rem = self->alloc - self->len;
-        if (size < rem) {
-            memcpy(self->buf + self->len, buf, size);
-            self->len += size;
-            return org_size;
-        }
+    if (self->error != 0) {
+        *errcode = self->error;
+        self->error = 0;
+        return MP_STREAM_ERROR;
+    }
 
+    // Using this form allows ensuring that we try to empty the buffer if it is
+    // full without trying to extract any data first.
+    while (true) {
         // Buffer flushing policy here is to flush entire buffer all the time.
         // This allows e.g. to have a block device as backing storage and write
         // entire block to it. memcpy below is not ideal and could be optimized
         // in some cases. But the way it is now it at least ensures that buffer
         // is word-aligned, to guard against obscure cases when it matters, e.g.
         // https://github.com/micropython/micropython/issues/1863
+        // Try to empty the buffer first
+        if (self->len == alloc) {
+            if ((*errcode = bufwriter_do_write(self)) != 0) {
+                // If this is the first write with data from the new buffer
+                // But no data from the new buffer was written out,
+                // then remove the data added from the new buffer and raise
+                // the error
+                if (org_size == size + rem && self->len >= rem) {
+                    // Remove the extra non-written data from the buffer and error
+                    self->len -= rem;
+                    return MP_STREAM_ERROR;
+                }
+                // Some data from the new buffer has been written rollback as much
+                // as possible and return what was written then raise an error
+                // on the next call.
+                size += self->len;
+                self->len = 0;
+                self->error = *errcode;
+                *errcode = 0;
+                return org_size - size;
+            }
+        }
+        // No data left to write
+        if (size == 0) {
+            return org_size;
+        }
+
+        rem = MIN(alloc - self->len, size);
         memcpy(self->buf + self->len, buf, rem);
+        self->len += rem;
         buf = (byte *)buf + rem;
         size -= rem;
-        mp_uint_t out_sz = mp_stream_write_exactly(self->stream, self->buf, self->alloc, errcode);
-        (void)out_sz;
-        if (*errcode != 0) {
-            return MP_STREAM_ERROR;
-        }
-        // TODO: try to recover from a case of non-blocking stream, e.g. move
-        // remaining chunk to the beginning of buffer.
-        assert(out_sz == self->alloc);
-        self->len = 0;
     }
-
-    return org_size;
 }
 
 static mp_obj_t bufwriter_flush(mp_obj_t self_in) {
     mp_obj_bufwriter_t *self = MP_OBJ_TO_PTR(self_in);
-
-    if (self->len != 0) {
-        int err;
-        mp_uint_t out_sz = mp_stream_write_exactly(self->stream, self->buf, self->len, &err);
-        (void)out_sz;
-        // TODO: try to recover from a case of non-blocking stream, e.g. move
-        // remaining chunk to the beginning of buffer.
-        assert(out_sz == self->len);
-        self->len = 0;
-        if (err != 0) {
-            mp_raise_OSError(err);
+    int err = self->error;
+    if (err == 0 && self->len != 0) {
+        err = bufwriter_do_write(self);
+        // If we couldn't flush the whole buffer notify the user.
+        if (err == 0 && self->len != 0) {
+            err = MP_EAGAIN;
         }
     }
-
+    // If there is an error raise it
+    if (err != 0) {
+        self->error = 0;
+        mp_raise_OSError(err);
+    }
     return mp_const_none;
 }
 static MP_DEFINE_CONST_FUN_OBJ_1(bufwriter_flush_obj, bufwriter_flush);
