RDMA 编程客户端client demo

/*
Author: ben.wei
Date: 01-25 2022
A RDMA read/write demo for test
client code
*/

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <rdma/rdma_cma.h>

#include <netdb.h>

#define TEST_NZ(x) \
	do { if ((x)) die("err: " #x " failed(return non-zero)");} while(0)

#define TEST_Z(x) \
	do { if (!(x)) die("err: " #x " failed(return null/zero)");} while(0)

#define TIMEOUT_IN_MS 500 // ms
			  //
#define M_WRITE 1
#define M_READ 2

#define M_WRITE_COMPLETED 3
#define M_READ_COMPLETED 4

int s_mode = 0;
struct message {
	int opcode;
	union {
		struct ibv_mr mr;
	}data;
};
struct connection {
	struct rdma_cm_id *id;
	struct ibv_qp *qp;

	int connected;

	struct ibv_mr *recv_mr;
	struct ibv_mr *send_mr;
	struct ibv_mr *rdma_local_mr;
	struct ibv_mr *rdma_remote_mr;

	struct ibv_mr peer_mr;

	struct message *recv_msg;
	struct message *send_msg;

	char* rdma_local_region;
	char* rdma_remote_region;
};

struct privdata {
	int version;
};
static void die(const char *reason)
{
	fprintf(stderr, "%s\n", reason);
	exit(EXIT_FAILURE);
}
int on_event(struct rdma_cm_event *event);
int main(int argc, char **argv)
{
	struct addrinfo *addr;
	struct rdma_cm_event *event = NULL;
	struct rdma_cm_id *conn = NULL;
	struct rdma_event_channel *ec = NULL;

	if(argc != 3) {
		die("usage client serverip serverport");
	}
	TEST_NZ(getaddrinfo(argv[1], argv[2], NULL, &addr));
	TEST_Z(ec = rdma_create_event_channel());
	TEST_NZ(rdma_create_id(ec,&conn,NULL,RDMA_PS_TCP));
	TEST_NZ(rdma_resolve_addr(conn, NULL, addr->ai_addr, TIMEOUT_IN_MS));
	freeaddrinfo(addr);

	while(rdma_get_cm_event(ec, &event)==0){
		struct rdma_cm_event event_copy;

		memcpy(&event_copy, event, sizeof(*event));	
		rdma_ack_cm_event(event);
		if(on_event(&event_copy)){
			break;
		}
	}
	rdma_destroy_event_channel(ec);

	return 0;
}
static int on_addr_resolved(struct rdma_cm_id *id);
static int on_route_resolved(struct rdma_cm_id *id);
static int on_disconnect(struct rdma_cm_id *id);
static int on_connection(void *context);

int on_event(struct rdma_cm_event *event)
{
	int r = 0;
	if(event->event == RDMA_CM_EVENT_ADDR_RESOLVED)
		r = on_addr_resolved(event->id);
	else if(event->event == RDMA_CM_EVENT_ROUTE_RESOLVED)
		r = on_route_resolved(event->id);
	else if(event->event == RDMA_CM_EVENT_ESTABLISHED)
		r = on_connection(event->id->context);
	else if(event->event = RDMA_CM_EVENT_DISCONNECTED)
		r = on_disconnect(event->id);
	else
		die("on event unknown event");

	return r;
}

struct context {
	struct ibv_context *ctx;
	struct ibv_pd *pd;
	struct ibv_cq *cq;
	struct ibv_comp_channel *comp_channel;

	pthread_t cq_poller_thread;
};
static struct context *s_ctx = NULL;
static void* poll_cq(void*);
void build_context(struct ibv_context *verbs)
{
	if(s_ctx){
		if(s_ctx->ctx != verbs)
			die("cannot handle events in more than one context");
		return;
	}
	s_ctx = (struct context*)malloc(sizeof(struct context));
	s_ctx->ctx = verbs;

	TEST_Z(s_ctx->pd = ibv_alloc_pd(s_ctx->ctx));
	TEST_Z(s_ctx->comp_channel = ibv_create_comp_channel(s_ctx->ctx));
	TEST_Z(s_ctx->cq = ibv_create_cq(s_ctx->ctx, 10, NULL, s_ctx->comp_channel, 0));
	TEST_NZ(ibv_req_notify_cq(s_ctx->cq, 0));

	TEST_NZ(pthread_create(&s_ctx->cq_poller_thread, NULL, poll_cq, NULL));
}
static void on_completion(struct ibv_wc *wc);
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;
}
void build_qp_attr(struct ibv_qp_init_attr *qp_attr)
{
	memset(qp_attr, 0 ,sizeof(*qp_attr));

	qp_attr->send_cq = s_ctx->cq;
	qp_attr->recv_cq = s_ctx->cq;
	qp_attr->qp_type = IBV_QPT_RC;

	qp_attr->cap.max_send_wr = 1;
	qp_attr->cap.max_recv_wr = 1;
	qp_attr->cap.max_send_sge = 1;
	qp_attr->cap.max_recv_sge = 1;
}
const int BUFFER_SIZE = 1024;
void register_memory(struct connection *conn)
{
	conn->send_msg = malloc(sizeof(struct message));
	conn->recv_msg = malloc(sizeof(struct message));
	conn->rdma_local_region = malloc(BUFFER_SIZE);
	conn->rdma_remote_region = malloc(BUFFER_SIZE);
	
	TEST_Z(conn->send_mr = ibv_reg_mr(s_ctx->pd, conn->send_msg, sizeof(struct message), IBV_ACCESS_LOCAL_WRITE|IBV_ACCESS_REMOTE_WRITE));	
	TEST_Z(conn->recv_mr = ibv_reg_mr(s_ctx->pd, conn->recv_msg, sizeof(struct message), IBV_ACCESS_LOCAL_WRITE|IBV_ACCESS_REMOTE_WRITE));	
	TEST_Z(conn->rdma_local_mr = ibv_reg_mr(s_ctx->pd, conn->rdma_local_region, BUFFER_SIZE, IBV_ACCESS_LOCAL_WRITE|IBV_ACCESS_REMOTE_WRITE|IBV_ACCESS_REMOTE_READ));	
	TEST_Z(conn->rdma_remote_mr = ibv_reg_mr(s_ctx->pd, conn->rdma_remote_region, BUFFER_SIZE, IBV_ACCESS_LOCAL_WRITE|IBV_ACCESS_REMOTE_READ|IBV_ACCESS_REMOTE_WRITE));	
}

void post_receives(struct connection *conn)
{
	struct ibv_recv_wr wr, *bad_wr = NULL;
	struct ibv_sge sge;

	wr.wr_id = (uintptr_t)conn;
	wr.next = NULL;
	wr.sg_list = &sge;
	wr.num_sge = 1;

	sge.addr = (uintptr_t)conn->recv_msg;
	sge.length = sizeof(struct message);
	sge.lkey = conn->recv_mr->lkey;

	TEST_NZ(ibv_post_recv(conn->qp, &wr, &bad_wr));
}
int on_addr_resolved(struct rdma_cm_id *id)
{
	struct ibv_qp_init_attr qp_attr;
	struct connection *conn;
	printf("addr resolved\n");

	build_context(id->verbs);
	build_qp_attr(&qp_attr);
	TEST_NZ(rdma_create_qp(id, s_ctx->pd, &qp_attr));

	id->context = conn = (struct connection*)malloc(sizeof(struct connection));
	conn->id = id;
	conn->qp = id->qp;
	register_memory(conn);
	post_receives(conn);
	TEST_NZ(rdma_resolve_route(id, TIMEOUT_IN_MS));
	return 0;
}
int on_route_resolved(struct rdma_cm_id *id)
{
	struct rdma_conn_param cm_param;
	struct privdata *p = (struct privdata*)malloc(sizeof(struct privdata));
	p->version = 0x0101;
	printf("on route resolved, version:%x\n", p->version);

	memset(&cm_param, 0, sizeof(cm_param));
	cm_param.private_data = p;
	cm_param.private_data_len = sizeof(struct privdata);
	TEST_NZ(rdma_connect(id, &cm_param));

	return 0;
}
void send_message(struct connection *conn)
{
	struct ibv_send_wr wr, *bad_wr = NULL;	
	struct ibv_sge sge;

	memset(&wr, 0, sizeof(wr));
	wr.wr_id = (uintptr_t)conn;
	wr.opcode=IBV_WR_SEND;
	wr.sg_list = &sge;
	wr.num_sge = 1;
	wr.send_flags = IBV_SEND_SIGNALED;

	sge.addr = (uintptr_t)conn->send_msg;
	sge.length = sizeof(struct message);
	sge.lkey = conn->send_mr->lkey;
	while(!conn->connected);

	TEST_NZ(ibv_post_send(conn->qp, &wr, &bad_wr));
}
void send_write_op(struct connection *conn)
{
//	struct connection *conn = (struct connection*)context;
	conn->send_msg->opcode = M_WRITE;
	strcpy(conn->rdma_remote_region, "message from client side");
	memcpy(&conn->send_msg->data.mr, conn->rdma_remote_mr, sizeof(struct ibv_mr));
	send_message(conn);
}
void send_read_op(struct connection *conn)
{
	//struct connection *conn = (struct connection*)context;
	conn->send_msg->opcode = M_READ;
	memcpy(&conn->send_msg->data.mr, conn->rdma_local_mr, sizeof(struct ibv_mr));
	send_message(conn);
}

int on_connection(void *context)
{
	struct connection *conn = (struct connection *)context;
	struct ibv_send_wr wr, *bad_wr = NULL;
	struct ibv_sge sge;

	//snprintf(conn->send_region, BUFFER_SIZE, "message from client size with pid: %d\n", getpid());
	printf("connected, posting  read send...\n");
	conn->connected = 1;
	//send_read_op(context);
	send_write_op(conn);
/*
	memset(&wr, 0, sizeof(wr));

	wr.wr_id = (uintptr_t)conn;
	wr.opcode = IBV_WR_SEND;
	wr.sg_list = &sge;
	wr.num_sge = 1;
	wr.send_flags = IBV_SEND_SIGNALED;

	sge.addr = (uintptr_t)conn->rdma_local_region;
	sge.length = BUFFER_SIZE;
	sge.lkey = conn->send_mr->lkey;

	TEST_NZ(ibv_post_send(conn->qp,&wr,&bad_wr));
	*/
	return 0;
}
int on_disconnect(struct rdma_cm_id *id)
{
	struct connection *conn = (struct connection*)id->context;
	printf("disconnect now \n");

	rdma_destroy_qp(id);
	ibv_dereg_mr(conn->send_mr);
	ibv_dereg_mr(conn->recv_mr);
	ibv_dereg_mr(conn->rdma_local_mr);
	ibv_dereg_mr(conn->rdma_remote_mr);

	free(conn->send_msg);
	free(conn->recv_msg);
	free(conn->rdma_local_region);
	free(conn->rdma_remote_region);
	free(conn);
	rdma_destroy_id(id);
	return 1;
}
void on_completion(struct ibv_wc *wc)
{
	struct connection *conn = (struct connection*)(uintptr_t)wc->wr_id;

	if(wc->status != IBV_WC_SUCCESS)
		die("on completion status is not success");
	if(wc->opcode & IBV_WC_RECV){
		printf("on_completion 1 received message op:%d\n", conn->recv_msg->opcode);
		if(conn->recv_msg->opcode == M_READ_COMPLETED){
			printf("on_completion read data from server: %s\n", conn->rdma_local_region);
		}
		printf("on_completion 2 received message op:%d\n", conn->recv_msg->opcode);

	}else if(wc->opcode == IBV_WC_SEND){
		printf("on_completion send successfully, msg op:%d\n", conn->send_msg->opcode);
		if(conn->send_msg->opcode == M_READ)
			send_write_op(conn);
		else if(conn->send_msg->opcode == M_WRITE){
			send_read_op(conn);
		}
	}
	else
		die("on completion: completion is not a send or receive");
}