[Sheepdog] [PATCH] count the number of requests to free a client_info safely

[MORITA Kazutaka]

When a connection is closed, we cannot free a client_info if server is
processing a request which references the client_info.

This patch introduces a reference count to check whether we can free
the client_info safely.

Signed-off-by: MORITA Kazutaka <morita.kazutaka at lab.ntt.co.jp>
---
 include/event.h    |    2 +-
 include/net.h      |    4 ++--
 lib/event.c        |    9 ++++++---
 lib/net.c          |    8 ++++----
 sheep/sdnet.c      |   36 ++++++++++++++++++++++++++++++++----
 sheep/sheep_priv.h |    2 ++
 6 files changed, 47 insertions(+), 14 deletions(-)

diff --git a/include/event.h b/include/event.h
index d3f0004..b16c97a 100644
--- a/include/event.h
+++ b/include/event.h
@@ -10,7 +10,7 @@ typedef void (*event_handler_t)(int fd, int events, void *data);
 int init_event(int nr);
 int register_event(int fd, event_handler_t h, void *data);
 void unregister_event(int fd);
-void modify_event(int fd, unsigned int events);
+int modify_event(int fd, unsigned int events);
 void event_loop(int timeout);
 
 struct timer {
diff --git a/include/net.h b/include/net.h
index b0bc399..539701e 100644
--- a/include/net.h
+++ b/include/net.h
@@ -25,8 +25,8 @@ struct connection {
 	struct sd_rsp tx_hdr;
 };
 
-void conn_tx_off(struct connection *conn);
-void conn_tx_on(struct connection *conn);
+int conn_tx_off(struct connection *conn);
+int conn_tx_on(struct connection *conn);
 int is_conn_dead(struct connection *conn);
 int do_read(int sockfd, void *buf, int len);
 int rx(struct connection *conn, enum conn_state next_state);
diff --git a/lib/event.c b/lib/event.c
index a24d3cb..6991b6c 100644
--- a/lib/event.c
+++ b/lib/event.c
@@ -141,7 +141,7 @@ void unregister_event(int fd)
 	free(ei);
 }
 
-void modify_event(int fd, unsigned int events)
+int modify_event(int fd, unsigned int events)
 {
 	int ret;
 	struct epoll_event ev;
@@ -150,7 +150,7 @@ void modify_event(int fd, unsigned int events)
 	ei = lookup_event(fd);
 	if (!ei) {
 		eprintf("can't find event info %d\n", fd);
-		return;
+		return 1;
 	}
 
 	memset(&ev, 0, sizeof(ev));
@@ -158,8 +158,11 @@ void modify_event(int fd, unsigned int events)
 	ev.data.ptr = ei;
 
 	ret = epoll_ctl(efd, EPOLL_CTL_MOD, fd, &ev);
-	if (ret)
+	if (ret) {
 		eprintf("can't del epoll event, %m\n");
+		return 1;
+	}
+	return 0;
 }
 
 void event_loop(int timeout)
diff --git a/lib/net.c b/lib/net.c
index 9b260ab..107408c 100644
--- a/lib/net.c
+++ b/lib/net.c
@@ -29,14 +29,14 @@
 #include "net.h"
 #include "logger.h"
 
-void conn_tx_off(struct connection *conn)
+int conn_tx_off(struct connection *conn)
 {
-	modify_event(conn->fd, EPOLLIN);
+	return modify_event(conn->fd, EPOLLIN);
 }
 
-void conn_tx_on(struct connection *conn)
+int conn_tx_on(struct connection *conn)
 {
-	modify_event(conn->fd, EPOLLIN|EPOLLOUT);
+	return modify_event(conn->fd, EPOLLIN|EPOLLOUT);
 }
 
 int is_conn_dead(struct connection *conn)
