From 49fa9573430db571b47d7c26781254311169e8e1 Mon Sep 17 00:00:00 2001
From: Mathias Buus <mathiasbuus@gmail.com>
Date: Mon, 9 Dec 2024 14:38:21 +0100
Subject: [PATCH] add idle hook and ref count lookups and events also (#231)

---
 include/udx.h         | 13 +++++++--
 src/udx.c             | 65 +++++++++++++++++++++++++++----------------
 test/lookup-invalid.c |  4 ++-
 test/lookup-ipv6.c    |  4 ++-
 test/lookup.c         |  4 ++-
 5 files changed, 61 insertions(+), 29 deletions(-)

diff --git a/include/udx.h b/include/udx.h
index cbf71fc6..c5c0afb9 100644
--- a/include/udx.h
+++ b/include/udx.h
@@ -83,6 +83,8 @@ typedef struct udx_lookup_s udx_lookup_t;
 
 typedef struct udx_interface_event_s udx_interface_event_t;
 
+typedef void (*udx_idle_cb)(udx_t *udx);
+
 typedef void (*udx_socket_send_cb)(udx_socket_send_t *req, int status);
 typedef void (*udx_socket_recv_cb)(udx_socket_t *socket, ssize_t read_len, const uv_buf_t *buf, const struct sockaddr *from);
 typedef void (*udx_socket_close_cb)(udx_socket_t *socket);
@@ -106,6 +108,7 @@ struct udx_s {
   uv_loop_t *loop;
 
   uint32_t refs;
+  udx_idle_cb on_idle;
 
   uint32_t sockets_len;
   uint32_t sockets_max_len;
@@ -375,6 +378,8 @@ struct udx_stream_send_s {
 
 struct udx_lookup_s {
   uv_getaddrinfo_t req;
+  udx_t *udx;
+
   struct addrinfo hints;
 
   udx_lookup_cb on_lookup;
@@ -385,6 +390,7 @@ struct udx_lookup_s {
 struct udx_interface_event_s {
   uv_timer_t timer;
   uv_loop_t *loop;
+  udx_t *udx;
 
   uv_interface_address_t *addrs;
   int addrs_len;
@@ -399,6 +405,9 @@ struct udx_interface_event_s {
 int
 udx_init (uv_loop_t *loop, udx_t *udx);
 
+void
+udx_idle (udx_t *udx, udx_idle_cb cb);
+
 int
 udx_socket_init (udx_t *udx, udx_socket_t *socket);
 
@@ -518,10 +527,10 @@ int
 udx_stream_destroy (udx_stream_t *stream);
 
 int
-udx_lookup (uv_loop_t *loop, udx_lookup_t *req, const char *host, unsigned int flags, udx_lookup_cb cb);
+udx_lookup (udx_t *udx, udx_lookup_t *req, const char *host, unsigned int flags, udx_lookup_cb cb);
 
 int
-udx_interface_event_init (uv_loop_t *loop, udx_interface_event_t *handle);
+udx_interface_event_init (udx_t *udx, udx_interface_event_t *handle);
 
 int
 udx_interface_event_start (udx_interface_event_t *handle, udx_interface_event_cb cb, uint64_t frequency);
diff --git a/src/udx.c b/src/udx.c
index 02e5875d..a96759d8 100644
--- a/src/udx.c
+++ b/src/udx.c
@@ -33,6 +33,8 @@
 #define UDX_DEFAULT_TTL         64
 #define UDX_DEFAULT_BUFFER_SIZE 212992
 
+#define UDX_DEFAULT_SET_SIZE 16
+
 #define UDX_MAX_RTO_TIMEOUTS 6
 
 #define UDX_CONG_C           400  // C=0.4 (inverse) in scaled 1000
@@ -151,24 +153,6 @@ send_window_in_packets (udx_stream_t *stream) {
 static void
 on_uv_poll (uv_poll_t *handle, int status, int events);
 
-static void
-ref_inc (udx_t *udx) {
-  udx->refs++;
-
-  if (udx->sockets == NULL) {
-    udx->sockets_len = 0;
-    udx->sockets_max_len = 16;
-    udx->sockets = malloc(udx->sockets_max_len * sizeof(udx_socket_t *));
-  }
-
-  if (udx->streams == NULL) {
-    udx->streams_len = 0;
-    udx->streams_max_len = 16;
-    udx->streams = malloc(udx->streams_max_len * sizeof(udx_stream_t *));
-    udx__cirbuf_init(&(udx->streams_by_id), 16);
-  }
-}
-
 static void
 ref_dec (udx_t *udx) {
   udx->refs--;
@@ -188,6 +172,10 @@ ref_dec (udx_t *udx) {
 
     udx__cirbuf_destroy(&(udx->streams_by_id));
   }
+
+  if (udx->on_idle != NULL) {
+    udx->on_idle(udx);
+  }
 }
 
 static void
@@ -1966,6 +1954,7 @@ on_uv_poll (uv_poll_t *handle, int status, int events) {
 int
 udx_init (uv_loop_t *loop, udx_t *udx) {
   udx->refs = 0;
+  udx->on_idle = NULL;
 
   udx->sockets_len = 0;
   udx->sockets_max_len = 0;
@@ -1986,9 +1975,20 @@ udx_init (uv_loop_t *loop, udx_t *udx) {
   return 0;
 }
 
+void
+udx_idle (udx_t *udx, udx_idle_cb cb) {
+  udx->on_idle = cb;
+}
+
 int
 udx_socket_init (udx_t *udx, udx_socket_t *socket) {
-  ref_inc(udx);
+  udx->refs++;
+
+  if (udx->sockets == NULL) {
+    udx->sockets_len = 0;
+    udx->sockets_max_len = UDX_DEFAULT_SET_SIZE;
+    udx->sockets = malloc(udx->sockets_max_len * sizeof(udx_socket_t *));
+  }
 
   socket->family = 0;
   socket->status = 0;
@@ -2259,7 +2259,14 @@ udx_socket_close (udx_socket_t *socket, udx_socket_close_cb cb) {
 
 int
 udx_stream_init (udx_t *udx, udx_stream_t *stream, uint32_t local_id, udx_stream_close_cb close_cb, udx_stream_finalize_cb finalize_cb) {
-  ref_inc(udx);
+  udx->refs++;
+
+  if (udx->streams == NULL) {
+    udx->streams_len = 0;
+    udx->streams_max_len = UDX_DEFAULT_SET_SIZE;
+    udx->streams = malloc(udx->streams_max_len * sizeof(udx_stream_t *));
+    udx__cirbuf_init(&(udx->streams_by_id), 16);
+  }
 
   stream->local_id = local_id;
   stream->remote_id = 0;
@@ -2743,13 +2750,18 @@ on_uv_getaddrinfo (uv_getaddrinfo_t *req, int status, struct addrinfo *res) {
   }
 
   uv_freeaddrinfo(res);
+
+  ref_dec(lookup->udx);
 }
 
 int
-udx_lookup (uv_loop_t *loop, udx_lookup_t *req, const char *host, unsigned int flags, udx_lookup_cb cb) {
+udx_lookup (udx_t *udx, udx_lookup_t *req, const char *host, unsigned int flags, udx_lookup_cb cb) {
+  req->udx = udx;
   req->on_lookup = cb;
   req->req.data = req;
 
+  udx->refs++;
+
   memset(&req->hints, 0, sizeof(struct addrinfo));
 
   int family = AF_UNSPEC;
@@ -2760,7 +2772,7 @@ udx_lookup (uv_loop_t *loop, udx_lookup_t *req, const char *host, unsigned int f
   req->hints.ai_family = family;
   req->hints.ai_socktype = SOCK_STREAM;
 
-  return uv_getaddrinfo(loop, &req->req, on_uv_getaddrinfo, host, NULL, &req->hints);
+  return uv_getaddrinfo(udx->loop, &req->req, on_uv_getaddrinfo, host, NULL, &req->hints);
 }
 
 static int
@@ -2828,13 +2840,18 @@ on_interface_event_close (uv_handle_t *handle) {
   if (event->on_close != NULL) {
     event->on_close(event);
   }
+
+  ref_dec(event->udx);
 }
 
 int
-udx_interface_event_init (uv_loop_t *loop, udx_interface_event_t *handle) {
-  handle->loop = loop;
+udx_interface_event_init (udx_t *udx, udx_interface_event_t *handle) {
+  handle->udx = udx;
+  handle->loop = udx->loop;
   handle->sorted = false;
 
+  udx->refs++;
+
   int err = uv_interface_addresses(&(handle->addrs), &(handle->addrs_len));
   if (err < 0) return err;
 
diff --git a/test/lookup-invalid.c b/test/lookup-invalid.c
index 57cb95a1..22d3ba46 100644
--- a/test/lookup-invalid.c
+++ b/test/lookup-invalid.c
@@ -4,6 +4,7 @@
 
 #include "../include/udx.h"
 
+udx_t udx;
 uv_loop_t loop;
 udx_lookup_t req;
 
@@ -21,8 +22,9 @@ main () {
   int e;
 
   uv_loop_init(&loop);
+  udx_init(&loop, &udx);
 
-  e = udx_lookup(&loop, &req, "example.invalid.", 0, on_lookup);
+  e = udx_lookup(&udx, &req, "example.invalid.", 0, on_lookup);
   assert(e == 0);
 
   uv_run(&loop, UV_RUN_DEFAULT);
diff --git a/test/lookup-ipv6.c b/test/lookup-ipv6.c
index a579d575..5bf75498 100644
--- a/test/lookup-ipv6.c
+++ b/test/lookup-ipv6.c
@@ -5,6 +5,7 @@
 #include "../include/udx.h"
 
 uv_loop_t loop;
+udx_t udx;
 udx_lookup_t req;
 
 void
@@ -26,8 +27,9 @@ main () {
   int e;
 
   uv_loop_init(&loop);
+  udx_init(&loop, &udx);
 
-  e = udx_lookup(&loop, &req, "localhost", UDX_LOOKUP_FAMILY_IPV6, on_lookup);
+  e = udx_lookup(&udx, &req, "localhost", UDX_LOOKUP_FAMILY_IPV6, on_lookup);
   assert(e == 0);
 
   uv_run(&loop, UV_RUN_DEFAULT);
diff --git a/test/lookup.c b/test/lookup.c
index aa349c83..7fe4fe87 100644
--- a/test/lookup.c
+++ b/test/lookup.c
@@ -5,6 +5,7 @@
 #include "../include/udx.h"
 
 uv_loop_t loop;
+udx_t udx;
 udx_lookup_t req;
 
 void
@@ -26,8 +27,9 @@ main () {
   int e;
 
   uv_loop_init(&loop);
+  udx_init(&loop, &udx);
 
-  e = udx_lookup(&loop, &req, "localhost", UDX_LOOKUP_FAMILY_IPV4, on_lookup);
+  e = udx_lookup(&udx, &req, "localhost", UDX_LOOKUP_FAMILY_IPV4, on_lookup);
   assert(e == 0);
 
   uv_run(&loop, UV_RUN_DEFAULT);