Execute user's callback after setting request status

build_and_run.sh

@@ -156,7 +156,9 @@ run_py_benchmark() {
 }
 
 if [[ $RUN_CPP_TESTS != 0 ]]; then
-  ${BINARY_PATH}/gtests/libucxx/UCXX_TEST
+  # UCX_TCP_CM_REUSEADDR=y to be able to bind immediately to the same port before
+  # `TIME_WAIT` timeout
+  UCX_TCP_CM_REUSEADDR=y ${BINARY_PATH}/gtests/libucxx/UCXX_TEST
 fi
 if [[ $RUN_CPP_BENCH != 0 ]]; then
   # run_cpp_benchmark PROGRESS_MODE

ci/test_cpp.sh

@@ -25,11 +25,11 @@ print_system_stats
 BINARY_PATH=${CONDA_PREFIX}/bin
 
 run_tests() {
-  CMD_LINE="UCX_TCP_CM_REUSEADDR=y timeout 10m ${BINARY_PATH}/gtests/libucxx/UCXX_TEST --gtest_filter=-*DelayedSubmission*ProgressTagMulti*:ListenerTest.CloseCallback:ListenerTest.IsAlive:ListenerTest.RaiseOnError"
+  CMD_LINE="UCX_TCP_CM_REUSEADDR=y timeout 10m ${BINARY_PATH}/gtests/libucxx/UCXX_TEST --gtest_filter=-ListenerTest.CloseCallback:ListenerTest.IsAlive:ListenerTest.RaiseOnError"
 
   rapids-logger "Running: \n  - ${CMD_LINE}"
 
-  UCX_TCP_CM_REUSEADDR=y timeout 10m ${BINARY_PATH}/gtests/libucxx/UCXX_TEST --gtest_filter=-*DelayedSubmission*ProgressTagMulti*:ListenerTest.CloseCallback:ListenerTest.IsAlive:ListenerTest.RaiseOnError
+  UCX_TCP_CM_REUSEADDR=y timeout 10m ${BINARY_PATH}/gtests/libucxx/UCXX_TEST --gtest_filter=-ListenerTest.CloseCallback:ListenerTest.IsAlive:ListenerTest.RaiseOnError
 }
 
 run_benchmark() {

cpp/include/ucxx/constructors.h

@@ -59,15 +59,14 @@ std::shared_ptr<RequestStream> createRequestStream(std::shared_ptr<Endpoint> end
                                                    size_t length,
                                                    const bool enablePythonFuture);
 
-std::shared_ptr<RequestTag> createRequestTag(
-  std::shared_ptr<Component> endpointOrWorker,
-  bool send,
-  void* buffer,
-  size_t length,
-  ucp_tag_t tag,
-  const bool enablePythonFuture,
-  std::function<void(std::shared_ptr<void>)> callbackFunction,
-  std::shared_ptr<void> callbackData);
+std::shared_ptr<RequestTag> createRequestTag(std::shared_ptr<Component> endpointOrWorker,
+                                             bool send,
+                                             void* buffer,
+                                             size_t length,
+                                             ucp_tag_t tag,
+                                             const bool enablePythonFuture,
+                                             RequestCallbackUserFunction callbackFunction,
+                                             RequestCallbackUserData callbackData);
 
 std::shared_ptr<RequestTagMulti> createRequestTagMultiSend(std::shared_ptr<Endpoint> endpoint,
                                                            const std::vector<void*>& buffer,

cpp/include/ucxx/endpoint.h

@@ -319,13 +319,12 @@ class Endpoint : public Component {
    *
    * @returns Request to be subsequently checked for the completion and its state.
    */
-  std::shared_ptr<Request> tagSend(
-    void* buffer,
-    size_t length,
-    ucp_tag_t tag,
-    const bool enablePythonFuture                               = false,
-    std::function<void(std::shared_ptr<void>)> callbackFunction = nullptr,
-    std::shared_ptr<void> callbackData                          = nullptr);
+  std::shared_ptr<Request> tagSend(void* buffer,
+                                   size_t length,
+                                   ucp_tag_t tag,
+                                   const bool enablePythonFuture                = false,
+                                   RequestCallbackUserFunction callbackFunction = nullptr,
+                                   RequestCallbackUserData callbackData         = nullptr);
 
   /**
    * @brief Enqueue a tag receive operation.
@@ -350,13 +349,12 @@ class Endpoint : public Component {
    *
    * @returns Request to be subsequently checked for the completion and its state.
    */
-  std::shared_ptr<Request> tagRecv(
-    void* buffer,
-    size_t length,
-    ucp_tag_t tag,
-    const bool enablePythonFuture                               = false,
-    std::function<void(std::shared_ptr<void>)> callbackFunction = nullptr,
-    std::shared_ptr<void> callbackData                          = nullptr);
+  std::shared_ptr<Request> tagRecv(void* buffer,
+                                   size_t length,
+                                   ucp_tag_t tag,
+                                   const bool enablePythonFuture                = false,
+                                   RequestCallbackUserFunction callbackFunction = nullptr,
+                                   RequestCallbackUserData callbackData         = nullptr);
 
   /**
    * @brief Enqueue a multi-buffer tag send operation.

cpp/include/ucxx/request.h

@@ -27,8 +27,8 @@ class Request : public Component {
   std::string _status_msg{};                          ///< Human-readable status message
   void* _request{nullptr};                            ///< Pointer to UCP request
   std::shared_ptr<Future> _future{nullptr};           ///< Future to notify upon completion
-  std::function<void(std::shared_ptr<void>)> _callback{nullptr};  ///< Completion callback
-  std::shared_ptr<void> _callbackData{nullptr};                   ///< Completion callback data
+  RequestCallbackUserFunction _callback{nullptr};     ///< Completion callback
+  RequestCallbackUserData _callbackData{nullptr};     ///< Completion callback data
   std::shared_ptr<Worker> _worker{
     nullptr};  ///< Worker that generated request (if not from endpoint)
   std::shared_ptr<Endpoint> _endpoint{

cpp/include/ucxx/request_tag.h

@@ -53,9 +53,9 @@ class RequestTag : public Request {
              void* buffer,
              size_t length,
              ucp_tag_t tag,
-             const bool enablePythonFuture                               = false,
-             std::function<void(std::shared_ptr<void>)> callbackFunction = nullptr,
-             std::shared_ptr<void> callbackData                          = nullptr);
+             const bool enablePythonFuture                = false,
+             RequestCallbackUserFunction callbackFunction = nullptr,
+             RequestCallbackUserData callbackData         = nullptr);
 
  public:
   /**
@@ -85,15 +85,14 @@ class RequestTag : public Request {
    *
    * @returns The `shared_ptr<ucxx::RequestTag>` object
    */
-  friend std::shared_ptr<RequestTag> createRequestTag(
-    std::shared_ptr<Component> endpointOrWorker,
-    bool send,
-    void* buffer,
-    size_t length,
-    ucp_tag_t tag,
-    const bool enablePythonFuture,
-    std::function<void(std::shared_ptr<void>)> callbackFunction,
-    std::shared_ptr<void> callbackData);
+  friend std::shared_ptr<RequestTag> createRequestTag(std::shared_ptr<Component> endpointOrWorker,
+                                                      bool send,
+                                                      void* buffer,
+                                                      size_t length,
+                                                      ucp_tag_t tag,
+                                                      const bool enablePythonFuture,
+                                                      RequestCallbackUserFunction callbackFunction,
+                                                      RequestCallbackUserData callbackData);
 
   virtual void populateDelayedSubmission();
 

cpp/include/ucxx/request_tag_multi.h

@@ -216,9 +216,10 @@ class RequestTagMulti : public std::enable_shared_from_this<RequestTagMulti> {
    * which will be later used to evaluate if all frames completed and set the final status
    * of the multi-transfer request and the Python future, if enabled.
    *
+   * @param[in] status the status of the request being completed.
    * @param[in] request the `ucxx::BufferRequest` object containing a single tag .
    */
-  void markCompleted(std::shared_ptr<void> request);
+  void markCompleted(ucs_status_t status, RequestCallbackUserData request);
 
   /**
    * @brief Callback to submit request to receive new header or frames.
@@ -231,9 +232,12 @@ class RequestTagMulti : public std::enable_shared_from_this<RequestTagMulti> {
    * containing the `next` flag set, then the next request is another header. Otherwise, the
    * next incoming message(s) is(are) frame(s).
    *
+   * When called, the callback receives a single argument, the status of the current request.
+   *
+   * @param[in] status the status of the request being completed.
    * @throws std::runtime_error if called by a send request.
    */
-  void callback();
+  void callback(ucs_status_t status);
 
   /**
    * @brief Return the status of the request.

cpp/include/ucxx/typedefs.h

@@ -5,6 +5,7 @@
 #pragma once
 
 #include <functional>
+#include <memory>
 #include <string>
 #include <unordered_map>
 
@@ -32,4 +33,7 @@ typedef enum {
 
 typedef std::unordered_map<std::string, std::string> ConfigMap;
 
+typedef std::function<void(ucs_status_t, std::shared_ptr<void>)> RequestCallbackUserFunction;
+typedef std::shared_ptr<void> RequestCallbackUserData;
+
 }  // namespace ucxx

cpp/include/ucxx/worker.h

@@ -544,13 +544,12 @@ class Worker : public Component {
    *
    * @returns Request to be subsequently checked for the completion and its state.
    */
-  std::shared_ptr<Request> tagRecv(
-    void* buffer,
-    size_t length,
-    ucp_tag_t tag,
-    const bool enableFuture                                     = false,
-    std::function<void(std::shared_ptr<void>)> callbackFunction = nullptr,
-    std::shared_ptr<void> callbackData                          = nullptr);
+  std::shared_ptr<Request> tagRecv(void* buffer,
+                                   size_t length,
+                                   ucp_tag_t tag,
+                                   const bool enableFuture                      = false,
+                                   RequestCallbackUserFunction callbackFunction = nullptr,
+                                   RequestCallbackUserData callbackData         = nullptr);
 
   /**
    * @brief Get the address of the UCX worker object.

cpp/src/endpoint.cpp

@@ -209,26 +209,24 @@ std::shared_ptr<Request> Endpoint::streamRecv(void* buffer,
     createRequestStream(endpoint, false, buffer, length, enablePythonFuture));
 }
 
-std::shared_ptr<Request> Endpoint::tagSend(
-  void* buffer,
-  size_t length,
-  ucp_tag_t tag,
-  const bool enablePythonFuture,
-  std::function<void(std::shared_ptr<void>)> callbackFunction,
-  std::shared_ptr<void> callbackData)
+std::shared_ptr<Request> Endpoint::tagSend(void* buffer,
+                                           size_t length,
+                                           ucp_tag_t tag,
+                                           const bool enablePythonFuture,
+                                           RequestCallbackUserFunction callbackFunction,
+                                           RequestCallbackUserData callbackData)
 {
   auto endpoint = std::dynamic_pointer_cast<Endpoint>(shared_from_this());
   return registerInflightRequest(createRequestTag(
     endpoint, true, buffer, length, tag, enablePythonFuture, callbackFunction, callbackData));
 }
 
-std::shared_ptr<Request> Endpoint::tagRecv(
-  void* buffer,
-  size_t length,
-  ucp_tag_t tag,
-  const bool enablePythonFuture,
-  std::function<void(std::shared_ptr<void>)> callbackFunction,
-  std::shared_ptr<void> callbackData)
+std::shared_ptr<Request> Endpoint::tagRecv(void* buffer,
+                                           size_t length,
+                                           ucp_tag_t tag,
+                                           const bool enablePythonFuture,
+                                           RequestCallbackUserFunction callbackFunction,
+                                           RequestCallbackUserData callbackData)
 {
   auto endpoint = std::dynamic_pointer_cast<Endpoint>(shared_from_this());
   return registerInflightRequest(createRequestTag(

cpp/src/request.cpp

@@ -98,17 +98,17 @@ bool Request::isCompleted() { return _status != UCS_INPROGRESS; }
 
 void Request::callback(void* request, ucs_status_t status)
 {
-  setStatus(status);
-
   ucxx_trace_req_f(_ownerString.c_str(),
                    request,
                    _operationName.c_str(),
                    "callback %p",
                    _callback.target<void (*)(void)>());
-  if (_callback) _callback(_callbackData);
+  if (_callback) _callback(status, _callbackData);
 
   ucp_request_free(request);
+
   ucxx_trace("Request completed: %p, handle: %p", this, request);
+  setStatus(status);
 }
 
 void Request::process()
@@ -138,13 +138,6 @@ void Request::process()
                    status,
                    ucs_status_string(status));
 
-  ucxx_trace_req_f(_ownerString.c_str(),
-                   _request,
-                   _operationName.c_str(),
-                   "callback %p",
-                   _callback.target<void (*)(void)>());
-  if (_callback) _callback(_callbackData);
-
   if (status != UCS_OK) {
     ucxx_error(
       "error on %s with status %d (%s)", _operationName.c_str(), status, ucs_status_string(status));
@@ -153,6 +146,13 @@ void Request::process()
       _ownerString.c_str(), _request, _operationName.c_str(), "completed immediately");
   }
 
+  ucxx_trace_req_f(_ownerString.c_str(),
+                   _request,
+                   _operationName.c_str(),
+                   "callback %p",
+                   _callback.target<void (*)(void)>());
+  if (_callback) _callback(status, _callbackData);
+
   setStatus(status);
 }
 

cpp/src/request_tag.cpp

@@ -13,15 +13,14 @@
 
 namespace ucxx {
 
-std::shared_ptr<RequestTag> createRequestTag(
-  std::shared_ptr<Component> endpointOrWorker,
-  bool send,
-  void* buffer,
-  size_t length,
-  ucp_tag_t tag,
-  const bool enablePythonFuture                               = false,
-  std::function<void(std::shared_ptr<void>)> callbackFunction = nullptr,
-  std::shared_ptr<void> callbackData                          = nullptr)
+std::shared_ptr<RequestTag> createRequestTag(std::shared_ptr<Component> endpointOrWorker,
+                                             bool send,
+                                             void* buffer,
+                                             size_t length,
+                                             ucp_tag_t tag,
+                                             const bool enablePythonFuture                = false,
+                                             RequestCallbackUserFunction callbackFunction = nullptr,
+                                             RequestCallbackUserData callbackData         = nullptr)
 {
   return std::shared_ptr<RequestTag>(new RequestTag(endpointOrWorker,
                                                     send,
@@ -39,8 +38,8 @@ RequestTag::RequestTag(std::shared_ptr<Component> endpointOrWorker,
                        size_t length,
                        ucp_tag_t tag,
                        const bool enablePythonFuture,
-                       std::function<void(std::shared_ptr<void>)> callbackFunction,
-                       std::shared_ptr<void> callbackData)
+                       RequestCallbackUserFunction callbackFunction,
+                       RequestCallbackUserData callbackData)
   : Request(endpointOrWorker,
             std::make_shared<DelayedSubmission>(send, buffer, length, tag),
             std::string(send ? "tagSend" : "tagRecv"),