int wait_receive_data() {
    /* Wait for receive completion */
    err = ibv_get_cq_event(comp_chan, &evt_cq, &cq_context);
    if (err) return 1;

    ibv_ack_cq_events(evt_cq, 1);

    err = ibv_req_notify_cq(cq, 0);
    if (err) return 1;

    n = ibv_poll_cq(cq, 1, &wc);
    if (n <= 0) return 1;
    if (wc.status != IBV_WC_SUCCESS) return 1;
    
    return 0;
}

void * poll_cq(void *ctx)
{
  struct ibv_cq *cq;
  struct ibv_wc wc;

  while (1) {
    TEST_NZ(ibv_get_cq_event(s_ctx->comp_channel, &cq, &ctx));
    ibv_ack_cq_events(cq, 1);
    TEST_NZ(ibv_req_notify_cq(cq, 0));

    while (ibv_poll_cq(cq, 1, &wc))
      on_completion(&wc);
  }

  return NULL;
}

static int get_thread_wc(struct thread_context_t *t_ctx, struct ibv_wc *wc, int is_send)
{
	struct ibv_cq           *cq;
	struct ibv_comp_channel *comp_channel;
	struct rdma_resource_t *rdma_resource;
	struct user_param_t *user_param;
	void *ectx;
	int rc = 0;

	rdma_resource = t_ctx->rdma_resource;
	user_param    = &(rdma_resource->user_param);

	if (is_send) {
		cq = t_ctx->send_cq;
		comp_channel = t_ctx->send_comp_channel;
	} else {
		cq = t_ctx->recv_cq;
		comp_channel = t_ctx->recv_comp_channel;
	}

	if (user_param->use_event) {
		rc = ibv_get_cq_event(comp_channel, &cq, &ectx);
		if (rc != 0) {
			ERROR("Failed to do ibv_get_cq_event.\n");
			return 1;
		}

		ibv_ack_cq_events(cq, 1);

		rc = ibv_req_notify_cq(cq, 0);
		if (rc != 0) {
			ERROR("Failed to do ibv_get_cq_event");
			return 1;
		}
	}

	do {
		rc = ibv_poll_cq(cq, 1, wc);
		if (rc < 0) {
			ERROR("Failed to poll CQ.\n");
			return 1;
		}
	} while (!user_param->use_event && (rc == 0)); /// need timeout

	return 0;
}

static int rdma_poll_cq(RdmaDeviceResources *rdma_dev_res, struct ibv_cq *ibcq)
{
    int i, ne, total_ne = 0;
    BackendCtx *bctx;
    struct ibv_wc wc[2];
    RdmaProtectedGSList *cqe_ctx_list;

    qemu_mutex_lock(&rdma_dev_res->lock);
    do {
        ne = ibv_poll_cq(ibcq, ARRAY_SIZE(wc), wc);

        trace_rdma_poll_cq(ne, ibcq);

        for (i = 0; i < ne; i++) {
            bctx = rdma_rm_get_cqe_ctx(rdma_dev_res, wc[i].wr_id);
            if (unlikely(!bctx)) {
                rdma_error_report("No matching ctx for req %"PRId64,
                                  wc[i].wr_id);
                continue;
            }

            comp_handler(bctx->up_ctx, &wc[i]);

            if (bctx->backend_qp) {
                cqe_ctx_list = &bctx->backend_qp->cqe_ctx_list;
            } else {
                cqe_ctx_list = &bctx->backend_srq->cqe_ctx_list;
            }

            rdma_protected_gslist_remove_int32(cqe_ctx_list, wc[i].wr_id);
            rdma_rm_dealloc_cqe_ctx(rdma_dev_res, wc[i].wr_id);
            g_free(bctx);
        }
        total_ne += ne;
    } while (ne > 0);
    atomic_sub(&rdma_dev_res->stats.missing_cqe, total_ne);
    qemu_mutex_unlock(&rdma_dev_res->lock);

    if (ne < 0) {
        rdma_error_report("ibv_poll_cq fail, rc=%d, errno=%d", ne, errno);
    }

    rdma_dev_res->stats.completions += total_ne;

    return total_ne;
}

