UCS/CALLBACKQ: fix recursive one-shot

openucx · Nov 19, 2020 · 1fb7e41 · 1fb7e41
1 parent 499581e
commit 1fb7e41
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 5 deletions.
diff --git a/src/ucs/datastruct/callbackq.c b/src/ucs/datastruct/callbackq.c
@@ -363,18 +363,20 @@ static unsigned ucs_callbackq_slow_proxy(void *arg)
 {
     ucs_callbackq_t      *cbq  = arg;
     ucs_callbackq_priv_t *priv = ucs_callbackq_priv(cbq);
+    unsigned num_slow_elems    = priv->num_slow_elems;
+    unsigned count             = 0;
     ucs_callbackq_elem_t *elem;
     unsigned UCS_V_UNUSED removed_idx;
     unsigned slow_idx, fast_idx;
     ucs_callbackq_elem_t tmp_elem;
-    unsigned count = 0;
 
     ucs_trace_poll("cbq=%p", cbq);
 
     ucs_callbackq_enter(cbq);
 
-    /* Execute and update slow-path callbacks */
-    for (slow_idx = 0; slow_idx < priv->num_slow_elems; ++slow_idx) {
+    /* Execute and update slow-path callbacks by num_slow_elems copy to avoid
+     * infinite loop if callback adds another one */
+    for (slow_idx = 0; slow_idx < num_slow_elems; ++slow_idx) {
         elem = &priv->slow_elems[slow_idx];
         if (elem->id == UCS_CALLBACKQ_ID_NULL) {
             continue;

diff --git a/test/gtest/ucs/test_callbackq.cc b/test/gtest/ucs/test_callbackq.cc
@@ -32,6 +32,7 @@ class test_callbackq :
         uint32_t                  count;
         int                       command;
         callback_ctx              *to_add;
+        unsigned                  flags;
         int                       key;
     };
 
@@ -81,9 +82,11 @@ class test_callbackq :
     void init_ctx(callback_ctx *ctx, int key = 0)
     {
         ctx->test        = this;
+        ctx->callback_id = UCS_CALLBACKQ_ID_NULL;
         ctx->count       = 0;
         ctx->command     = COMMAND_NONE;
-        ctx->callback_id = UCS_CALLBACKQ_ID_NULL;
+        ctx->to_add      = NULL;
+        ctx->flags       = 0;
         ctx->key         = key;
     }
 
@@ -95,7 +98,7 @@ class test_callbackq :
     {
         ctx->callback_id = ucs_callbackq_add(&m_cbq, callback_proxy,
                                              reinterpret_cast<void*>(ctx),
-                                             cb_flags() | flags);
+                                             ctx->flags | cb_flags() | flags);
     }
 
     void remove(int callback_id)
@@ -337,6 +340,24 @@ UCS_TEST_F(test_callbackq_noflags, oneshot) {
     EXPECT_EQ(1u, ctx.count);
 }
 
+UCS_TEST_F(test_callbackq_noflags, oneshot_recursive) {
+    callback_ctx ctx;
+
+    init_ctx(&ctx);
+    ctx.command = COMMAND_ADD_ANOTHER;
+    ctx.flags   = UCS_CALLBACKQ_FLAG_ONESHOT;
+    ctx.to_add  = &ctx;
+
+    add(&ctx);
+
+    for (unsigned i = 0; i < 10; ++i) {
+        dispatch(1);
+        EXPECT_LE(i + 1, ctx.count);
+    }
+
+    remove(ctx.callback_id);
+}
+
 UCS_TEST_F(test_callbackq_noflags, remove_if) {
     const size_t count = 1000;
     const int num_keys = 10;