udp: add uv_udp_using_recvmmsg query

Allows for determining if a buffer large enough for multiple dgrams should be allocated in alloc_cb of uv_udp_recvstart, for example. Contributes towards libuv#2822. PR-URL: libuv#2830 Reviewed-By: Ben Noordhuis <info@bnoordhuis.nl> Reviewed-By: Jameson Nash <vtjnash@gmail.com>
escherstair · Jul 29, 2020 · 6b5aa66 · 6b5aa66
1 parent 540d723
commit 6b5aa66
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 12 deletions.
diff --git a/docs/src/udp.rst b/docs/src/udp.rst
@@ -391,6 +391,16 @@ API
     .. versionchanged:: 1.37.0 :man:`recvmmsg(2)` support is no longer enabled implicitly,
                         it must be explicitly requested by passing the `UV_UDP_RECVMMSG` flag to
                         :c:func:`uv_udp_init_ex`.
+    .. versionchanged:: 1.39.0 :c:func:`uv_udp_using_recvmmsg` can be used in `alloc_cb` to
+                        determine if a buffer sized for use with :man:`recvmmsg(2)` should be 
+                        allocated for the current handle/platform.
+
+.. c:function:: int uv_udp_using_recvmmsg(uv_udp_t* handle)
+
+    Returns 1 if the UDP handle was created with the `UV_UDP_RECVMMSG` flag
+    and the platform supports :man:`recvmmsg(2)`, 0 otherwise.
+
+    .. versionadded:: 1.39.0
 
 .. c:function:: int uv_udp_recv_stop(uv_udp_t* handle)
 

diff --git a/include/uv.h b/include/uv.h
@@ -693,6 +693,7 @@ UV_EXTERN int uv_udp_try_send(uv_udp_t* handle,
 UV_EXTERN int uv_udp_recv_start(uv_udp_t* handle,
                                 uv_alloc_cb alloc_cb,
                                 uv_udp_recv_cb recv_cb);
+UV_EXTERN int uv_udp_using_recvmmsg(const uv_udp_t* handle);
 UV_EXTERN int uv_udp_recv_stop(uv_udp_t* handle);
 UV_EXTERN size_t uv_udp_get_send_queue_size(const uv_udp_t* handle);
 UV_EXTERN size_t uv_udp_get_send_queue_count(const uv_udp_t* handle);

diff --git a/src/unix/udp.c b/src/unix/udp.c
@@ -270,14 +270,11 @@ static void uv__udp_recvmsg(uv_udp_t* handle) {
     assert(buf.base != NULL);
 
 #if HAVE_MMSG
-    if (handle->flags & UV_HANDLE_UDP_RECVMMSG) {
-      uv_once(&once, uv__udp_mmsg_init);
-      if (uv__recvmmsg_avail) {
-        nread = uv__udp_recvmmsg(handle, &buf);
-        if (nread > 0)
-          count -= nread;
-        continue;
-      }
+    if (uv_udp_using_recvmmsg(handle)) {
+      nread = uv__udp_recvmmsg(handle, &buf);
+      if (nread > 0)
+        count -= nread;
+      continue;
     }
 #endif
 
@@ -976,6 +973,17 @@ int uv__udp_init_ex(uv_loop_t* loop,
 }
 
 
+int uv_udp_using_recvmmsg(const uv_udp_t* handle) {
+#if HAVE_MMSG
+  if (handle->flags & UV_HANDLE_UDP_RECVMMSG) {
+    uv_once(&once, uv__udp_mmsg_init);
+    return uv__recvmmsg_avail;
+  }
+#endif
+  return 0;
+}
+
+
 int uv_udp_open(uv_udp_t* handle, uv_os_sock_t sock) {
   int err;
 

diff --git a/src/win/udp.c b/src/win/udp.c
@@ -189,6 +189,11 @@ void uv_udp_endgame(uv_loop_t* loop, uv_udp_t* handle) {
 }
 
 
+int uv_udp_using_recvmmsg(const uv_udp_t* handle) {
+  return 0;
+}
+
+
 static int uv_udp_maybe_bind(uv_udp_t* handle,
                              const struct sockaddr* addr,
                              unsigned int addrlen,

diff --git a/test/test-udp-mmsg.c b/test/test-udp-mmsg.c
@@ -30,7 +30,7 @@
   ASSERT((uv_udp_t*)(handle) == &recver || (uv_udp_t*)(handle) == &sender)
 
 #define BUFFER_MULTIPLIER 4
-#define BUFFER_SIZE (BUFFER_MULTIPLIER * 64 * 1024)
+#define MAX_DGRAM_SIZE (64 * 1024)
 #define NUM_SENDS 8
 #define EXPECTED_MMSG_ALLOCS (NUM_SENDS / BUFFER_MULTIPLIER)
 
@@ -44,11 +44,18 @@ static int alloc_cb_called;
 static void alloc_cb(uv_handle_t* handle,
                      size_t suggested_size,
                      uv_buf_t* buf) {
+  size_t buffer_size;
   CHECK_HANDLE(handle);
+
+  /* Only allocate enough room for multiple dgrams if we can actually recv them */
+  buffer_size = MAX_DGRAM_SIZE;
+  if (uv_udp_using_recvmmsg((uv_udp_t*)handle))
+    buffer_size *= BUFFER_MULTIPLIER;
+
   /* Actually malloc to exercise free'ing the buffer later */
-  buf->base = malloc(BUFFER_SIZE);
+  buf->base = malloc(buffer_size);
   ASSERT(buf->base != NULL);
-  buf->len = BUFFER_SIZE;
+  buf->len = buffer_size;
   alloc_cb_called++;
 }
 
@@ -119,7 +126,10 @@ TEST_IMPL(udp_mmsg) {
   printf("%d allocs for %d recvs\n", alloc_cb_called, recv_cb_called);
 
   /* On platforms that don't support mmsg, each recv gets its own alloc */
-  ASSERT(alloc_cb_called == EXPECTED_MMSG_ALLOCS || alloc_cb_called == recv_cb_called);
+  if (uv_udp_using_recvmmsg(&recver))
+    ASSERT_EQ(alloc_cb_called, EXPECTED_MMSG_ALLOCS);
+  else
+    ASSERT_EQ(alloc_cb_called, recv_cb_called);
 
   MAKE_VALGRIND_HAPPY();
   return 0;