From c20df2b20d018fef79fb7d28f4c1862f010824b3 Mon Sep 17 00:00:00 2001
From: Mikhail Brinskii <mikhailb@nvidia.com>
Date: Thu, 12 Aug 2021 18:24:26 +0300
Subject: [PATCH] UCP/TAG: Fix mt bug in probe

---
 src/ucp/tag/probe.c | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/ucp/tag/probe.c b/src/ucp/tag/probe.c
index 971c57fd1f6..093e826e045 100644
--- a/src/ucp/tag/probe.c
+++ b/src/ucp/tag/probe.c
@@ -28,6 +28,7 @@ UCS_PROFILE_FUNC(ucp_tag_message_h, ucp_tag_probe_nb,
 
     UCP_CONTEXT_CHECK_FEATURE_FLAGS(worker->context, UCP_FEATURE_TAG,
                                     return NULL);
+    UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(worker);
 
     ucs_trace_req("probe_nb tag %"PRIx64"/%"PRIx64" remove=%d", tag, tag_mask,
                   rem);
@@ -47,6 +48,7 @@ UCS_PROFILE_FUNC(ucp_tag_message_h, ucp_tag_probe_nb,
         }
     }
 
+    UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(worker);
 
     return rdesc;
 }