/**
 * Polling for events on a inner thread allows processing of management messages
 * like buffer connection immediately, even if the user is not polling.
 * Otherwise buffer constructors would block indefinitely.
 *
 * Deep learning workloads are about sending small numbers of large messages,
 * in which case this model works great. If the library was to be used to
 * exchange large numbers of short messages, it would be useful to split
 * management and data messages over two different queue pairs. User threads
 * could then wait or poll on the data queue pair directly.
 */
void RDMAAdapter::InternalThreadEntry() {
  while (!must_stop()) {
    ibv_cq* cq;
    void* cq_context;
    CHECK(!ibv_get_cq_event(channel_, &cq, &cq_context));
    CHECK(cq == cq_);
    ibv_ack_cq_events(cq, 1);
    CHECK(!ibv_req_notify_cq(cq_, 0));

    int ne = ibv_poll_cq(cq_, MAX_CONCURRENT_WRITES * 2,
      static_cast<ibv_wc*>(wc_));
    CHECK_GE(ne, 0);

    for (int i = 0; i < ne; ++i) {
      CHECK(wc_[i].status == IBV_WC_SUCCESS) << "Failed status \n"
                                             << ibv_wc_status_str(wc_[i].status)
                                             << " " << wc_[i].status << " "
                                             << static_cast<int>(wc_[i].wr_id)
                                             << " "<< wc_[i].vendor_err;

      if (wc_[i].opcode == IBV_WC_RECV_RDMA_WITH_IMM) {
        // Data message, add it to user received queue
        RDMAChannel* channel = reinterpret_cast<RDMAChannel*>(wc_[i].wr_id);
        channel->recv();
        int id = wc_[i].imm_data;
        if (id >= CTRL_ID_OFFSET) {
        // ctrl signal
          ctrl_received_.push(channel->buffers_[id - CTRL_ID_OFFSET]);
        } else {
        // data
          received_.push(channel->buffers_[id]);
        }
      } else {
        if (wc_[i].opcode & IBV_WC_RECV) {
          // Buffer connection message
          RDMAChannel* channel = reinterpret_cast<RDMAChannel*>(wc_[i].wr_id);
          int id = wc_[i].imm_data;
          channel->memory_regions_queue_.push(channel->memory_regions_[id]);
          CHECK(id == channel->memory_regions_received_++);
          CHECK(!ibv_dereg_mr(channel->region_regions_[id]));
        }
      }
    }
  }
}

static int send_qp_num_for_ah(struct pingpong_context *ctx,
		struct perftest_parameters *user_param)
{
	struct ibv_send_wr wr;
	struct ibv_send_wr *bad_wr;
	struct ibv_sge list;
	struct ibv_wc wc;
	int ne;

	memcpy(ctx->buf,&ctx->qp[0]->qp_num,sizeof(uint32_t));

	list.addr   = (uintptr_t)ctx->buf;
	list.length = sizeof(uint32_t);
	list.lkey   = ctx->mr->lkey;

	wr.wr_id      = 0;
	wr.sg_list    = &list;
	wr.num_sge    = 1;
	wr.opcode     = IBV_WR_SEND_WITH_IMM;
	wr.send_flags = IBV_SEND_SIGNALED;
	wr.next       = NULL;
	wr.imm_data   = htonl(ctx->qp[0]->qp_num);

	wr.wr.ud.ah = ctx->ah[0];
	wr.wr.ud.remote_qpn  = user_param->rem_ud_qpn;
	wr.wr.ud.remote_qkey = user_param->rem_ud_qkey;


	if (ibv_post_send(ctx->qp[0],&wr,&bad_wr)) {
		fprintf(stderr, "Function ibv_post_send failed\n");
		return 1;
	}

