issue #233: ffi.init_once()

5 files changed, 224 insertions, 1 deletions

diff --git a/c/_cffi_backend.c b/c/_cffi_backend.c
index 1fc210b..4ea2c47 100644
--- a/c/_cffi_backend.c
+++ b/c/_cffi_backend.c
@@ -103,7 +103,11 @@
 #endif
 
 #if PY_MAJOR_VERSION < 3
-#define PyCapsule_New(pointer, name, destructor)        \
+# undef PyCapsule_GetPointer
+# undef PyCapsule_New
+# define PyCapsule_GetPointer(capsule, name) \
+    (PyCObject_AsVoidPtr(capsule))
+# define PyCapsule_New(pointer, name, destructor) \
     (PyCObject_FromVoidPtr(pointer, destructor))
 #endif
 
diff --git a/c/ffi_obj.c b/c/ffi_obj.c
index 737a0d9..f276e21 100644
--- a/c/ffi_obj.c
+++ b/c/ffi_obj.c
@@ -24,6 +24,7 @@
 struct FFIObject_s {
     PyObject_HEAD
     PyObject *gc_wrefs, *gc_wrefs_freelist;
+    PyObject *init_once_cache;
     struct _cffi_parse_info_s info;
     char ctx_is_static, ctx_is_nonempty;
     builder_c_t types_builder;
@@ -52,6 +53,7 @@ static FFIObject *ffi_internal_new(PyTypeObject *ffitype,
     }
     ffi->gc_wrefs = NULL;
     ffi->gc_wrefs_freelist = NULL;
+    ffi->init_once_cache = NULL;
     ffi->info.ctx = &ffi->types_builder.ctx;
     ffi->info.output = internal_output;
     ffi->info.output_size = FFI_COMPLEXITY_OUTPUT;
@@ -65,6 +67,7 @@ static void ffi_dealloc(FFIObject *ffi)
     PyObject_GC_UnTrack(ffi);
     Py_XDECREF(ffi->gc_wrefs);
     Py_XDECREF(ffi->gc_wrefs_freelist);
+    Py_XDECREF(ffi->init_once_cache);
 
     free_builder_c(&ffi->types_builder, ffi->ctx_is_static);
 
@@ -881,6 +884,130 @@ PyDoc_STRVAR(ffi_memmove_doc,
 #define ffi_memmove  b_memmove     /* ffi_memmove() => b_memmove()
                                       from _cffi_backend.c */
 
+#ifdef WITH_THREAD
+# include "pythread.h"
+#else
+typedef void *PyThread_type_lock;
+# define PyThread_allocate_lock()        ((void *)-1)
+# define PyThread_free_lock(lock)        ((void)(lock))
+# define PyThread_acquire_lock(lock, _)  ((void)(lock))
+# define PyThread_release_lock(lock)     ((void)(lock))
+#endif
+
+PyDoc_STRVAR(ffi_init_once_doc,
+             "XXX document me");
+
+#if PY_MAJOR_VERSION < 3
+/* PyCapsule_New is redefined to be PyCObject_FromVoidPtr in _cffi_backend,
+   which gives 2.6 compatibility; but the destructor signature is different */
+static void _free_init_once_lock(void *lock)
+{
+    PyThread_free_lock((PyThread_type_lock)lock);
+}
+#else
+static void _free_init_once_lock(PyObject *capsule)
+{
+    PyThread_type_lock lock;
+    lock = PyCapsule_GetPointer(capsule, "cffi_init_once_lock");
+    if (lock != NULL)
+        PyThread_free_lock(lock);
+}
+#endif
+
+static PyObject *ffi_init_once(FFIObject *self, PyObject *args, PyObject *kwds)
+{
+    static char *keywords[] = {"func", "tag", NULL};
+    PyObject *cache, *func, *tag, *tup, *res, *x, *lockobj;
+    PyThread_type_lock lock;
+
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO", keywords, &func, &tag))
+        return NULL;
+
+    /* a lot of fun with reference counting and error checking
+       in this function */
+
+    /* atomically get or create a new dict (no GIL release) */
+    cache = self->init_once_cache;
+    if (cache == NULL) {
+        cache = PyDict_New();
+        if (cache == NULL)
+            return NULL;
+        self->init_once_cache = cache;
+    }
+
+    /* get the tuple from cache[tag], or make a new one: (False, lock) */
+    tup = PyDict_GetItem(cache, tag);
+    if (tup == NULL) {
+        lock = PyThread_allocate_lock();
+        if (lock == NULL)
+            return NULL;
+        x = PyCapsule_New(lock, "cffi_init_once_lock", _free_init_once_lock);
+        if (x == NULL) {
+            PyThread_free_lock(lock);
+            return NULL;
+        }
+        tup = PyTuple_Pack(2, Py_False, x);
+        Py_DECREF(x);
+        if (tup == NULL)
+            return NULL;
+        x = tup;
+
+        /* Possible corner case if 'tag' is an object overriding __eq__
+           in pure Python: the GIL may be released when we are running it.
+           We really need to call dict.setdefault(). */
+        tup = PyObject_CallMethod(cache, "setdefault", "OO", tag, x);
+        Py_DECREF(x);
+        if (tup == NULL)
+            return NULL;
+
+        Py_DECREF(tup);   /* there is still a ref inside the dict */
+    }
+
+    res = PyTuple_GET_ITEM(tup, 1);
+    Py_INCREF(res);
+
+    if (PyTuple_GET_ITEM(tup, 0) == Py_True) {
+        /* tup == (True, result): return the result. */
+        return res;
+    }
+
+    /* tup == (False, lock) */
+    lockobj = res;
+    lock = (PyThread_type_lock)PyCapsule_GetPointer(lockobj,
+                                                    "cffi_init_once_lock");
+    if (lock == NULL) {
+        Py_DECREF(lockobj);
+        return NULL;
+    }
+
+    Py_BEGIN_ALLOW_THREADS
+    PyThread_acquire_lock(lock, WAIT_LOCK);
+    Py_END_ALLOW_THREADS
+
+    x = PyDict_GetItem(cache, tag);
+    if (x != NULL && PyTuple_GET_ITEM(x, 0) == Py_True) {
+        /* the real result was put in the dict while we were waiting
+           for PyThread_acquire_lock() above */
+        res = PyTuple_GET_ITEM(x, 1);
+        Py_INCREF(res);
+    }
+    else {
+        res = PyObject_CallFunction(func, "");
+        if (res != NULL) {
+            tup = PyTuple_Pack(2, Py_True, res);
+            if (tup == NULL || PyDict_SetItem(cache, tag, tup) < 0) {
+                Py_XDECREF(tup);
+                Py_DECREF(res);
+                res = NULL;
+            }
+        }
+    }
+
+    PyThread_release_lock(lock);
+    Py_DECREF(lockobj);
+    return res;
+}
+
 
 #define METH_VKW  (METH_VARARGS | METH_KEYWORDS)
 static PyMethodDef ffi_methods[] = {
@@ -898,6 +1025,7 @@ static PyMethodDef ffi_methods[] = {
 #ifdef MS_WIN32
  {"getwinerror",(PyCFunction)ffi_getwinerror,METH_VKW,     ffi_getwinerror_doc},
 #endif
+ {"init_once",  (PyCFunction)ffi_init_once,  METH_VKW,     ffi_init_once_doc},
  {"integer_const",(PyCFunction)ffi_int_const,METH_VKW,     ffi_int_const_doc},
  {"memmove",    (PyCFunction)ffi_memmove,    METH_VKW,     ffi_memmove_doc},
  {"new",        (PyCFunction)ffi_new,        METH_VKW,     ffi_new_doc},
diff --git a/cffi/api.py b/cffi/api.py
index 0a98e05..bcc23af 100644
--- a/cffi/api.py
+++ b/cffi/api.py
@@ -72,6 +72,7 @@ class FFI(object):
         self._cdefsources = []
         self._included_ffis = []
         self._windows_unicode = None
+        self._init_once_cache = {}
         if hasattr(backend, 'set_ffi'):
             backend.set_ffi(self)
         for name in backend.__dict__:
@@ -598,6 +599,30 @@ class FFI(object):
         return recompile(self, module_name, source, tmpdir=tmpdir,
                          source_extension=source_extension, **kwds)
 
+    def init_once(self, func, tag):
+        # Read _init_once_cache[tag], which is either (False, lock) if
+        # we're calling the function now in some thread, or (True, result).
+        # Don't call setdefault() in most cases, to avoid allocating and
+        # immediately freeing a lock; but still use setdefaut() to avoid
+        # races.
+        try:
+            x = self._init_once_cache[tag]
+        except KeyError:
+            x = self._init_once_cache.setdefault(tag, (False, allocate_lock()))
+        # Common case: we got (True, result), so we return the result.
+        if x[0]:
+            return x[1]
+        # Else, it's a lock.  Acquire it to serialize the following tests.
+        with x[1]:
+            # Read again from _init_once_cache the current status.
+            x = self._init_once_cache[tag]
+            if x[0]:
+                return x[1]
+            # Call the function and store the result back.
+            result = func()
+            self._init_once_cache[tag] = (True, result)
+        return result
+
 
 def _load_backend_lib(backend, name, flags):
     if name is None:
diff --git a/testing/cffi0/backend_tests.py b/testing/cffi0/backend_tests.py
index 23b925f..2a6af1b 100644
--- a/testing/cffi0/backend_tests.py
+++ b/testing/cffi0/backend_tests.py
@@ -1809,3 +1809,35 @@ class BackendTests:
         assert lib.EE1 == 0
         assert lib.EE2 == 0
         assert lib.EE3 == 1
+
+    def test_init_once(self):
+        def do_init():
+            seen.append(1)
+            return 42
+        ffi = FFI()
+        seen = []
+        for i in range(3):
+            res = ffi.init_once(do_init, "tag1")
+            assert res == 42
+            assert seen == [1]
+        for i in range(3):
+            res = ffi.init_once(do_init, "tag2")
+            assert res == 42
+            assert seen == [1, 1]
+
+    def test_init_once_multithread(self):
+        import thread, time
+        def do_init():
+            seen.append('init!')
+            time.sleep(1)
+            seen.append('init done')
+            return 7
+        ffi = FFI()
+        seen = []
+        for i in range(6):
+            def f():
+                res = ffi.init_once(do_init, "tag")
+                seen.append(res)
+            thread.start_new_thread(f, ())
+        time.sleep(1.5)
+        assert seen == ['init!', 'init done'] + 6 * [7]
diff --git a/testing/cffi1/test_ffi_obj.py b/testing/cffi1/test_ffi_obj.py
index 7cb67cc..354ad6e 100644
--- a/testing/cffi1/test_ffi_obj.py
+++ b/testing/cffi1/test_ffi_obj.py
@@ -415,3 +415,37 @@ def test_cast_from_int_type_to_bool():
             assert int(ffi.cast("_Bool", ffi.cast(type, 42))) == 1
             assert int(ffi.cast("bool", ffi.cast(type, 42))) == 1
             assert int(ffi.cast("_Bool", ffi.cast(type, 0))) == 0
+
+def test_init_once():
+    def do_init():
+        seen.append(1)
+        return 42
+    ffi = _cffi1_backend.FFI()
+    seen = []
+    for i in range(3):
+        res = ffi.init_once(do_init, "tag1")
+        assert res == 42
+        assert seen == [1]
+    for i in range(3):
+        res = ffi.init_once(do_init, "tag2")
+        assert res == 42
+        assert seen == [1, 1]
+
+def test_init_once_multithread():
+    import thread, time
+    def do_init():
+        print 'init!'
+        seen.append('init!')
+        time.sleep(1)
+        seen.append('init done')
+        print 'init done'
+        return 7
+    ffi = _cffi1_backend.FFI()
+    seen = []
+    for i in range(6):
+        def f():
+            res = ffi.init_once(do_init, "tag")
+            seen.append(res)
+        thread.start_new_thread(f, ())
+    time.sleep(1.5)
+    assert seen == ['init!', 'init done'] + 6 * [7]