static void psmx_set_epaddr_context(struct psmx_fid_domain *domain,
psm_epid_t epid, psm_epaddr_t epaddr)
{
struct psmx_epaddr_context *context;
context = (void *)psm_epaddr_getctxt(epaddr);
if (context) {
if (context->domain != domain || context->epid != epid) {
PSMX_WARN("%s: domain or epid doesn't match\n", __func__);
context = NULL;
}
}
if (context)
return;
context = malloc(sizeof *context);
if (!context) {
PSMX_WARN("%s: cannot allocate context\n", __func__);
return;
}
context->domain = domain;
context->epid = epid;
psm_epaddr_setctxt(epaddr, context);
}
int psmx_am_msg_handler(psm_am_token_t token, psm_epaddr_t epaddr,
psm_amarg_t *args, int nargs, void *src, uint32_t len)
{
psm_amarg_t rep_args[8];
struct psmx_am_request *req;
struct psmx_cq_event *event;
struct psmx_epaddr_context *epaddr_context;
struct psmx_fid_domain *domain;
int copy_len;
uint64_t offset;
int cmd, eom;
int err = 0;
int op_error = 0;
struct psmx_unexp *unexp;
epaddr_context = psm_epaddr_getctxt(epaddr);
if (!epaddr_context) {
PSMX_WARN("%s: NULL context for epaddr %p\n", __func__, epaddr);
return -FI_EIO;
}
domain = epaddr_context->domain;
cmd = args[0].u32w0 & PSMX_AM_OP_MASK;
eom = args[0].u32w0 & PSMX_AM_EOM;
switch (cmd) {
case PSMX_AM_REQ_SEND:
assert(len == args[0].u32w1);
offset = args[3].u64;
if (offset == 0) {
/* this is the first packet */
req = psmx_am_search_and_dequeue_recv(domain, (const void *)epaddr);
if (req) {
copy_len = MIN(len, req->recv.len);
memcpy(req->recv.buf, src, len);
req->recv.len_received += copy_len;
}
else {
unexp = malloc(sizeof(*unexp) + len);
if (!unexp) {
op_error = -FI_ENOSPC;
}
else {
memcpy(unexp->buf, src, len);
unexp->sender_addr = epaddr;
unexp->sender_context = args[1].u64;
unexp->len_received = len;
unexp->done = !!eom;
unexp->list_entry.next = NULL;
psmx_am_enqueue_unexp(domain, unexp);
if (!eom) {
/* stop here. will reply when recv is posted */
break;
}
}
}
if (!op_error && !eom) {
/* reply w/ recv req to be used for following packets */
rep_args[0].u32w0 = PSMX_AM_REP_SEND;
rep_args[0].u32w1 = 0;
rep_args[1].u64 = args[1].u64;
rep_args[2].u64 = (uint64_t)(uintptr_t)req;
err = psm_am_reply_short(token, PSMX_AM_MSG_HANDLER,
rep_args, 3, NULL, 0, 0,
NULL, NULL );
}
}
else {
req = (struct psmx_am_request *)(uintptr_t)args[2].u64;
if (req) {
copy_len = MIN(req->recv.len + offset, len);
memcpy(req->recv.buf + offset, src, copy_len);
req->recv.len_received += copy_len;
}
else {
PSMX_WARN("%s: NULL recv_req in follow-up packets.\n", __func__);
op_error = -FI_ENOMSG;
}
}
if (eom && req) {
if (req->ep->recv_cq && !req->no_event) {
event = psmx_cq_create_event(
req->ep->recv_cq,
req->recv.context,
req->recv.buf,
req->cq_flags,
req->recv.len_received,
0, /* data */
0, /* tag */
req->recv.len - req->recv.len_received,
0 /* err */);
if (event)
psmx_cq_enqueue_event(req->ep->recv_cq, event);
else
err = -FI_ENOMEM;
}
if (req->ep->recv_cntr)
psmx_cntr_inc(req->ep->recv_cntr);
free(req);
}
if (eom || op_error) {
rep_args[0].u32w0 = PSMX_AM_REP_SEND;
rep_args[0].u32w1 = op_error;
rep_args[1].u64 = args[1].u64;
rep_args[2].u64 = 0; /* done */
err = psm_am_reply_short(token, PSMX_AM_MSG_HANDLER,
rep_args, 3, NULL, 0, 0,
NULL, NULL );
}
break;
case PSMX_AM_REP_SEND:
req = (struct psmx_am_request *)(uintptr_t)args[1].u64;
op_error = (int)args[0].u32w1;
assert(req->op == PSMX_AM_REQ_SEND);
if (args[2].u64) { /* more to send */
req->send.peer_context = (void *)(uintptr_t)args[2].u64;
#if PSMX_AM_USE_SEND_QUEUE
/* psm_am_request_short() can't be called inside the handler.
* put the request into a queue and process it later.
*/
psmx_am_enqueue_send(req->ep->domain, req);
if (req->ep->domain->progress_thread)
pthread_cond_signal(&req->ep->domain->progress_cond);
#else
req->send.peer_ready = 1;
#endif
}
else { /* done */
if (req->ep->send_cq && !req->no_event) {
event = psmx_cq_create_event(
req->ep->send_cq,
req->send.context,
req->send.buf,
req->cq_flags,
req->send.len,
0, /* data */
0, /* tag */
0, /* olen */
op_error);
if (event)
psmx_cq_enqueue_event(req->ep->send_cq, event);
else
err = -FI_ENOMEM;
}
if (req->ep->send_cntr)
psmx_cntr_inc(req->ep->send_cntr);
if (req->state == PSMX_AM_STATE_QUEUED)
req->state = PSMX_AM_STATE_DONE;
else
free(req);
}
break;
default:
err = -FI_EINVAL;
}
return err;
}
