properly handle non-blocking mode of socket

addressing #171

src/allocator_thread.c

@@ -176,24 +176,61 @@ static int wait_data(int readfd) {
 	return 1;
 }
 
+static int trywrite(int fd, void* buf, size_t bytes) {
+	ssize_t ret;
+	unsigned char *out = buf;
+again:
+	ret = write(fd, out, bytes);
+	switch(ret) {
+		case -1:
+			if(errno == EINTR) goto again;
+		case  0:
+			return 0;
+		default:
+			if(ret == bytes || !bytes) return 1;
+			out += ret;
+			bytes -= ret;
+			goto again;
+	}
+}
+
 static int sendmessage(enum at_direction dir, struct at_msghdr *hdr, void* data) {
 	static int* destfd[ATD_MAX] = { [ATD_SERVER] = &req_pipefd[1], [ATD_CLIENT] = &resp_pipefd[1] };
-	int ret = write(*destfd[dir], hdr, sizeof *hdr) == sizeof *hdr;
+	int ret = trywrite(*destfd[dir], hdr, sizeof *hdr);
 	if(ret && hdr->datalen) {
 		assert(hdr->datalen <= MSG_LEN_MAX);
-		ret = write(*destfd[dir], data, hdr->datalen) == hdr->datalen;
+		ret = trywrite(*destfd[dir], data, hdr->datalen);
 	}
 	return ret;
 }
 
+static int tryread(int fd, void* buf, size_t bytes) {
+	ssize_t ret;
+	unsigned char *out = buf;
+again:
+	ret = read(fd, out, bytes);
+	switch(ret) {
+		case -1:
+			if(errno == EINTR) goto again;
+		case  0:
+			return 0;
+		default:
+			if(ret == bytes || !bytes) return 1;
+			out += ret;
+			bytes -= ret;
+			goto again;
+	}
+}
+
 static int getmessage(enum at_direction dir, struct at_msghdr *hdr, void* data) {
 	static int* readfd[ATD_MAX] = { [ATD_SERVER] = &req_pipefd[0], [ATD_CLIENT] = &resp_pipefd[0] };
-	int ret;
+	ssize_t ret;
 	if((ret = wait_data(*readfd[dir]))) {
-		ret = read(*readfd[dir], hdr, sizeof *hdr) == sizeof(*hdr);
+		if(!tryread(*readfd[dir], hdr, sizeof *hdr))
+			return 0;
 		assert(hdr->datalen <= MSG_LEN_MAX);
-		if(ret && hdr->datalen) {
-			ret = read(*readfd[dir], data, hdr->datalen) == hdr->datalen;
+		if(hdr->datalen) {
+			ret = tryread(*readfd[dir], data, hdr->datalen);
 		}
 	}
 	return ret;

src/core.c

@@ -154,10 +154,19 @@ static int timed_connect(int sock, const struct sockaddr *addr, socklen_t len) {
 
 	pfd[0].fd = sock;
 	pfd[0].events = POLLOUT;
-	fcntl(sock, F_SETFL, O_NONBLOCK);
+	int flags = fcntl(sock, F_GETFL, 0);
+	/* put socket temporarily into nonblocking mode so we can enforce
+	 * the timeout. */
+	if(!(flags & O_NONBLOCK))
+		fcntl(sock, F_SETFL, flags | O_NONBLOCK);
 	ret = true_connect(sock, addr, len);
 	PDEBUG("\nconnect ret=%d\n", ret);
-	
+
+	/* if the socket was already non-blocking, we assume the app takes
+	 * care of handling the timeouts itself. */
+	if(flags & O_NONBLOCK)
+		return ret;
+
 	if(ret == -1 && errno == EINPROGRESS) {
 		ret = poll_retry(pfd, 1, tcp_connect_time_out);
 		PDEBUG("\npoll ret=%d\n", ret);
@@ -181,7 +190,7 @@ static int timed_connect(int sock, const struct sockaddr *addr, socklen_t len) {
 			ret = -1;
 	}
 
-	fcntl(sock, F_SETFL, !O_NONBLOCK);
+	fcntl(sock, F_SETFL, flags);
 	return ret;
 }