BINDINGS/GO: Bug fix - make gobindings thread safe.

ofirfarjun7 · ofirfarjun7 · commit ba5cc7e061c0 · 2025-05-04T15:38:08.000Z
diff --git a/bindings/go/src/ucx/am_data.go b/bindings/go/src/ucx/am_data.go
@@ -20,9 +20,6 @@ type UcpAmData struct {
 	flags   UcpAmRecvAttrs
 }
 
-// To connect callback id with worker, to use in AmData.Receive()
-var idToWorker = make(map[uint64]*UcpWorker)
-
 // Whether actual data is received or need to call UcpAmData.Receive()
 func (d *UcpAmData) IsDataValid() bool {
 	return (d.flags & UCP_AM_RECV_ATTR_FLAG_RNDV) == 0
diff --git a/bindings/go/src/ucx/callbacks.go b/bindings/go/src/ucx/callbacks.go
@@ -8,8 +8,8 @@ package ucx
 // #include "goucx.h"
 import "C"
 import (
-	"sync"
 	"unsafe"
+	"runtime/cgo"
 )
 
 type UcpCallback interface{}
@@ -23,46 +23,48 @@ type UcpAmDataRecvCallback = func(request *UcpRequest, status UcsStatus, length
 type UcpAmRecvCallback = func(header unsafe.Pointer, headerSize uint64,
 	data *UcpAmData, replyEp *UcpEp) UcsStatus
 
+type UcpAmRecvCallbackBundle struct {
+	cb     UcpAmRecvCallback
+	worker *UcpWorker
+}	
+
 // This callback routine is invoked on the server side to handle incoming
 // connections from remote clients.
 type UcpListenerConnectionHandler = func(connRequest *UcpConnectionRequest)
 
-// Map from the callback id that is passed to C to the actual go callback.
-var callback_map = make(map[uint64]UcpCallback)
+type PackedCallback cgo.Handle
+func packCallback(cb UcpCallback) PackedCallback {
+    if cb == nil {
+        return 0
+    }
+
+    return PackedCallback(cgo.NewHandle(cb))
+}
 
-// Unique index for each go callback, that passes to user_data.
-var callback_id uint64 = 1
+func (pc PackedCallback) unpackInternal(freeHandle bool) UcpCallback {
+    if pc == 0 {
+        return nil
+    }
 
-var mu sync.Mutex
+    h := cgo.Handle(pc)
+    if freeHandle {
+        defer h.Delete()
+    }
 
-// Associates go callback with a unique id
-func register(cb UcpCallback) uint64 {
-	mu.Lock()
-	defer mu.Unlock()
-	callback_id++
-	callback_map[callback_id] = cb
-	return callback_id
+    return h.Value().(UcpCallback)
 }
 
-// Atomically removes registered callback by it's id
-func deregister(id uint64) (UcpCallback, bool) {
-	mu.Lock()
-	defer mu.Unlock()
-	val, ret := callback_map[id]
-	delete(callback_map, id)
-	return val, ret
+func (pc PackedCallback) unpack() UcpCallback {
+    return pc.unpackInternal(false)
 }
 
-func getCallback(id uint64) (UcpCallback, bool) {
-	mu.Lock()
-	defer mu.Unlock()
-	val, ret := callback_map[id]
-	return val, ret
+func (pc PackedCallback) unpackAndFree() UcpCallback {
+    return pc.unpackInternal(true)
 }
 
 //export ucxgo_completeGoSendRequest
-func ucxgo_completeGoSendRequest(request unsafe.Pointer, status C.ucs_status_t, callbackId unsafe.Pointer) {
-	if callback, found := deregister(uint64(uintptr(callbackId))); found {
+func ucxgo_completeGoSendRequest(request unsafe.Pointer, status C.ucs_status_t, packedCb unsafe.Pointer) {
+	if callback := PackedCallback(packedCb).unpackAndFree(); callback != nil {
 		callback.(UcpSendCallback)(&UcpRequest{
 			request: request,
 			Status:  UcsStatus(status),
@@ -71,8 +73,8 @@ func ucxgo_completeGoSendRequest(request unsafe.Pointer, status C.ucs_status_t,
 }
 
 //export ucxgo_completeGoTagRecvRequest
-func ucxgo_completeGoTagRecvRequest(request unsafe.Pointer, status C.ucs_status_t, tag_info *C.ucp_tag_recv_info_t, callbackId unsafe.Pointer) {
-	if callback, found := deregister(uint64(uintptr(callbackId))); found {
+func ucxgo_completeGoTagRecvRequest(request unsafe.Pointer, status C.ucs_status_t, tag_info *C.ucp_tag_recv_info_t, packedCb unsafe.Pointer) {
+	if callback := PackedCallback(packedCb).unpackAndFree(); callback != nil {
 		callback.(UcpTagRecvCallback)(&UcpRequest{
 			request: request,
 			Status:  UcsStatus(status),
@@ -84,30 +86,30 @@ func ucxgo_completeGoTagRecvRequest(request unsafe.Pointer, status C.ucs_status_
 }
 
 //export ucxgo_amRecvCallback
-func ucxgo_amRecvCallback(calbackId unsafe.Pointer, header unsafe.Pointer, headerSize C.size_t,
+func ucxgo_amRecvCallback(packedCb unsafe.Pointer, header unsafe.Pointer, headerSize C.size_t,
 	data unsafe.Pointer, dataSize C.size_t, params *C.ucp_am_recv_param_t) C.ucs_status_t {
-	cbId := uint64(uintptr(calbackId))
-	if callback, found := getCallback(cbId); found {
+	if callback := PackedCallback(packedCb).unpack(); callback != nil {
+		bundle := callback.(*UcpAmRecvCallbackBundle)
 		var replyEp *UcpEp
 		if (params.recv_attr & C.UCP_AM_RECV_ATTR_FIELD_REPLY_EP) != 0 {
 			replyEp = &UcpEp{ep: params.reply_ep}
 		}
 		amData := &UcpAmData{
-			worker:  idToWorker[cbId],
+			worker:  bundle.worker,
 			flags:   UcpAmRecvAttrs(params.recv_attr),
 			dataPtr: data,
 			length:  uint64(dataSize),
 		}
-		return C.ucs_status_t(callback.(UcpAmRecvCallback)(header, uint64(headerSize), amData, replyEp))
+		return C.ucs_status_t(bundle.cb(header, uint64(headerSize), amData, replyEp))
 	}
 	return C.UCS_OK
 }
 
 //export ucxgo_completeAmRecvData
 func ucxgo_completeAmRecvData(request unsafe.Pointer, status C.ucs_status_t,
-	length C.size_t, callbackId unsafe.Pointer) {
+	length C.size_t, packedCb unsafe.Pointer) {
 
-	if callback, found := deregister(uint64(uintptr(callbackId))); found {
+	if callback := PackedCallback(packedCb).unpackAndFree(); callback != nil {
 		callback.(UcpAmDataRecvCallback)(&UcpRequest{
 			request: request,
 			Status:  UcsStatus(status),
diff --git a/bindings/go/src/ucx/context.go b/bindings/go/src/ucx/context.go
@@ -5,9 +5,11 @@
 
 package ucx
 
+// #include <stdlib.h>
 // #include <ucp/api/ucp.h>
 // #include <ucs/type/status.h>
 import "C"
+import "unsafe"
 
 // UCP application context (or just a context) is an opaque handle that holds a
 // UCP communication instance's global information. It represents a single UCP
@@ -27,10 +29,44 @@ type UcpContext struct {
 	context C.ucp_context_h
 }
 
+type UcpConfig struct {
+	config *C.ucp_config_t
+}
+
+func ReadConfig() (UcpConfig, error) {
+	var config *C.ucp_config_t
+	if status := C.ucp_config_read(nil, nil, &config); status != C.UCS_OK {
+		return UcpConfig{}, newUcxError(status)
+	}
+	return UcpConfig{config: config}, nil
+}
+
+func (c *UcpConfig) Modify(name string, value string) error {
+	cName := C.CString(name)
+	cValue := C.CString(value)
+	defer C.free(unsafe.Pointer(cName))
+	defer C.free(unsafe.Pointer(cValue))
+	if status := C.ucp_config_modify(c.config, cName, cValue); status != C.UCS_OK {
+		return newUcxError(status)
+	}
+	return nil
+}
+
+func (c *UcpConfig) Close() {
+	C.ucp_config_release(c.config)
+}
+
 func NewUcpContext(contextParams *UcpParams) (*UcpContext, error) {
 	var ucp_context C.ucp_context_h
 
-	if status := C.ucp_init(&contextParams.params, nil, &ucp_context); status != C.UCS_OK {
+	config, err := ReadConfig()
+	if err != nil {
+		return nil, err
+	}
+	defer config.Close()
+	config.Modify("GVA_ENABLE", "auto")
+
+	if status := C.ucp_init(&contextParams.params, config.config, &ucp_context); status != C.UCS_OK {
 		return nil, newUcxError(status)
 	}
 
@@ -92,6 +128,10 @@ func (c *UcpContext) Query(attrs ...UcpContextAttr) (*C.ucp_context_attr_t, erro
 func (c *UcpContext) NewWorker(workerParams *UcpWorkerParams) (*UcpWorker, error) {
 	var ucp_worker C.ucp_worker_h
 
+	if ((workerParams.params.field_mask & C.UCP_WORKER_PARAM_FIELD_THREAD_MODE) == 0) {
+		workerParams.SetThreadMode(UCS_THREAD_MODE_MULTI)
+	}
+
 	if status := C.ucp_worker_create(c.context, &workerParams.params, &ucp_worker); status != C.UCS_OK {
 		return nil, newUcxError(status)
 	}
diff --git a/bindings/go/src/ucx/endpoint.go b/bindings/go/src/ucx/endpoint.go
@@ -18,7 +18,7 @@ type UcpEp struct {
 
 var errorHandles = make(map[C.ucp_ep_h]UcpEpErrHandler)
 
-func setSendParams(goRequestParams *UcpRequestParams, cRequestParams *C.ucp_request_param_t) uint64 {
+func setSendParams(goRequestParams *UcpRequestParams, cRequestParams *C.ucp_request_param_t) unsafe.Pointer {
 	return packParams(goRequestParams, cRequestParams, unsafe.Pointer(C.ucxgo_completeGoSendRequest))
 }
 
@@ -28,22 +28,20 @@ func setSendParams(goRequestParams *UcpRequestParams, cRequestParams *C.ucp_requ
 func (e *UcpEp) FlushNonBlocking(params *UcpRequestParams) (*UcpRequest, error) {
 	var requestParams C.ucp_request_param_t
 
-	cbId := setSendParams(params, &requestParams)
-
+	callback := setSendParams(params, &requestParams)
 	request := C.ucp_ep_flush_nbx(e.ep, &requestParams)
-	return NewRequest(request, cbId, nil)
+	return newRequest(request, callback, nil)
 }
 
 func (e *UcpEp) CloseNonBlocking(mode C.uint, params *UcpRequestParams) (*UcpRequest, error) {
 	var requestParams C.ucp_request_param_t
 	requestParams.op_attr_mask = C.UCP_OP_ATTR_FIELD_FLAGS
 	requestParams.flags = mode
 
-	cbId := setSendParams(params, &requestParams)
-
+	callback := setSendParams(params, &requestParams)
 	request := C.ucp_ep_close_nbx(e.ep, &requestParams)
 	delete(errorHandles, e.ep)
-	return NewRequest(request, cbId, nil)
+	return newRequest(request, callback, nil)
 }
 
 // Non-blocking endpoint closure. Releases the endpoint without any
@@ -68,10 +66,9 @@ func (e *UcpEp) SendTagNonBlocking(tag uint64, address unsafe.Pointer, size uint
 	params *UcpRequestParams) (*UcpRequest, error) {
 	var requestParams C.ucp_request_param_t
 
-	cbId := setSendParams(params, &requestParams)
-
+	callback := setSendParams(params, &requestParams)
 	request := C.ucp_tag_send_nbx(e.ep, address, C.size_t(size), C.ucp_tag_t(tag), &requestParams)
-	return NewRequest(request, cbId, nil)
+	return newRequest(request, callback, nil)
 }
 
 // This routine sends an Active Message to an ep.
@@ -81,11 +78,10 @@ func (e *UcpEp) SendAmNonBlocking(id uint, header unsafe.Pointer, headerSize uin
 	data unsafe.Pointer, dataSize uint64, flags UcpAmSendFlags, params *UcpRequestParams) (*UcpRequest, error) {
 	var requestParams C.ucp_request_param_t
 
-	cbId := setSendParams(params, &requestParams)
-
+	callback := setSendParams(params, &requestParams)
 	requestParams.op_attr_mask |= C.UCP_OP_ATTR_FIELD_FLAGS
 	requestParams.flags = C.uint(flags)
 
 	request := C.ucp_am_send_nbx(e.ep, C.uint(id), header, C.size_t(headerSize), data, C.size_t(dataSize), &requestParams)
-	return NewRequest(request, cbId, nil)
+	return newRequest(request, callback, nil)
 }
diff --git a/bindings/go/src/ucx/listener.go b/bindings/go/src/ucx/listener.go
@@ -7,24 +7,23 @@ package ucx
 
 // #include <ucp/api/ucp.h>
 import "C"
-import "net"
+import (
+	"net"
+)
 
 type UcpListener struct {
-	listener      C.ucp_listener_h
-	connHandlerId uint64
+	listener C.ucp_listener_h
+	callback UcpListenerConnectionHandler
+	packedCb PackedCallback
 }
 
-// Needed to call connHandler.Reject() rather than listener.Reject(connHandler)
-var connHandles2Listener = make(map[uint64]C.ucp_listener_h)
-
 type UcpListenerAttributes struct {
 	Address *net.TCPAddr
 }
 
 func (l *UcpListener) Close() {
 	C.ucp_listener_destroy(l.listener)
-	deregister(l.connHandlerId)
-	delete(connHandles2Listener, l.connHandlerId)
+	l.packedCb.unpackAndFree()
 }
 
 func (l *UcpListener) Query(attrs ...UcpListenerAttribute) (*UcpListenerAttributes, error) {
diff --git a/bindings/go/src/ucx/listener_params.go b/bindings/go/src/ucx/listener_params.go
@@ -16,18 +16,17 @@ import (
 
 // Tuning parameters for the UCP listener.
 type UcpListenerParams struct {
-	params        C.ucp_listener_params_t
-	connHandlerId uint64
+	params   C.ucp_listener_params_t
+	callback UcpListenerConnectionHandler
 }
 
 //export ucxgo_completeConnHandler
-func ucxgo_completeConnHandler(connRequest C.ucp_conn_request_h, cbId unsafe.Pointer) {
-	id := uint64(uintptr((cbId)))
-	if callback, found := getCallback(id); found {
-		listener := connHandles2Listener[id]
-		callback.(UcpListenerConnectionHandler)(&UcpConnectionRequest{
+func ucxgo_completeConnHandler(connRequest C.ucp_conn_request_h, packedCb unsafe.Pointer) {
+	if callback := PackedCallback(packedCb).unpack(); callback != nil {
+		listener := callback.(*UcpListener)
+		listener.callback(&UcpConnectionRequest{
 			connRequest: connRequest,
-			listener:    listener,
+			listener:    listener.listener,
 		})
 	}
 }
@@ -49,10 +48,8 @@ func (p *UcpListenerParams) SetSocketAddress(a *net.TCPAddr) (*UcpListenerParams
 // Handler of an incoming connection request in a client-server connection flow.
 func (p *UcpListenerParams) SetConnectionHandler(connHandler UcpListenerConnectionHandler) *UcpListenerParams {
 	var ucpConnHndl C.ucp_listener_conn_handler_t
-	cbId := register(connHandler)
 
-	p.connHandlerId = cbId
-	ucpConnHndl.arg = unsafe.Pointer(uintptr(cbId))
+	p.callback = connHandler
 	ucpConnHndl.cb = (C.ucp_listener_conn_callback_t)(C.ucxgo_completeConnHandler)
 	p.params.field_mask |= C.UCP_LISTENER_PARAM_FIELD_CONN_HANDLER
 	p.params.conn_handler = ucpConnHndl
diff --git a/bindings/go/src/ucx/request.go b/bindings/go/src/ucx/request.go
diff --git a/bindings/go/src/ucx/ucp_contsants.go b/bindings/go/src/ucx/ucp_contsants.go
diff --git a/bindings/go/src/ucx/worker.go b/bindings/go/src/ucx/worker.go

Original file line number	Diff line number	Diff line change
`@@ -20,9 +20,6 @@ type UcpAmData struct {`
`20`	`20`	`flags UcpAmRecvAttrs`
`21`	`21`	`}`
`22`	`22`
`23`		`-// To connect callback id with worker, to use in AmData.Receive()`
`24`		`-var idToWorker = make(map[uint64]*UcpWorker)`
`25`		`-`
`26`	`23`	`// Whether actual data is received or need to call UcpAmData.Receive()`
`27`	`24`	`func (d *UcpAmData) IsDataValid() bool {`
`28`	`25`	`return (d.flags & UCP_AM_RECV_ATTR_FLAG_RNDV) == 0`