	do {
		ne = ibv_poll_cq(ctx->send_cq, 1,&wc);
	} while (ne == 0);

	if (wc.status || wc.opcode != IBV_WC_SEND || wc.wr_id != 0) {
		fprintf(stderr, " Couldn't post send my QP number %d\n",(int)wc.status);
		return 1;
	}

	return 0;

}

inline void cfio_rdma_client_wait(void *ctx)
{
    struct ibv_cq *cq;
    struct ibv_wc wc;

    while (request_stack_size) {
        // rdma_debug("get cq event ...");
        TEST_NZ(ibv_get_cq_event(s_ctx->comp_channel, &cq, &ctx));
        // rdma_debug("ibv_ack_cq_events...");
        ibv_ack_cq_events(cq, 1);
        TEST_NZ(ibv_req_notify_cq(cq, 0));

        while (ibv_poll_cq(cq, 1, &wc)) {
            // rdma_debug("handle cq ...");
            on_completion(&wc);
        }
    }
}

void event_handler(struct ibv_cq *cq)
{
  int ret;
 
  while(1) {
     /* int ibv_poll_cq(a,b,c):
      *	    a: command queue to poll
      *	    b: max number of completions to return
      *	    c: array of at least (b) entries of ibv_wc where these
      *		completion events will be returned.
      */
    ret = ibv_poll_cq(cq, 1, &wc);

     if(ret == 0) {
        LOGPRINTF(("Empty completion queue, requesting next notification"));
        ibv_req_notify_cq(r_cq_hndl, 0);  /* ... explained in prev line.. */
        return;
     } else if(ret < 0) {
        fprintf(stderr, "Error in event_handler (polling cq)\n");
        exit(-1);
     } else if(wc.status != IBV_WC_SUCCESS) {
        fprintf(stderr, "Error in event_handler, on returned work completion "
		"status: %d\n", wc.status);
        exit(-1);
     }
     
     LOGPRINTF(("Retrieved work completion"));

     /* For ping-pong mode at least, this check shouldn't be needed for
      * normal operation, but it will help catch any bugs with multiple
      * sends coming through when we're only expecting one.
      */
     if(receive_complete == 1) {

        while(receive_complete != 0) sched_yield();

     }

     receive_complete = 1;

  }
  
}

static int poll_cqs(enum CQ_INDEX index)
{
	struct ibv_wc wc[8];
	int done, i, ret;

	for (i = 0; i < connections; i++) {
		if (!test.nodes[i].connected)
			continue;

		for (done = 0; done < message_count; done += ret) {
			ret = ibv_poll_cq(test.nodes[i].cq[index], 8, wc);
			if (ret < 0) {
				printf("cmatose: failed polling CQ: %d\n", ret);
				return ret;
			}
		}
	}
	return 0;
}

void *
poll_cq(void *ctx)
{
    struct ibv_cq *cq;
    struct ibv_wc wc;
    IbvConnection *conn = (IbvConnection *)ctx;

    while (1) {
        TEST_NZ(ibv_get_cq_event(conn->comp_channel, &cq, &ctx));
        ibv_ack_cq_events(cq, 1);
        TEST_NZ(ibv_req_notify_cq(cq, 0));

        while (ibv_poll_cq(cq, 1, &wc)) {
            (OnCompletionHandler)(&wc);
        }
    }

    return NULL;
}

/*-----------------------------------------------------------------------------------*/
static err_t
ibvif_thread(struct netif *netif)
{
  struct ibvif *ibvif;
  int ne, i;
  struct ibv_wc wc[PBUF_READ_DEPTH];

  ibvif = (struct ibvif *)netif->state;

  ne = ibv_poll_cq(ibvif->send_cq, PBUF_READ_DEPTH, wc);

  for (i=0; i<ne; i++) {
    if (wc[i].status != IBV_WC_SUCCESS) {
      perror("tapif: write 2");
    }
  }

  /* Wait for a packet to arrive. */
  low_level_input(netif);
}

