diff --git a/src/uct/tcp/tcp.h b/src/uct/tcp/tcp.h index fa097e8e83a..6fc553f3e68 100644 --- a/src/uct/tcp/tcp.h +++ b/src/uct/tcp/tcp.h @@ -205,11 +205,13 @@ typedef struct uct_tcp_am_hdr { */ typedef enum uct_tcp_ep_am_id { /* AM ID reserved for TCP internal Connection Manager messages */ - UCT_TCP_EP_CM_AM_ID = UCT_AM_ID_MAX, + UCT_TCP_EP_CM_AM_ID = UCT_AM_ID_MAX, /* AM ID reserved for TCP internal PUT REQ message */ - UCT_TCP_EP_PUT_REQ_AM_ID = UCT_AM_ID_MAX + 1, + UCT_TCP_EP_PUT_REQ_AM_ID = UCT_AM_ID_MAX + 1, /* AM ID reserved for TCP internal PUT ACK message */ - UCT_TCP_EP_PUT_ACK_AM_ID = UCT_AM_ID_MAX + 2 + UCT_TCP_EP_PUT_ACK_AM_ID = UCT_AM_ID_MAX + 2, + /* AM ID reserved for TCP internal PUT ACK message */ + UCT_TCP_EP_KEEPALIVE_AM_ID = UCT_AM_ID_MAX + 3 } uct_tcp_ep_am_id_t; @@ -513,6 +515,9 @@ void uct_tcp_ep_pending_purge(uct_ep_h tl_ep, uct_pending_purge_callback_t cb, ucs_status_t uct_tcp_ep_flush(uct_ep_h tl_ep, unsigned flags, uct_completion_t *comp); +ucs_status_t +uct_tcp_ep_check(uct_ep_h tl_ep, unsigned flags, uct_completion_t *comp); + ucs_status_t uct_tcp_cm_send_event(uct_tcp_ep_t *ep, uct_tcp_cm_conn_event_t event, int log_error); diff --git a/src/uct/tcp/tcp_cm.c b/src/uct/tcp/tcp_cm.c index 172629120b4..c22adcb7b63 100644 --- a/src/uct/tcp/tcp_cm.c +++ b/src/uct/tcp/tcp_cm.c @@ -182,7 +182,7 @@ ucs_status_t uct_tcp_cm_send_event(uct_tcp_ep_t *ep, pkt_buf = ucs_alloca(pkt_length); pkt_hdr = (uct_tcp_am_hdr_t*)(UCS_PTR_BYTE_OFFSET(pkt_buf, magic_number_length)); - pkt_hdr->am_id = UCT_AM_ID_MAX; + pkt_hdr->am_id = UCT_TCP_EP_CM_AM_ID; pkt_hdr->length = cm_pkt_length; if (event == UCT_TCP_CM_CONN_REQ) { diff --git a/src/uct/tcp/tcp_ep.c b/src/uct/tcp/tcp_ep.c index 6d73a551baf..bf91f5c1dfa 100644 --- a/src/uct/tcp/tcp_ep.c +++ b/src/uct/tcp/tcp_ep.c @@ -1295,6 +1295,9 @@ static unsigned uct_tcp_ep_progress_am_rx(uct_tcp_ep_t *ep) ucs_assert(hdr->length == sizeof(uint32_t)); uct_tcp_ep_handle_put_ack(ep, (uct_tcp_ep_put_ack_hdr_t*)(hdr + 1)); handled++; + } else if (hdr->am_id == UCT_TCP_EP_KEEPALIVE_AM_ID) { + /* just ignore keepalive requests */ + handled++; } else { ucs_assert(hdr->am_id == UCT_TCP_EP_CM_AM_ID); handled += 1 + uct_tcp_cm_handle_conn_pkt(&ep, hdr + 1, hdr->length); @@ -1920,3 +1923,23 @@ ucs_status_t uct_tcp_ep_flush(uct_ep_h tl_ep, unsigned flags, return UCS_OK; } +ucs_status_t +uct_tcp_ep_check(uct_ep_h tl_ep, unsigned flags, uct_completion_t *comp) +{ + uct_tcp_ep_t *ep = ucs_derived_of(tl_ep, uct_tcp_ep_t); + uct_tcp_iface_t *iface = ucs_derived_of(tl_ep->iface, uct_tcp_iface_t); + uct_tcp_am_hdr_t *hdr = NULL; /* init to suppress build warning */ + ucs_status_t status; + + UCT_CHECK_PARAM(comp == NULL, "Unsupported completion on ep_check"); + UCT_CHECK_PARAM(flags == 0, "Unsupported flags: %x", flags); + + status = uct_tcp_ep_am_prepare(iface, ep, UCT_TCP_EP_KEEPALIVE_AM_ID, &hdr); + if (status != UCS_OK) { + return (status == UCS_ERR_NO_RESOURCE) ? UCS_OK : status; + } + + ucs_assert(hdr != NULL); + hdr->length = 0; + return uct_tcp_ep_am_send(ep, hdr); +} diff --git a/src/uct/tcp/tcp_iface.c b/src/uct/tcp/tcp_iface.c index cc80420dd6d..89bd1c6655f 100644 --- a/src/uct/tcp/tcp_iface.c +++ b/src/uct/tcp/tcp_iface.c @@ -160,6 +160,7 @@ static ucs_status_t uct_tcp_iface_query(uct_iface_h tl_iface, uct_iface_attr_t * UCT_IFACE_FLAG_AM_BCOPY | UCT_IFACE_FLAG_PENDING | UCT_IFACE_FLAG_CB_SYNC | + UCT_IFACE_FLAG_EP_CHECK | UCT_IFACE_FLAG_ERRHANDLE_PEER_FAILURE; attr->cap.event_flags = UCT_IFACE_FLAG_EVENT_SEND_COMP | UCT_IFACE_FLAG_EVENT_RECV | @@ -342,6 +343,7 @@ static uct_iface_ops_t uct_tcp_iface_ops = { .ep_pending_purge = uct_tcp_ep_pending_purge, .ep_flush = uct_tcp_ep_flush, .ep_fence = uct_base_ep_fence, + .ep_check = uct_tcp_ep_check, .ep_create = uct_tcp_ep_create, .ep_destroy = uct_tcp_ep_destroy, .ep_get_address = uct_tcp_ep_get_address,