Merge pull request #708 from fsbs/patch-allow-threads

Fix multi callback issues: easy.pause(), easy.close(), multi.close()

5 files changed, 126 insertions, 4 deletions

diff --git a/src/easy.c b/src/easy.c
index a22ad9c..89d159d 100644
--- a/src/easy.c
+++ b/src/easy.c
@@ -137,14 +137,17 @@ util_curl_xdecref(CurlObject *self, int flags, CURL *handle)
         /* Decrement refcount for multi_stack. */
         if (self->multi_stack != NULL) {
             CurlMultiObject *multi_stack = self->multi_stack;
-            self->multi_stack = NULL;
             if (multi_stack->multi_handle != NULL && handle != NULL) {
                 /* TODO this is where we could remove the easy object
                 from the multi object's easy_object_dict, but this
                 requires us to have a reference to the multi object
                 which right now we don't. */
+                /* Allow threads because callbacks can be invoked */
+                PYCURL_BEGIN_ALLOW_THREADS_EASY
                 (void) curl_multi_remove_handle(multi_stack->multi_handle, handle);
+                PYCURL_END_ALLOW_THREADS_EASY
             }
+            self->multi_stack = NULL;
             Py_DECREF(multi_stack);
         }
     }
diff --git a/src/easyperform.c b/src/easyperform.c
index 44f21c0..5326df3 100644
--- a/src/easyperform.c
+++ b/src/easyperform.c
@@ -100,7 +100,7 @@ do_curl_pause(CurlObject *self, PyObject *args)
 #ifdef WITH_THREAD
     /* Save handle to current thread (used as context for python callbacks) */
     saved_state = self->state;
-    PYCURL_BEGIN_ALLOW_THREADS
+    PYCURL_BEGIN_ALLOW_THREADS_EASY
 
     /* We must allow threads here because unpausing a handle can cause
        some of its callbacks to be invoked immediately, from inside
@@ -110,7 +110,7 @@ do_curl_pause(CurlObject *self, PyObject *args)
     res = curl_easy_pause(self->handle, bitmask);
 
 #ifdef WITH_THREAD
-    PYCURL_END_ALLOW_THREADS
+    PYCURL_END_ALLOW_THREADS_EASY
 
     /* Restore the thread-state to whatever it was on entry */
     self->state = saved_state;
diff --git a/src/multi.c b/src/multi.c
index 1ec5e19..7162564 100644
--- a/src/multi.c
+++ b/src/multi.c
@@ -94,11 +94,11 @@ util_multi_close(CurlMultiObject *self)
     
     if (self->multi_handle != NULL) {
         CURLM *multi_handle = self->multi_handle;
-        self->multi_handle = NULL;
         /* Allow threads because callbacks can be invoked */
         PYCURL_BEGIN_ALLOW_THREADS
         curl_multi_cleanup(multi_handle);
         PYCURL_END_ALLOW_THREADS
+        self->multi_handle = NULL;
     }
 }
 
diff --git a/src/pycurl.h b/src/pycurl.h
index 42c7d57..1e3fff3 100644
--- a/src/pycurl.h
+++ b/src/pycurl.h
@@ -266,6 +266,19 @@ PYCURL_INTERNAL void pycurl_ssl_cleanup(void);
 #  define PYCURL_END_ALLOW_THREADS \
        Py_END_ALLOW_THREADS \
        self->state = NULL;
+#  define PYCURL_BEGIN_ALLOW_THREADS_EASY \
+       if (self->multi_stack == NULL) { \
+           self->state = PyThreadState_Get(); \
+           assert(self->state != NULL); \
+       } else { \
+           self->multi_stack->state = PyThreadState_Get(); \
+           assert(self->multi_stack->state != NULL); \
+       } \
+       Py_BEGIN_ALLOW_THREADS
+#  define PYCURL_END_ALLOW_THREADS_EASY \
+       PYCURL_END_ALLOW_THREADS \
+       if (self->multi_stack != NULL) \
+           self->multi_stack->state = NULL;
 #else
 #  define PYCURL_DECLARE_THREAD_STATE
 #  define PYCURL_ACQUIRE_THREAD() (1)