static UCS_F_ALWAYS_INLINE ucs_status_t
uct_ud_verbs_iface_poll_rx(uct_ud_verbs_iface_t *iface)
{
    uct_ib_iface_recv_desc_t *desc;
    struct ibv_wc wc[UCT_IB_MAX_WC];
    int i, ret;
    char *packet;
    ucs_status_t status;


    ret = ibv_poll_cq(iface->super.super.recv_cq, UCT_IB_MAX_WC, wc);
    if (ret == 0) {
        status = UCS_ERR_NO_PROGRESS;
        goto out;
    } 
    if (ucs_unlikely(ret < 0)) {
        ucs_fatal("Failed to poll receive CQ");
    }

    for (i = 0; i < ret; ++i) {
        if (ucs_unlikely(wc[i].status != IBV_WC_SUCCESS)) {
            ucs_fatal("Receive completion with error: %s", 
                      ibv_wc_status_str(wc[i].status));
        }

        desc = (void*)wc[i].wr_id;
        ucs_trace_data("pkt rcvd: buf=%p len=%d", desc, wc[i].byte_len);
        packet = uct_ib_iface_recv_desc_hdr(&iface->super.super, desc);
        VALGRIND_MAKE_MEM_DEFINED(packet, wc[i].byte_len);

        uct_ud_ep_process_rx(&iface->super, 
                             (uct_ud_neth_t *)(packet + UCT_IB_GRH_LEN),
                             wc[i].byte_len - UCT_IB_GRH_LEN,
                             (uct_ud_recv_skb_t *)desc); 
    }
    iface->super.rx.available += ret;
    status = UCS_OK;
out:
    uct_ud_verbs_iface_post_recv(iface);
    return status;
}

static int __xfer_rdma_poll_cq(struct xfer_context *ctx, struct ibv_wc *ret_wc, int sleep)
{
	//void *ctx_ptr;
	int ne;

	/*
	if (ibv_get_cq_event(ctx->ch, &ctx->cq, &ctx_ptr)) {
		fprintf(stderr, "Failed to get cq_event\n");
		return -1;
	}
	
	ibv_ack_cq_events(ctx->cq, 1);

	if (ibv_req_notify_cq(ctx->cq, 0)) {
		fprintf(stderr, "Couldn't request CQ notification\n");
		return -1;
	}
	*/
	
	do {
		ne = ibv_poll_cq(ctx->cq, 1, ret_wc);
		if (ne < 0) {
			fprintf(stderr, "Failed to poll completions from the CQ\n");
			return -1;
		}
		
		if (!ne && sleep)
			usleep(100);
		
	} while (ne == 0);
	
	//printf("got events: %d, opcode: %d\n", ne, ret_wc->opcode);

	if (ret_wc->status != IBV_WC_SUCCESS) {
		fprintf(stderr, "Completion with status 0x%x was found\n", ret_wc->status);
		return -1;
	}
	
	return 0;
}

void mvdev_flush_qp(mv_qp_pool_entry *rqp, int num_to_flush) 
{
    struct ibv_qp_attr qp_attr;
    struct ibv_wc wc;
    int ne;

    memset(&qp_attr, 0, sizeof(qp_attr));
    qp_attr.qp_state = IBV_QPS_ERR;

    /* need to transition to the error state so we can flush
     * all the posted buffers
     */
    if(ibv_modify_qp(rqp->ud_qp, &qp_attr, IBV_QP_STATE)) {
        error_abort_all(IBV_RETURN_ERR, "Error changing to the err state\n");
    }

    /* pull failed completions */
    {
        int total_pulled = 0;
        do {
            ne = ibv_poll_cq(rqp->ud_cq, 1, &wc);
            total_pulled += ne;
        } while(total_pulled < num_to_flush);
    }

    {
        struct ibv_qp_attr attr;
        memset(&attr, 0, sizeof(struct ibv_qp_attr));

        attr.qp_state = IBV_QPS_RESET;

        if (ibv_modify_qp(rqp->ud_qp, &attr, IBV_QP_STATE)) {
            error_abort_all(IBV_RETURN_ERR,
                    "Failed to modify QP to RESET");
        }       
    }

    /* now we need to re-transition it back to the RTS phase */
    MV_Transition_UD_QP(&mvdev.rndv_si, rqp->ud_qp);
}

