From dfb56ec1703b52cd62d0c37678d239e97ecba030 Mon Sep 17 00:00:00 2001
From: Thomas Vegas <tvegas@nvidia.com>
Date: Mon, 7 Oct 2024 09:49:13 +0300
Subject: [PATCH] UCS/ASYNC: Use heap to process expired timers in batch

---
 src/ucs/async/async.c  | 19 ++++++++++++++-----
 src/ucs/sys/compiler.h | 20 ++++++++++++++++++++
 2 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/src/ucs/async/async.c b/src/ucs/async/async.c
index f6be081231a..8e019ffc718 100644
--- a/src/ucs/async/async.c
+++ b/src/ucs/async/async.c
@@ -332,12 +332,19 @@ ucs_status_t ucs_async_dispatch_handlers(int *handler_ids, size_t count,
 ucs_status_t ucs_async_dispatch_timerq(ucs_timer_queue_t *timerq,
                                        ucs_time_t current_time)
 {
-    size_t max_timers, num_timers = 0;
+    size_t num_timers = 0;
+    ucs_status_t status;
+    size_t max_timers, size;
     int *expired_timers;
     ucs_timer_t *timer;
 
-    max_timers     = ucs_max(1, ucs_timerq_size(timerq));
-    expired_timers = ucs_alloca(max_timers * sizeof(*expired_timers));
+    max_timers = ucs_max(1, ucs_timerq_size(timerq));
+    size       = max_timers * sizeof(*expired_timers);
+
+    expired_timers = ucs_alloc_on_stack(size, "async_dispatch_timerq");
+    if (expired_timers == NULL) {
+        return UCS_ERR_NO_MEMORY;
+    }
 
     ucs_timerq_for_each_expired(timer, timerq, current_time, {
         expired_timers[num_timers++] = timer->id;
@@ -346,8 +353,10 @@ ucs_status_t ucs_async_dispatch_timerq(ucs_timer_queue_t *timerq,
         }
     })
 
-    return ucs_async_dispatch_handlers(expired_timers, num_timers,
-                                       UCS_ASYNC_EVENT_DUMMY);
+    status = ucs_async_dispatch_handlers(expired_timers, num_timers,
+                                         UCS_ASYNC_EVENT_DUMMY);
+    ucs_free_on_stack(expired_timers, size);
+    return status;
 }
 
 ucs_status_t ucs_async_context_init(ucs_async_context_t *async, ucs_async_mode_t mode)
diff --git a/src/ucs/sys/compiler.h b/src/ucs/sys/compiler.h
index 089bc36fb1c..9ebb8e89fa3 100644
--- a/src/ucs/sys/compiler.h
+++ b/src/ucs/sys/compiler.h
@@ -56,6 +56,26 @@
         } \
     }
 
+/**
+ * Use stack or heap allocation depending on size. Free function @ref
+ * ucs_free_on_stack must always be called.
+ */
+#define ucs_alloc_on_stack(_size, _name) \
+    ({ \
+        size_t _sz = (_size); \
+        (_sz <= UCS_ALLOCA_MAX_SIZE) ? alloca(_sz) : ucs_malloc(_sz, _name); \
+    })
+
+/**
+ * Release function to be paired with @ref ucs_alloc_on_stack.
+ */
+#define ucs_free_on_stack(_ptr, _size) \
+    ({ \
+        if ((_size) > UCS_ALLOCA_MAX_SIZE) { \
+            ucs_free(_ptr); \
+        } \
+    })
+
 /**
  * alloca which makes sure the size is small enough.
  */