From de7fc1fcf850ab44eb33b6a285f2a455dde67e73 Mon Sep 17 00:00:00 2001
From: Min RK <benjaminrk@gmail.com>
Date: Tue, 27 Sep 2016 14:08:02 +0200
Subject: [PATCH] add n_events argument to zmq_poller_wait_all

avoids unnecessary heap allocations, races on the number of items
---
 include/zmq.h         |  2 +-
 src/socket_poller.cpp |  4 ++--
 src/socket_poller.hpp |  2 +-
 src/zmq.cpp           | 29 ++++++++++-------------------
 src/zmq_draft.h       |  2 +-
 5 files changed, 15 insertions(+), 24 deletions(-)

diff --git a/include/zmq.h b/include/zmq.h
index 108e8644..c9fe749f 100644
--- a/include/zmq.h
+++ b/include/zmq.h
@@ -577,7 +577,7 @@ ZMQ_EXPORT int  zmq_poller_add (void *poller, void *socket, void *user_data, sho
 ZMQ_EXPORT int  zmq_poller_modify (void *poller, void *socket, short events);
 ZMQ_EXPORT int  zmq_poller_remove (void *poller, void *socket);
 ZMQ_EXPORT int  zmq_poller_wait (void *poller, zmq_poller_event_t *event, long timeout);
-ZMQ_EXPORT int  zmq_poller_wait_all (void *poller, zmq_poller_event_t *events, long timeout);
+ZMQ_EXPORT int  zmq_poller_wait_all (void *poller, zmq_poller_event_t *events, int n_events, long timeout);
 
 #if defined _WIN32
 ZMQ_EXPORT int zmq_poller_add_fd (void *poller, SOCKET fd, void *user_data, short events);
diff --git a/src/socket_poller.cpp b/src/socket_poller.cpp
index 70a7fe1e..c4a720ef 100644
--- a/src/socket_poller.cpp
+++ b/src/socket_poller.cpp
@@ -380,7 +380,7 @@ int zmq::socket_poller_t::rebuild ()
     return 0;
 }
 
-int zmq::socket_poller_t::wait (zmq::socket_poller_t::event_t *events_, long timeout_)
+int zmq::socket_poller_t::wait (zmq::socket_poller_t::event_t *events_, int n_events_, long timeout_)
 {
     if (need_rebuild)
         if (rebuild () == -1)
@@ -441,7 +441,7 @@ int zmq::socket_poller_t::wait (zmq::socket_poller_t::event_t *events_, long tim
 
         //  Check for the events.
         int i=0;
-        for (items_t::iterator it = items.begin (); it != items.end (); ++i, ++it) {
+        for (items_t::iterator it = items.begin (); it != items.end () && i < n_events_; ++i, ++it) {
 
             events_[i].socket = NULL;
             events_[i].fd = 0;
diff --git a/src/socket_poller.hpp b/src/socket_poller.hpp
index 5fe43c06..46f4f9bf 100644
--- a/src/socket_poller.hpp
+++ b/src/socket_poller.hpp
@@ -73,7 +73,7 @@ namespace zmq
         int modify_fd (fd_t fd, short events);
         int remove_fd (fd_t fd);
 
-        int wait (event_t *event, long timeout);
+        int wait (event_t *event, int n_events, long timeout);
 
         inline int size (void) { return items.size (); };
 
diff --git a/src/zmq.cpp b/src/zmq.cpp
index ada3f7f9..c791258d 100644
--- a/src/zmq.cpp
+++ b/src/zmq.cpp
@@ -1222,46 +1222,37 @@ int zmq_poller_wait (void *poller_, zmq_poller_event_t *event, long timeout_)
 
     zmq_assert (event != NULL);
 
-    int n_items = ((zmq::socket_poller_t*)poller_)->size ();
-    zmq_poller_event_t *events;
-    events = new zmq_poller_event_t[n_items];
-    alloc_assert(events);
+    int rc = zmq_poller_wait_all(poller_, event, 1, timeout_);
 
-    int rc = zmq_poller_wait_all(poller_, events, timeout_);
-
-    if (rc >= 0) {
-        *event = events[0];
-    } else {
+    if (rc < 0) {
         memset (event, 0, sizeof(zmq_poller_event_t));
     }
-    delete [] events;
 
     return rc;
 }
 
-int zmq_poller_wait_all (void *poller_, zmq_poller_event_t *events, long timeout_)
+int zmq_poller_wait_all (void *poller_, zmq_poller_event_t *events, int n_events, long timeout_)
 {
     if (!poller_ || !((zmq::socket_poller_t*)poller_)->check_tag ()) {
         errno = EFAULT;
         return -1;
     }
+    if (n_events < 0) {
+        errno = EINVAL;
+        return -1;
+    }
 
     zmq_assert (events != NULL);
+    zmq::socket_poller_t::event_t evts[n_events];
 
-    int n_items = ((zmq::socket_poller_t*)poller_)->size ();
-    zmq::socket_poller_t::event_t *evts;
-    evts = new zmq::socket_poller_t::event_t[n_items];
-    alloc_assert(evts);
+    int rc = ((zmq::socket_poller_t*)poller_)->wait (evts, n_events, timeout_);
 
-    int rc = ((zmq::socket_poller_t*)poller_)->wait (evts, timeout_);
-
-    for(int i = 0; i < n_items; ++i) {
+    for(int i = 0; i < n_events; ++i) {
         events[i].socket = evts[i].socket;
         events[i].fd = evts[i].fd;
         events[i].user_data = evts[i].user_data;
         events[i].events = evts[i].events;
     }
-    delete [] evts;
 
     return rc;
 }
diff --git a/src/zmq_draft.h b/src/zmq_draft.h
index 907f0b83..343f8ef6 100644
--- a/src/zmq_draft.h
+++ b/src/zmq_draft.h
@@ -80,7 +80,7 @@ int  zmq_poller_add (void *poller, void *socket, void *user_data, short events);
 int  zmq_poller_modify (void *poller, void *socket, short events);
 int  zmq_poller_remove (void *poller, void *socket);
 int  zmq_poller_wait (void *poller, zmq_poller_event_t *event, long timeout);
-int  zmq_poller_wait_all (void *poller, zmq_poller_event_t *events, long timeout);
+int  zmq_poller_wait_all (void *poller, zmq_poller_event_t *events, int n_events, long timeout);
 
 #if defined _WIN32
 int zmq_poller_add_fd (void *poller, SOCKET fd, void *user_data, short events);