Maintain a list of BufferedWriter objects. Flush them on exit.

In Python 3, the buffer and the underlying file object are separate
and so the order in which objects are finalized matters.  This is
unlike Python 2 where the file and buffer were a single object and
finalization was done for both at the same time.  In Python 3, if
the file is finalized and closed before the buffer then the data in
the buffer is lost.

This change adds a doubly linked list of open file buffers.  An atexit
hook ensures they are flushed before proceeding with interpreter
shutdown.  This is addition does not remove the need to properly close
files as there are other reasons why buffered data could get lost during
finalization.

Initial patch by Armin Rigo.

Author: nascheme

Lib/_pyio.py

@@ -1185,6 +1185,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()
@@ -2574,3 +2575,23 @@ def encoding(self):
     def detach(self):
         # This doesn't make sense on StringIO.
         self._unsupported("detach")
+
+
+# ____________________________________________________________
+
+import atexit, weakref
+
+_all_writers = weakref.WeakKeyDictionary()
+
+def _register_writer(w):
+    # keep weak-ref to buffered writer
+    _all_writers[w] = True
+
+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:
+        w.flush()
+atexit.register(_flush_all_writers)

Modules/_io/_iomodule.c

@@ -766,6 +766,8 @@ PyInit__io(void)
         !(_PyIO_empty_bytes = PyBytes_FromStringAndSize(NULL, 0)))
         goto fail;
 
+    _Py_PyAtExit(_PyIO_atexit_flush);
+
     state->initialized = 1;
 
     return m;

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

@@ -196,6 +196,10 @@ bufferediobase_write(PyObject *self, PyObject *args)
 }
 
 
+struct doubly_linked_s {
+    struct doubly_linked_s *prev, *next;
+};
+
 typedef struct {
     PyObject_HEAD
 
@@ -240,8 +244,15 @@ typedef struct {
 
     PyObject *dict;
     PyObject *weakreflist;
+
+    /* a doubly-linked chained list of "buffered" objects that need to
+       be flushed when the process exits */
+    struct doubly_linked_s buffered_writers_list;
 } buffered;
 
+static struct doubly_linked_s doubly_linked_end = {
+    &doubly_linked_end, &doubly_linked_end };
+
 /*
     Implementation notes:
 
@@ -386,6 +397,15 @@ _enter_buffered_busy(buffered *self)
         (self->buffer_size * (size / self->buffer_size)))
 
 
+static void
+remove_from_linked_list(buffered *self)
+{
+    self->buffered_writers_list.next->prev = self->buffered_writers_list.prev;
+    self->buffered_writers_list.prev->next = self->buffered_writers_list.next;
+    self->buffered_writers_list.prev = NULL;
+    self->buffered_writers_list.next = NULL;
+}
+
 static void
 buffered_dealloc(buffered *self)
 {
@@ -394,6 +414,8 @@ buffered_dealloc(buffered *self)
         return;
     _PyObject_GC_UNTRACK(self);
     self->ok = 0;
+    if (self->buffered_writers_list.next != NULL)
+        remove_from_linked_list(self);
     if (self->weakreflist != NULL)
         PyObject_ClearWeakRefs((PyObject *)self);
     Py_CLEAR(self->raw);
@@ -1817,10 +1839,34 @@ _io_BufferedWriter___init___impl(buffered *self, PyObject *raw,
     self->fast_closed_checks = (Py_TYPE(self) == &PyBufferedWriter_Type &&
                                 Py_TYPE(raw) == &PyFileIO_Type);
 
+    if (self->buffered_writers_list.next == NULL) {
+        self->buffered_writers_list.prev = &doubly_linked_end;
+        self->buffered_writers_list.next = doubly_linked_end.next;
+        doubly_linked_end.next->prev = &self->buffered_writers_list;
+        doubly_linked_end.next = &self->buffered_writers_list;
+    }
+
     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 (doubly_linked_end.next != &doubly_linked_end) {
+        buffered *buf = (buffered *)(((char *)doubly_linked_end.next) -
+                                     offsetof(buffered, buffered_writers_list));
+        remove_from_linked_list(buf);
+        buffered_flush(buf, NULL);
+        PyErr_Clear();
+    }
+}
+
 static Py_ssize_t
 _bufferedwriter_raw_write(buffered *self, char *start, Py_ssize_t len)
 {