Function ReceiveData

GPU-Virtual-Service/gpu-remoting/include/ucpUtil.h:283–347

}

static ucs_status_t ReceiveData(ucp_dt_iov_t* iov, size_t* iovCnt, void* dataBuffer, uint64_t amID, ucp_worker_h* dataWorker, ucs_memory_type_t memType, bool* isClosed) {
    // iovCnt is set by ActiveMessageRecvCallback
    const char* myName = "ReceiveData";
    ucs_status_t status = UCS_OK;
    ActiveMessageDesc_t am_request_ctx = { .complete = 0, .is_rndv = 0, .mem_type = memType, .desc = NULL, 
                                           .iov = iov, .iov_num = 0, .recv_buf = dataBuffer, .recv_length = 0};
    Request_t recv_request_ctx = {.type = 1, .complete = 0};

    ucp_am_handler_param_t param1;
    param1.field_mask   = UCP_AM_HANDLER_PARAM_FIELD_ID |
                          UCP_AM_HANDLER_PARAM_FIELD_CB |
                          UCP_AM_HANDLER_PARAM_FIELD_ARG;
    param1.id           = amID; // todo: TEST_AM_ID
    param1.cb           = ActiveMessageRecvCallback;
    param1.arg          = &am_request_ctx;
    status = ucp_worker_set_am_recv_handler(*dataWorker, &param1);
    if (status != UCS_OK) {
        tool::Logging(LOG_ERROR, myName, "failed to set am handler.\n");
        return status;
    }

    tool::Logging(LOG_COMM, myName, "waiting for the client to send a message.\n");
    while (!am_request_ctx.complete) { // waiting ActiveMessageRecvCallback() to be invoked
        if (isClosed != NULL && *isClosed) {
            tool::Logging(LOG_ERROR, myName, "the connection has been closed.\n");
            return UCS_ERR_CONNECTION_RESET;
        }
        ucp_worker_progress(*dataWorker);
    }

    *iovCnt = am_request_ctx.iov_num;
    size_t msg_length = am_request_ctx.recv_length;
    if (!am_request_ctx.is_rndv) {
        tool::Logging(LOG_COMM, myName, "Eager request has arrived\n");
    }
    else {
        tool::Logging(LOG_COMM, myName, "Rendezvous request has arrived\n");

        ucp_request_param_t param2;
        param2.op_attr_mask     = UCP_OP_ATTR_FIELD_CALLBACK |
                                  UCP_OP_ATTR_FIELD_DATATYPE |
                                  UCP_OP_ATTR_FIELD_USER_DATA|
                                  UCP_OP_ATTR_FIELD_MEMORY_TYPE;
        param2.op_attr_mask    |= UCP_OP_ATTR_FLAG_NO_IMM_CMPL;
        // param2.datatype     = (*iovCnt == 1) ? ucp_dt_make_contig(1) : UCP_DATATYPE_IOV;
        param2.datatype         = ucp_dt_make_contig(1);
        param2.user_data        = &recv_request_ctx;
        param2.cb.recv_am       = (ucp_am_recv_data_nbx_callback_t)RecvCallBack;
        param2.memory_type      = memType;
        Request_t* rndv_request = (Request_t*)ucp_am_recv_data_nbx(*dataWorker,
                                              am_request_ctx.desc,
                                              am_request_ctx.recv_buf, am_request_ctx.recv_length, //recv data
                                             &param2);

        status = Wait(rndv_request, &recv_request_ctx, dataWorker); 
        if (status != UCS_OK) {
            tool::Logging(LOG_ERROR, myName, "ucp_am_recv_data_nbx failed: %s\n", ucs_status_string(status));
        }