diff --git a/sheep/sdnet.c b/sheep/sdnet.c
index 79ba3ff..0d3d4b1 100644
--- a/sheep/sdnet.c
+++ b/sheep/sdnet.c
@@ -238,6 +238,9 @@ static void queue_request(struct request *req)
 	start_cpg_event_work();
 }
 
+static void client_incref(struct client_info *ci);
+static void client_decref(struct client_info *ci);
+
 static struct request *alloc_request(struct client_info *ci, int data_length)
 {
 	struct request *req;
@@ -247,6 +250,7 @@ static struct request *alloc_request(struct client_info *ci, int data_length)
 		return NULL;
 
 	req->ci = ci;
+	client_incref(ci);
 	if (data_length)
 		req->data = (char *)req + sizeof(*req);
 
@@ -258,6 +262,7 @@ static struct request *alloc_request(struct client_info *ci, int data_length)
 
 static void free_request(struct request *req)
 {
+	client_decref(req->ci);
 	list_del(&req->r_siblings);
 	free(req);
 }
@@ -265,7 +270,10 @@ static void free_request(struct request *req)
 static void req_done(struct request *req)
 {
 	list_add(&req->r_wlist, &req->ci->done_reqs);
-	conn_tx_on(&req->ci->conn);
+	if (conn_tx_on(&req->ci->conn)) {
+		dprintf("connection seems to be dead\n");
+		free_request(req);
+	}
 }
 
 static void init_rx_hdr(struct client_info *ci)
@@ -317,7 +325,12 @@ static void client_rx_handler(struct client_info *ci)
 		eprintf("BUG: unknown state %d\n", conn->c_rx_state);
 	}
 
-	if (is_conn_dead(conn) || conn->c_rx_state != C_IO_END)
+	if (is_conn_dead(conn) && ci->rx_req) {
+		free_request(ci->rx_req);
+		return;
+	}
+
+	if (conn->c_rx_state != C_IO_END)
 		return;
 
 	/* now we have a complete command */
@@ -406,8 +419,10 @@ again:
 	opt = 0;
 	setsockopt(ci->conn.fd, SOL_TCP, TCP_CORK, &opt, sizeof(opt));
 
-	if (is_conn_dead(&ci->conn) || ci->conn.c_tx_state != C_IO_END)
+	if (is_conn_dead(&ci->conn)) {
+		free_request(ci->tx_req);
 		return;
+	}
 
 	if (ci->conn.c_tx_state == C_IO_END) {
 		free_request(ci->tx_req);
@@ -422,6 +437,18 @@ static void destroy_client(struct client_info *ci)
 	free(ci);
 }
 
+static void client_incref(struct client_info *ci)
+{
+	if (ci)
+		ci->refcnt++;
+}
+
+static void client_decref(struct client_info *ci)
+{
+	if (ci && --ci->refcnt == 0)
+		destroy_client(ci);
+}
+
 static struct client_info *create_client(int fd, struct cluster_info *cluster)
 {
 	struct client_info *ci;
@@ -431,6 +458,7 @@ static struct client_info *create_client(int fd, struct cluster_info *cluster)
 		return NULL;
 
 	ci->conn.fd = fd;
+	ci->refcnt = 1;
 
 	INIT_LIST_HEAD(&ci->reqs);
 	INIT_LIST_HEAD(&ci->done_reqs);
@@ -453,7 +481,7 @@ static void client_handler(int fd, int events, void *data)
 	if (is_conn_dead(&ci->conn)) {
 		dprintf("closed a connection, %d\n", fd);
 		unregister_event(fd);
-		destroy_client(ci);
+		client_decref(ci);
 	}
 }
 
diff --git a/sheep/sheep_priv.h b/sheep/sheep_priv.h
index 8d1c6c0..47c83c1 100644
--- a/sheep/sheep_priv.h
+++ b/sheep/sheep_priv.h
@@ -59,6 +59,8 @@ struct client_info {
 
 	struct list_head reqs;
 	struct list_head done_reqs;
+
+	int refcnt;
 };
 
 struct request;
-- 
1.5.6.5