socket.onclose() (#1364)

* Add socket.onclose() and close socket immediately when socket channel closed by peer

* Revert read_socket(), Add SOCKET_MORE, See #1358

* remove warning log , because it's not a rare case any more. See #1358

lualib/skynet/socket.lua

@@ -123,6 +123,9 @@ socket_message[3] = function(id)
 	end
 	s.connected = false
 	wakeup(s)
+	if s.on_close then
+		s.on_close(id)
+	end
 end
 
 -- SKYNET_SOCKET_TYPE_ACCEPT = 4
@@ -485,4 +488,10 @@ function socket.warning(id, callback)
 	obj.on_warning = callback
 end
 
+function socket.onclose(id, callback)
+	local obj = socket_pool[id]
+	assert(obj)
+	obj.on_close = callback
+end
+
 return socket

lualib/skynet/socketchannel.lua

@@ -128,6 +128,16 @@ local function pop_response(self)
 	end
 end
 
+-- on close callback
+local function autoclose_cb(self, fd)
+	local sock = self.__sock
+	if self.__wait_response and sock and sock[1] == fd then
+		-- closed by peer
+		skynet.error("socket closed by peer : ", self.__host, self.__port)
+		close_channel_socket(self)
+	end
+end
+
 local function push_response(self, response, co)
 	if self.__response then
 		-- response is session
@@ -170,7 +180,15 @@ local function dispatch_by_order(self)
 			wakeup_all(self, "channel_closed")
 			break
 		end
-		local ok, result_ok, result_data = pcall(get_response, func, self.__sock)
+		local sock = self.__sock
+		if not sock then
+			-- closed by peer
+			self.__result[co] = socket_error
+			skynet.wakeup(co)
+			wakeup_all(self)
+			break
+		end
+		local ok, result_ok, result_data = pcall(get_response, func, sock)
 		if ok then
 			self.__result[co] = result_ok
 			if result_ok and self.__result_data[co] then
@@ -197,6 +215,9 @@ local function dispatch_function(self)
 	if self.__response then
 		return dispatch_by_session
 	else
+		socket.onclose(self.__sock[1], function(fd)
+			autoclose_cb(self, fd)
+		end)
 		return dispatch_by_order
 	end
 end

skynet-src/socket_server.c

@@ -1408,74 +1408,17 @@ ctrl_cmd(struct socket_server *ss, struct socket_message *result) {
 	return -1;
 }
 
-struct stream_buffer {
-	char * buf;
-	int sz;
-};
-
-static char *
-reserve_buffer(struct stream_buffer *buffer, int sz) {
-	if (buffer->buf == NULL) {
-		buffer->buf = (char *)MALLOC(sz);
-		return buffer->buf;
-	} else {
-		char * newbuffer = (char *)MALLOC(sz + buffer->sz);
-		memcpy(newbuffer, buffer->buf, buffer->sz);
-		char * ret = newbuffer + buffer->sz;
-		FREE(buffer->buf);
-		buffer->buf = newbuffer;
-		return ret;
-	}
-}
-
-static int
-read_socket(struct socket *s, struct stream_buffer *buffer) {
-	int sz = s->p.size;
-	buffer->buf = NULL;
-	buffer->sz = 0;
-	int rsz = sz;
-	for (;;) {
-		char *buf = reserve_buffer(buffer, rsz);
-		int n = (int)read(s->fd, buf, rsz);
-		if (n <= 0) {
-			if (buffer->sz == 0) {
-				// read nothing
-				FREE(buffer->buf);
-				return n;
-			} else {
-				// ignore the error or hang up, returns buffer
-				// If socket is hang up, SOCKET_CLOSE will be send later.
-				//    (buffer->sz should be 0 next time)
-				break;
-			}
-		}
-		buffer->sz += n;
-		if (n < rsz) {
-			break;
-		}
-		// n == rsz, read again ( and read more )
-		rsz *= 2;
-	}
-	int r = buffer->sz;
-	if (r > sz) {
-		s->p.size = sz * 2;
-	} else if (sz > MIN_READ_BUFFER && r*2 < sz) {
-		s->p.size = sz / 2;
-	}
-	return r;
-}
-
 // return -1 (ignore) when error
 static int
 forward_message_tcp(struct socket_server *ss, struct socket *s, struct socket_lock *l, struct socket_message * result) {
-	struct stream_buffer buf;
+	int sz = s->p.size;
-	int n = read_socket(s, &buf);
+	char * buffer = MALLOC(sz);
+	int n = (int)read(s->fd, buffer, sz);
 	if (n<0) {
+		FREE(buffer);
 		switch(errno) {
 		case EINTR:
-			break;
 		case AGAIN_WOULDBLOCK:
-			skynet_error(NULL, "socket-server: EAGAIN capture.");
 			break;
 		default:
 			// close when error
@@ -1486,6 +1429,7 @@ forward_message_tcp(struct socket_server *ss, struct socket *s, struct socket_lo
 		return -1;
 	}
 	if (n==0) {
+		FREE(buffer);
 		if (s->closing) {
 			// Rare case : if s->closing is true, reading event is disable, and SOCKET_CLOSE is raised.
 			if (nomore_sending_data(s)) {
@@ -1509,7 +1453,7 @@ forward_message_tcp(struct socket_server *ss, struct socket *s, struct socket_lo
 
 	if (halfclose_read(s)) {
 		// discard recv data (Rare case : if socket is HALFCLOSE_READ, reading event is disable.)
-		FREE(buf.buf);
+		FREE(buffer);
 		return -1;
 	}
 
@@ -1518,7 +1462,15 @@ forward_message_tcp(struct socket_server *ss, struct socket *s, struct socket_lo
 	result->opaque = s->opaque;
 	result->id = s->id;
 	result->ud = n;
-	result->data = buf.buf;
+	result->data = buffer;
+
+	if (n == sz) {
+		s->p.size *= 2;
+		return SOCKET_MORE;
+	} else if (sz > MIN_READ_BUFFER && n*2 < sz) {
+		s->p.size /= 2;
+	}
+
 	return SOCKET_DATA;
 }
 
@@ -1757,6 +1709,10 @@ socket_server_poll(struct socket_server *ss, struct socket_message * result, int
 				int type;
 				if (s->protocol == PROTOCOL_TCP) {
 					type = forward_message_tcp(ss, s, &l, result);
+					if (type == SOCKET_MORE) {
+						--ss->event_index;
+						return SOCKET_DATA;
+					}
 				} else {
 					type = forward_message_udp(ss, s, &l, result);
 					if (type == SOCKET_UDP) {

skynet-src/socket_server.h

@@ -13,7 +13,10 @@
 #define SOCKET_EXIT 5
 #define SOCKET_UDP 6
 #define SOCKET_WARNING 7
-#define SOCKET_RST 8 // Only for internal use
+
+// Only for internal use
+#define SOCKET_RST 8
+#define SOCKET_MORE 9
 
 struct socket_server;