GTEST: Synch last_wqe callback with main thread

openucx · Jun 10, 2021 · d492e5c · d492e5c
1 parent 2b81e34
commit d492e5c
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/test/gtest/uct/ib/test_ib_event.cc b/test/gtest/uct/ib/test_ib_event.cc
@@ -57,13 +57,20 @@ class uct_test_event_base : public uct_p2p_test {
     struct event_ctx {
         uct_ib_async_event_wait_t super;
         volatile bool             got;
+        uct_ib_device_t           *dev;
     };
 
     static unsigned last_wqe_check_cb(void *arg) {
         event_ctx *event = (event_ctx *)arg;
-        event->got       = true;
-        ucs_callbackq_remove_safe(event->super.cbq, event->super.cb_id);
+        int cb_id;
+
+        ucs_spin_lock(&event->dev->async_event_lock);
+        cb_id              = event->super.cb_id;
         event->super.cb_id = UCS_CALLBACKQ_ID_NULL;
+        ucs_spin_unlock(&event->dev->async_event_lock);
+
+        event->got = true;
+        ucs_callbackq_remove_safe(event->super.cbq, cb_id);
         return 1;
     }
 
@@ -81,6 +88,7 @@ class uct_test_event_base : public uct_p2p_test {
         m_event.got       = false;
         m_event.super.cb  = last_wqe_check_cb;
         m_event.super.cbq = &e.worker()->progress_q;
+        m_event.dev       = dev(e);
 
         if (before) {
             /* move QP to error state before scheduling event callback */