Incrementally flush large cPickle output to file

Modules/cPickle.c

 #include "Python.h"
 #include "structmember.h"
 #include "cPickle.h"
 
 PyDoc_STRVAR(cPickle_module_documentation,
 "C implementation and optimization of the Python pickle module.");
 
 #ifndef Py_eval_input
 #include <graminit.h>
 #define Py_eval_input eval_input
 #endif /* Py_eval_input */
 
 #define DEL_LIST_SLICE(list, from, to) (PyList_SetSlice(list, from, to, NULL))
 
 /* Bump this when new opcodes are added to the pickle protocol. */
 #define HIGHEST_PROTOCOL 2
 
+/* Maximum size we ever allow the output buffer to get before flushing. */
+#define PICKLER_MAX_BUFSIZE 64*1024
+
 /*
  * Pickle opcodes.  These must be kept in synch with pickle.py.  Extensive
  * docs are in pickletools.py.
  */
 #define MARK        '('
 #define STOP        '.'
 #define POP         '0'
 #define POP_MARK    '1'
 #define DUP         '2'
 #define FLOAT       'F'
@@ skipping 536 matching lines @@
         int proto;
 
         /* bool, true if proto > 0 */
         int bin;
 
         /* Does the pickle contain any of GET, BINGET or LONG_BINGET? */
         int uses_gets;
 
         /* Write into a local buffer before flushing out to file. output_len
            tracks the current size; when output_len >= max_output_len, we
-           PyMem_RESIZE. */
+           PyMem_RESIZE.  In any case if we are pickling to a file we flush
+           the buffer after it grows to more than PICKLER_MAX_BUFSIZE to
+           avoid eating up all of memory in that case.  */
         Py_ssize_t max_output_len;
         Py_ssize_t output_len;
         char *output_buffer;
 
         int fast; /* Fast mode doesn't save in memo, don't use if circ ref */
         PyObject *dispatch_table;
         int fast_container; /* count nested container dumps */
         PyObject *fast_memo;
 } Picklerobject;
 
@@ skipping 64 matching lines @@
                 if (args) retval=args;
                 else {
                         PyErr_SetObject(ErrType,Py_None);
                         return NULL;
                 }
         PyErr_SetObject(ErrType,retval);
         Py_DECREF(retval);
         return NULL;
 }
 
+/* forward declarations needed by _Pickler_Write */
+static int _Pickler_FlushToFile(Picklerobject *);
+static int _Pickler_ClearBuffer(Picklerobject *);
+
 static int
 _Pickler_Write(Picklerobject *self, const char *s, Py_ssize_t n)
 {
         Py_ssize_t i;
         char *buffer;
 
         if (self->output_len + n > self->max_output_len) {
                 self->max_output_len = (self->output_len + n) * 2;
                 buffer = self->output_buffer;
                 PyMem_RESIZE(buffer, char, self->max_output_len);
                 if (buffer == NULL) {
                         PyErr_NoMemory();
                         return -1;
                 }
                 self->output_buffer = buffer;
         }
         /* This is faster than memcpy for the char *s we're copying. */
         for (i = 0; i < n; i++) {
                 self->output_buffer[self->output_len + i] = s[i];
         }
         self->output_len += n;
+
+        if (self.file && self->output_len >= PICKLER_MAX_BUFSIZE) {
+                /* suppress optimization since the optimizer won't get a
+                   look at the entire pickled object */

[Antoine Pitrou, 2010/01/13 13:08:12]

You could probably do an early test on `n >= PICKLER_MAX_BUFSIZE` as well. If
the chunk to be written is big, there is no point in buffering it and flushing
it immediately after.

[Skip, 2010/01/13 13:47:16]

Sure, but this seems ever-so-slightly easier to get correct.

[Skip, 2010/01/13 13:48:59]

Done.

+                self->uses_gets = 1;

[Antoine Pitrou, 2010/01/13 13:08:12]

I suppose this is unladen-swallow-specific? I see no trace of it in mainline
cPickle.

[Skip, 2010/01/13 13:47:16]

Correct.  See _Pickler_Optimize.  It removes PUT ops if the pickle doesn't use
any GET ops.  Given that I amy possibly flushing GET ops to the file and thus
making them unavailable to the optimizer I simply pretend that the pickle does
use GET.  That just defeats the optimizer when it's called at the end.  I'm sure
someone smarter than me could solve that problem.  I'm not intimiately familiar
with the pickle virtual machine.

+                if (_Pickler_FlushToFile(self) < 0)
+                        return -1;
+                if (_Pickler_ClearBuffer(self) < 0)
+                        return -1;
+        }
+
         return n;
 }
 
 static PyObject *
 _Pickler_GetString(Picklerobject *self)
 {
         return PyString_FromStringAndSize(self->output_buffer,
                                           self->output_len);
 }
 
@@ skipping 5877 matching lines @@
                 "1.0",  /* Original protocol 0 */
                 "1.1",  /* Protocol 0 + INST */
                 "1.2",  /* Original protocol 1 */
                 "1.3",  /* Protocol 1 + BINFLOAT */
                 "2.0"); /* Original protocol 2 */
         PyDict_SetItemString(d, "format_version", format_version);
         PyDict_SetItemString(d, "compatible_formats", compatible_formats);
         Py_XDECREF(format_version);
         Py_XDECREF(compatible_formats);
 }