bpo-17852: Fix PR #3372, flush BufferedWriter objects on exit.

Lib/_pyio.py

@@ -1182,6 +1182,7 @@ def __init__(self, raw, buffer_size=DEFAULT_BUFFER_SIZE):
         self.buffer_size = buffer_size
         self._write_buf = bytearray()
         self._write_lock = Lock()
+        _register_writer(self)
 
     def writable(self):
         return self.raw.writable()
@@ -2586,3 +2587,26 @@ def encoding(self):
     def detach(self):
         # This doesn't make sense on StringIO.
         self._unsupported("detach")
+
+
+# ____________________________________________________________
+
+import atexit, weakref
+
+_all_writers = weakref.WeakSet()
+
+def _register_writer(w):
+    # keep weak-ref to buffered writer
+    _all_writers.add(w)
+
+def _flush_all_writers():
+    # Ensure all buffered writers are flushed before proceeding with
+    # normal shutdown.  Otherwise, if the underlying file objects get
+    # finalized before the buffered writer wrapping it then any buffered
+    # data will be lost.
+    for w in _all_writers:
+        try:
+            w.flush()
+        except Exception:
+            pass
+atexit.register(_flush_all_writers)

Lib/io.py

@@ -97,3 +97,6 @@ class TextIOBase(_io._TextIOBase, IOBase):
     pass
 else:
     RawIOBase.register(_WindowsConsoleIO)
+
+import atexit
+atexit.register(_io._flush_all_buffers)

Misc/NEWS.d/next/Library/2017-12-13-00-00-37.bpo-17852.Q8BP8N.rst

@@ -1 +1,2 @@
-Revert incorrect fix based on misunderstanding of _Py_PyAtExit() semantics.
+Fix PR #3372, flush BufferedWriter objects on exit.  Use atexit.register()
+and not _Py_PyAtExit().

Modules/_io/_iomodule.c

@@ -609,6 +609,21 @@ iomodule_free(PyObject *mod) {
     iomodule_clear(mod);
 }
 
+/*[clinic input]
+_io._flush_all_buffers
+
+Flushes all buffered io objects.  Called by atexit.
+
+[clinic start generated code]*/
+
+static PyObject *
+_io__flush_all_buffers_impl(PyObject *module)
+/*[clinic end generated code: output=a242f507481504e1 input=8b6276ac61894717]*/
+{
+    _PyIO_atexit_flush();
+    Py_INCREF(Py_None);
+    return Py_None;
+}
 
 /*
  * Module definition
@@ -618,6 +633,7 @@ iomodule_free(PyObject *mod) {
 
 static PyMethodDef module_methods[] = {
     _IO_OPEN_METHODDEF
+    _IO__FLUSH_ALL_BUFFERS_METHODDEF
     {NULL, NULL}
 };
 

Modules/_io/_iomodule.h

@@ -183,3 +183,5 @@ extern PyObject *_PyIO_empty_str;
 extern PyObject *_PyIO_empty_bytes;
 
 extern PyTypeObject _PyBytesIOBuffer_Type;
+
+extern void _PyIO_atexit_flush(void);

Modules/_io/bufferedio.c

@@ -197,7 +197,7 @@ bufferediobase_write(PyObject *self, PyObject *args)
 }
 
 
-typedef struct {
+typedef struct _buffered {
     PyObject_HEAD
 
     PyObject *raw;
@@ -239,8 +239,18 @@ typedef struct {
 
     PyObject *dict;
     PyObject *weakreflist;
+
+    /* a doubly-linked chained list of "buffered" objects that need to
+       be flushed when the process exits */
+    struct _buffered *next, *prev;
 } buffered;
 
+/* the actual list of buffered objects */
+static buffered buffer_list_end = {
+    .next = &buffer_list_end,
+    .prev = &buffer_list_end
+};
+
 /*
     Implementation notes:
 
@@ -379,10 +389,21 @@ _enter_buffered_busy(buffered *self)
         (self->buffer_size * (size / self->buffer_size)))
 
 
+static void
+remove_from_linked_list(buffered *self)
+{
+    self->next->prev = self->prev;
+    self->prev->next = self->next;
+    self->prev = NULL;
+    self->next = NULL;
+}
+
 static void
 buffered_dealloc(buffered *self)
 {
     self->finalizing = 1;
+    if (self->next != NULL)
+        remove_from_linked_list(self);
     if (_PyIOBase_finalize((PyObject *) self) < 0)
         return;
     _PyObject_GC_UNTRACK(self);
@@ -1806,10 +1827,38 @@ _io_BufferedWriter___init___impl(buffered *self, PyObject *raw,
     self->fast_closed_checks = (Py_TYPE(self) == &PyBufferedWriter_Type &&
                                 Py_TYPE(raw) == &PyFileIO_Type);
 
+    if (self->next == NULL) {
+        self->prev = &buffer_list_end;
+        self->next = buffer_list_end.next;
+        buffer_list_end.next->prev = self;
+        buffer_list_end.next = self;
+    }
+
     self->ok = 1;
     return 0;
 }
 
+/*
+* Ensure all buffered writers are flushed before proceeding with
+* normal shutdown.  Otherwise, if the underlying file objects get
+* finalized before the buffered writer wrapping it then any buffered
+* data will be lost.
+*/
+void _PyIO_atexit_flush(void)
+{
+    while (buffer_list_end.next != &buffer_list_end) {
+        buffered *buf = buffer_list_end.next;
+        remove_from_linked_list(buf);
+        if (buf->ok && !buf->finalizing) {
+            /* good state and not finalizing */
+            Py_INCREF(buf);
+            buffered_flush(buf, NULL);
+            Py_DECREF(buf);
+            PyErr_Clear();
+        }
+    }
+}
+
 static Py_ssize_t
 _bufferedwriter_raw_write(buffered *self, char *start, Py_ssize_t len)
 {

Modules/_io/clinic/_iomodule.c.h

@@ -158,4 +158,22 @@ _io_open(PyObject *module, PyObject **args, Py_ssize_t nargs, PyObject *kwnames)
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=7e0ab289d8465580 input=a9049054013a1b77]*/
+
+PyDoc_STRVAR(_io__flush_all_buffers__doc__,
+"_flush_all_buffers($module, /)\n"
+"--\n"
+"\n"
+"Flushes all buffered io objects.  Called by atexit.");
+
+#define _IO__FLUSH_ALL_BUFFERS_METHODDEF    \
+    {"_flush_all_buffers", (PyCFunction)_io__flush_all_buffers, METH_NOARGS, _io__flush_all_buffers__doc__},
+
+static PyObject *
+_io__flush_all_buffers_impl(PyObject *module);
+
+static PyObject *
+_io__flush_all_buffers(PyObject *module, PyObject *Py_UNUSED(ignored))
+{
+    return _io__flush_all_buffers_impl(module);
+}
+/*[clinic end generated code: output=c4901164cf35b7a2 input=a9049054013a1b77]*/