static UCS_F_ALWAYS_INLINE void 
uct_rc_verbs_iface_poll_tx(uct_rc_verbs_iface_t *iface)
{
    struct ibv_wc wc[UCT_IB_MAX_WC];
    uct_rc_verbs_ep_t *ep;
    uct_rc_iface_send_op_t *op;
    unsigned count;
    uint16_t sn;
    int i, ret;

    ret = ibv_poll_cq(iface->super.super.send_cq, UCT_IB_MAX_WC, wc);
    if (ucs_unlikely(ret <= 0)) {
        if (ucs_unlikely(ret < 0)) {
            ucs_fatal("Failed to poll send CQ");
        }
        return;
    }

    for (i = 0; i < ret; ++i) {
        if (ucs_unlikely(wc[i].status != IBV_WC_SUCCESS)) {
            ucs_fatal("Send completion with error: %s", ibv_wc_status_str(wc[i].status));
        }

        UCS_STATS_UPDATE_COUNTER(iface->super.stats, UCT_RC_IFACE_STAT_TX_COMPLETION, 1);

        ep = ucs_derived_of(uct_rc_iface_lookup_ep(&iface->super, wc[i].qp_num), uct_rc_verbs_ep_t);
        ucs_assert(ep != NULL);

        count = wc[i].wr_id + 1; /* Number of sends with WC completes in batch */
        ep->super.available         += count;
        ep->tx.completion_count     += count;
        ++iface->super.tx.cq_available;

        sn = ep->tx.completion_count;
        ucs_queue_for_each_extract(op, &ep->super.outstanding, queue,
                                   UCS_CIRCULAR_COMPARE16(op->sn, <=, sn)) {
            op->handler(op);
        }
    }
}

void RDMABuffer::Write(bool data) {
  struct ibv_sge list;
  list.addr = (uint64_t) addr_;
  list.length = size_;
  list.lkey = self_->lkey;

  struct ibv_send_wr wr;
  caffe_memset(sizeof(wr), 0, &wr);
  wr.wr_id = (uint64_t) this;
  wr.sg_list = &list;
  wr.num_sge = 1;
  wr.opcode = IBV_WR_RDMA_WRITE_WITH_IMM;
  wr.send_flags = IBV_SEND_SIGNALED;
  wr.imm_data = id_;
  if (!data) {
  // ctrl signal
    wr.imm_data += CTRL_ID_OFFSET;
  }

  wr.wr.rdma.remote_addr = (uint64_t) peer_->addr;
  wr.wr.rdma.rkey = peer_->rkey;

  struct ibv_send_wr *bad_wr;

  // lock the channel since there may be multiple threads calling write()
  boost::mutex::scoped_lock lock(channel_->mutex_);
  CHECK(!ibv_post_send(channel_->qp_, &wr, &bad_wr)) << "Failed to post send";

  // TODO poll only every N writes to improve performance
  for (;;) {
    ibv_wc wc;
    int ne = ibv_poll_cq(channel_->write_cq_, 1, &wc);
    CHECK_GE(ne, 0);
    if (ne) {
      CHECK(wc.wr_id == (uint64_t)this) << "Oops. Polled a Work Completion belongs to a different buffer";
      break;
    }
  }
}

