throw exception instead of assertion OK

JuliaParallel · Feb 1, 2021 · 345d179 · 345d179
1 parent 3286156
commit 345d179
Showing 1 changed file with 24 additions and 35 deletions.
diff --git a/src/UCX.jl b/src/UCX.jl
@@ -30,11 +30,19 @@ UCS_PTR_STATUS(ptr::Ptr{Cvoid}) = API.ucs_status_t(reinterpret(UInt, ptr))
 UCS_PTR_IS_ERR(ptr::Ptr{Cvoid}) = uintptr_t(ptr) >= uintptr_t(API.UCS_ERR_LAST)
 UCS_PTR_IS_PTR(ptr::Ptr{Cvoid}) = (uintptr_t(ptr) - 1) < (uintptr_t(API.UCS_ERR_LAST) - 1)
 
-struct UCXError
-    ctx::String
+struct UCXException <: Exception
     status::API.ucs_status_t
 end
 
+macro check(ex)
+    quote
+        status = $(esc(ex))
+        if status != API.UCS_OK
+            throw(UCXException(status))
+        end
+    end
+end
+
 function version()
     major = Ref{Cuint}()
     minor = Ref{Cuint}()
@@ -48,8 +56,7 @@ mutable struct UCXConfig
 
     function UCXConfig(;kwargs...)
         r_handle = Ref{Ptr{API.ucp_config_t}}()
-        status = API.ucp_config_read(C_NULL, C_NULL, r_handle) # XXX: Prefix is broken
-        @assert status == API.UCS_OK
+        @check API.ucp_config_read(C_NULL, C_NULL, r_handle) # XXX: Prefix is broken
 
         config = new(r_handle[])
         finalizer(config) do config
@@ -65,8 +72,7 @@ mutable struct UCXConfig
 end
 
 function Base.setindex!(config::UCXConfig, value::String, key::Union{String, Symbol})
-    status = API.ucp_config_modify(config.handle, key, value)
-    @assert status == API.UCS_OK
+    @check API.ucp_config_modify(config.handle, key, value)
     return value
 end
 
@@ -119,9 +125,8 @@ mutable struct UCXContext
 
         r_handle = Ref{API.ucp_context_h}()
         # UCP.ucp_init is a header function so we call, UCP.ucp_init_version
-        status = API.ucp_init_version(API.UCP_API_MAJOR, API.UCP_API_MINOR,
-                                      params, config.handle, r_handle)
-        @assert status === API.UCS_OK
+        @check API.ucp_init_version(API.UCP_API_MAJOR, API.UCP_API_MINOR,
+                                    params, config.handle, r_handle)
 
         context = new(r_handle[], parse(Dict, config))
 
@@ -167,8 +172,7 @@ mutable struct UCXWorker
         set!(params, :thread_mode, thread_mode)
 
         r_handle = Ref{API.ucp_worker_h}()
-        status = API.ucp_worker_create(context.handle, params, r_handle)
-        @assert status === API.UCS_OK
+        @check API.ucp_worker_create(context.handle, params, r_handle)
 
         worker = new(r_handle[], context)
         finalizer(worker) do worker
@@ -182,15 +186,6 @@ function progress(worker::UCXWorker)
     API.ucp_worker_progress(worker.handle) !== 0
 end
 
-function arm(worker::UCXWorker)
-    status = API.ucp_worker_arm(worker.handle)
-    if status == API.UCS_ERR_BUSY
-        return false
-    end
-    @assert status == API.UCS_OK
-    return true
-end
-
 struct UCXConnectionRequest
     handle::API.ucp_conn_request_h
 end
@@ -210,7 +205,7 @@ mutable struct UCXEndpoint
                 end
                 API.ucp_request_free(status)
             else
-                @assert UCS_PTR_STATUS(status) == API.UCS_OK
+                @check UCS_PTR_STATUS(status)
             end
         end
         endpoint
@@ -240,8 +235,7 @@ function UCXEndpoint(worker::UCXWorker, ip::IPv4, port)
 
         # TODO: Error callback
 
-        status = API.ucp_ep_create(worker.handle, params, r_handle)
-        @assert status == API.UCS_OK
+        @check API.ucp_ep_create(worker.handle, params, r_handle)
     end
 
     UCXEndpoint(worker, r_handle[])
@@ -261,8 +255,7 @@ function UCXEndpoint(worker::UCXWorker, conn_request::UCXConnectionRequest)
     # TODO: Error callback
 
     r_handle = Ref{API.ucp_ep_h}()
-    status = API.ucp_ep_create(worker.handle, params, r_handle)
-    @assert status == API.UCS_OK
+    @check API.ucp_ep_create(worker.handle, params, r_handle)
 
     UCXEndpoint(worker, r_handle[])
 end
@@ -295,8 +288,7 @@ mutable struct UCXListener
             set!(params, :sockaddr, ucs_sockaddr)
             set!(params, :conn_handler, conn_handler)
 
-            status = API.ucp_listener_create(worker.handle, params, r_handle)
-            @assert status === API.UCS_OK
+            @check API.ucp_listener_create(worker.handle, params, r_handle)
         end  
 
         listener = new(r_handle[], worker, port)
@@ -307,8 +299,7 @@ mutable struct UCXListener
 end
 
 function reject(listener::UCXListener, conn_request::UCXConnectionRequest)
-    status = API.ucp_listener_reject(listener.handle, conn_request.handle)
-    @assert status === API.UCS_OK
+    @check API.ucp_listener_reject(listener.handle, conn_request.handle)
 end
 
 function ucp_dt_make_contig(elem_size)
@@ -330,20 +321,18 @@ end
 # Current implementation is blocking
 handle_request(ep::UCXEndpoint, ptr) = handle_request(ep.worker, ptr)
 function handle_request(worker::UCXWorker, ptr)
-    if ptr === C_NULL
-        return API.UCS_OK
-    elseif UCS_PTR_IS_ERR(ptr)
-        return UCS_PTR_STATUS(ptr)
-    else
+    if UCS_PTR_IS_PTR(ptr)
         status = API.ucp_request_check_status(ptr)
         while status === API.UCS_INPROGRESS
             progress(worker)
             yield()
             status = API.ucp_request_check_status(ptr)
         end
         API.ucp_request_free(ptr)
-        return status
+    else
+        status = UCS_PTR_STATUS(ptr)
     end
+    @check status
 end