diff --git a/tests/multi_callback_test.py b/tests/multi_callback_test.py
new file mode 100644
index 0000000..c82bc38
--- /dev/null
+++ b/tests/multi_callback_test.py
@@ -0,0 +1,106 @@
+#! /usr/bin/env python
+# -*- coding: utf-8 -*-
+# vi:ts=4:et
+
+from . import localhost
+import pycurl
+import unittest
+
+from . import appmanager
+from . import util
+
+setup_module, teardown_module = appmanager.setup(('app', 8380))
+
+class MultiCallbackTest(unittest.TestCase):
+    def setUp(self):
+        self.easy = util.DefaultCurl()
+        self.easy.setopt(pycurl.URL, 'http://%s:8380/success' % localhost)
+        self.multi = pycurl.CurlMulti()
+        self.multi.setopt(pycurl.M_SOCKETFUNCTION, self.socket_callback)
+        self.multi.setopt(pycurl.M_TIMERFUNCTION, self.timer_callback)
+        self.timer_result = None
+        self.socket_result = None
+        self.socket_action = None
+
+    def tearDown(self):
+        self.multi.close()
+        self.easy.close()
+
+    def socket_callback(self, ev_bitmask, sock_fd, multi, data):
+        self.socket_result = (sock_fd, ev_bitmask)
+        if ev_bitmask & pycurl.POLL_REMOVE:
+            pass
+        else:
+            self.socket_action = (sock_fd, ev_bitmask)
+
+    def timer_callback(self, timeout_ms):
+        self.timer_result = timeout_ms
+        self.socket_action = (pycurl.SOCKET_TIMEOUT, 0)
+
+    # multi.socket_action must call both SOCKETFUNCTION and TIMERFUNCTION at
+    # various points during the transfer (at least at the start and end)
+    def test_multi_socket_action(self):
+        self.multi.add_handle(self.easy)
+        self.timer_result = None
+        self.socket_result = None
+        while self.multi.socket_action(*self.socket_action)[1]:
+            # Without real event loop we just use blocking select call instead
+            self.multi.select(0.1)
+        # both callbacks should be invoked multiple times by socket_action
+        assert self.socket_result is not None
+        assert self.timer_result is not None
+
+    # multi.add_handle must call TIMERFUNCTION to schedule a kick-start
+    def test_multi_add_handle(self):
+        assert self.timer_result == None
+        self.multi.add_handle(self.easy)
+        assert self.timer_result is not None
+
+    # (mid-transfer) multi.remove_handle must call SOCKETFUNCTION to remove sockets
+    def test_multi_remove_handle(self):
+        self.multi.add_handle(self.easy)
+        while self.multi.socket_action(*self.socket_action)[1]:
+            if self.socket_result:
+                # libcurl informed us about new sockets
+                break
+            # Without real event loop we just use blocking select call instead
+            self.multi.select(0.1)
+        self.socket_result = None
+        # libcurl will now inform us that we should remove those sockets
+        self.multi.remove_handle(self.easy)
+        assert self.socket_result is not None
+    
+    # (mid-transfer) easy.pause(PAUSE_ALL) must call SOCKETFUNCTION to remove sockets
+    # (mid-transfer) easy.pause(PAUSE_CONT) must call TIMERFUNCTION to resume
+    def test_easy_pause_unpause(self):
+        self.multi.add_handle(self.easy)
+        while self.multi.socket_action(*self.socket_action)[1]:
+            if self.socket_result:
+                # libcurl informed us about new sockets
+                break
+            # Without real event loop we just use blocking select call instead
+            self.multi.select(0.1)
+        self.socket_result = None
+        # libcurl will now inform us that we should remove those sockets
+        self.easy.pause(pycurl.PAUSE_ALL)
+        assert self.socket_result is not None
+        self.timer_result = None
+        # libcurl will now tell us to schedule a kickstart
+        self.easy.pause(pycurl.PAUSE_CONT)
+        assert self.timer_result is not None
+
+    # (mid-transfer) easy.close() must call SOCKETFUNCTION to remove sockets
+    # NOTE: doing easy.close() during transfer is considered wrong,
+    # but libcurl still invokes callbacks to inform your event loop
+    def test_easy_close(self):
+        self.multi.add_handle(self.easy)
+        while self.multi.socket_action(*self.socket_action)[1]:
+            if self.socket_result:
+                # libcurl informed us about new sockets
+                break
+            # Without real event loop we just use blocking select call instead
+            self.multi.select(0.1)
+        self.socket_result = None
+        # libcurl will now inform us that we should remove those sockets
+        self.easy.close()
+        assert self.socket_result is not None