static inline void uct_ud_verbs_iface_poll_tx(uct_ud_verbs_iface_t *iface)
{
    struct ibv_wc wc;
    int ret;

    ret = ibv_poll_cq(iface->super.super.send_cq, 1, &wc);
    if (ucs_unlikely(ret < 0)) {
        ucs_fatal("Failed to poll send CQ");
        return;
    }

    if (ret == 0) {
        return;
    }

    if (ucs_unlikely(wc.status != IBV_WC_SUCCESS)) {
        ucs_fatal("Send completion (wr_id=0x%0X with error: %s ", (unsigned)wc.wr_id, ibv_wc_status_str(wc.status));
        return;
    }

    iface->super.tx.available += UCT_UD_TX_MODERATION + 1;
}

static int rdma_read_keys(struct pingpong_dest *rem_dest,
		struct perftest_comm *comm)
{
	#ifdef HAVE_ENDIAN
	struct pingpong_dest a_rem_dest;
	#endif
	struct ibv_wc wc;
	int ne;

	do {
		ne = ibv_poll_cq(comm->rdma_ctx->recv_cq,1,&wc);
	} while (ne == 0);

	if (wc.status || !(wc.opcode & IBV_WC_RECV) || wc.wr_id != SYNC_SPEC_ID) {
		fprintf(stderr, "Bad wc status -- %d -- %d \n",(int)wc.status,(int)wc.wr_id);
		return 1;
	}

	#ifdef HAVE_ENDIAN
	memcpy(&a_rem_dest,comm->rdma_ctx->buf,sizeof(struct pingpong_dest));
	rem_dest->lid   = ntohl(a_rem_dest.lid);
	rem_dest->out_reads     = ntohl(a_rem_dest.out_reads);
	rem_dest->qpn   = ntohl(a_rem_dest.qpn);
	rem_dest->psn   = ntohl(a_rem_dest.psn);
	rem_dest->rkey  = ntohl(a_rem_dest.rkey);

	rem_dest->vaddr         = be64toh(a_rem_dest.vaddr);
	memcpy(rem_dest->gid.raw, &(a_rem_dest.gid), 16*sizeof(uint8_t));
	#else
	memcpy(&rem_dest,comm->rdma_ctx->buf,sizeof(struct pingpong_dest));
	#endif

	if (post_one_recv_wqe(comm->rdma_ctx)) {
		fprintf(stderr, "Couldn't post send \n");
		return 1;
	}

	return 0;
}

void RDMAWriteSocket::send_close() {

    Buffer send_buf = this->rsock->get_send_buf();

    // clear send cq
    struct ibv_wc wc[PACKET_WINDOW_SIZE];
    //this->rsock->poll_send_cq(PACKET_WINDOW_SIZE, wc);
    
    int ret = ibv_poll_cq(this->rsock->client_id->send_cq, PACKET_WINDOW_SIZE, wc);
    if (ret < 0) {
	perror("ibv_poll_cq");
        exit(1);
    }
    
    // send close msg
    MessageHeader header(MessageType::CLOSE, 0);
    int is_arrived = 0xffffffff;
    send_buf.write(header).write(is_arrived);
    this->rsock->post_write(send_buf, this->rka);

    // check send
    struct ibv_wc close_wc;
    this->rsock->poll_send_cq(1, &close_wc);
}

void  poll_cq(void *ctx)
{
  struct ibv_cq *cq;
  struct ibv_wc wc;
  int ne;

    TEST_NZ(ibv_get_cq_event(s_ctx->comp_channel, &cq, &ctx));//block by default
    ibv_ack_cq_events(cq, 1);
    TEST_NZ(ibv_req_notify_cq(cq, 0));

    do {
        ne = ibv_poll_cq(cq, 1, &wc);
        if(ne < 0){
            printf("fail to poll completion from the CQ. ret = %d\n", ne);
            return;
        }
        else if(ne == 0)
            continue;
        else
            on_completion(&wc);
    } while (ne == 0);

  return;
}