1 files changed, 32 insertions, 18 deletions

diff --git a/src/zmq.cpp b/src/zmq.cpp
index 03064bc..1ed96df 100644
--- a/src/zmq.cpp
+++ b/src/zmq.cpp
@@ -324,27 +324,20 @@ int zmq_poll (zmq_pollitem_t *items_, int nitems_, long timeout_)
         npollfds++;
     }
 
+    //  First iteration just check for events, don't block. Waiting would
+    //  prevent exiting on any events that may already been signaled on
+    //  0MQ sockets.
+    int rc = poll (pollfds, npollfds, 0);
+    if (rc == -1 && errno == EINTR) {
+        free (pollfds);
+        return 0;
+    }
+    errno_assert (rc >= 0);
+
     int timeout = timeout_ > 0 ? timeout_ / 1000 : -1;
     int nevents = 0;
-    bool initial = true;
-    while (!nevents) {
-
-        //  Wait for activity. In the first iteration just check for events,
-        //  don't wait. Waiting would prevent exiting on any events that may
-        //  already be signaled on 0MQ sockets.
-        int rc = poll (pollfds, npollfds, initial ? 0 : timeout);
-        if (rc == -1 && errno == EINTR)
-            break;
-        errno_assert (rc >= 0);
-
-        //  If timeout was hit with no events signaled, return zero.
-        if (!initial && rc == 0) {
-            free (pollfds);
-            return 0;
-        }
 
-        //  From now on, perform blocking polling.
-        initial = false;
+    while (true) {
 
         //  Process 0MQ commands if needed.
         if (nsockets && pollfds [npollfds -1].revents & POLLIN)
@@ -376,6 +369,27 @@ int zmq_poll (zmq_pollitem_t *items_, int nitems_, long timeout_)
             if (items_ [i].revents)
                 nevents++;
         }
+
+        //  If there's at least one event, or if we are asked not to block,
+        //  return immediately.
+        if (nevents || !timeout)
+            break;
+
+        //  Wait for events.
+        rc = poll (pollfds, npollfds, timeout);
+        if (rc == -1 && errno == EINTR)
+            break;
+        errno_assert (rc >= 0);
+
+        //  If timeout was hit with no events signaled, return zero.
+        if (rc == 0)
+            break;
+
+        //  If timeout was already applied, we don't want to poll anymore.
+        //  Setting timeout to zero will cause termination of the function
+        //  once the events we've got are processed.
+        if (timeout > 0)
+            timeout = 0;
     }
 
     free (pollfds);