r231: Make client sockets non-blocking, too.

[nbd.git] / nbd-server.c

diff --git a/nbd-server.c b/nbd-server.c
index 986aa4b..281a14f 100644 (file)
--- a/nbd-server.c
+++ b/nbd-server.c
@@ -295,12 +295,33 @@ int authorized_client(CLIENT *opts) {
  **/
 inline void readit(int f, void *buf, size_t len) {
        ssize_t res;
+       gboolean tried = FALSE;
+
        while (len > 0) {
                DEBUG("*");
-               if ((res = read(f, buf, len)) <= 0)
-                       err("Read failed: %m");
-               len -= res;
-               buf += res;
+               if ((res = read(f, buf, len)) <= 0) {
+                       if(!tried && errno==EAGAIN) {
+                               /* Assume the connection will work some time in
+                                * the future, but don't run away with CPU time
+                                * in case it doesn't */
+                               fd_set set;
+                               struct timeval tv;
+
+                               DEBUG("Read failed, trying again");
+                               tried=TRUE;
+                               FD_ZERO(&set);
+                               FD_SET(f, &set);
+                               tv.tv_sec=30;
+                               tv.tv_usec=0;
+                               select(f+1, &set, NULL, NULL, &tv);
+                       } else {
+                               err("Read failed: %m");
+                       }
+               } else {
+                       len -= res;
+                       buf += res;
+                       tried=FALSE;
+               }
        }
 }
 
@@ -313,12 +334,33 @@ inline void readit(int f, void *buf, size_t len) {
  **/
 inline void writeit(int f, void *buf, size_t len) {
        ssize_t res;
+       gboolean tried=FALSE;
+
        while (len > 0) {
                DEBUG("+");
-               if ((res = write(f, buf, len)) <= 0)
-                       err("Send failed: %m");
-               len -= res;
-               buf += res;
+               if ((res = write(f, buf, len)) <= 0) {
+                       if(!tried && errno==EAGAIN) {
+                               /* Assume the connection will work some time in
+                                * the future, but don't run away with CPU time
+                                * in case it doesn't */
+                               fd_set set;
+                               struct timeval tv;
+
+                               DEBUG("Write failed, trying again");
+                               tried=TRUE;
+                               FD_ZERO(&set);
+                               FD_SET(f, &set);
+                               tv.tv_sec=30;
+                               tv.tv_usec=0;
+                               select(f+1, NULL, &set, NULL, &tv);
+                       } else {
+                               err("Send failed: %m");
+                       }
+               } else {
+                       len -= res;
+                       buf += res;
+                       tried=FALSE;
+               }
        }
 }
 
@@ -1292,7 +1334,7 @@ void daemonize(SERVER* serve) {
                        strncpy(pidftemplate, "/var/run/server.pid", 255);
                }
        }
-       snprintf(pidfname, 255, pidftemplate, serve->port);
+       snprintf(pidfname, 255, pidftemplate, serve ? serve->port : 0);
        pidf=fopen(pidfname, "w");
        if(pidf) {
                fprintf(pidf,"%d\n", (int)getpid());
@@ -1337,7 +1379,7 @@ void setup_serve(SERVER *serve) {
                err("fcntl F_GETFL");
        }
        if (fcntl(serve->socket, F_SETFL, sock_flags | O_NONBLOCK) == -1) {
-               err("fcntl F_SETFL O_NONBLOCK");
+               err("fcntl F_SETFL O_NONBLOCK on server socket");
        }
 
        DEBUG("Waiting for connections... bind, ");
@@ -1359,7 +1401,6 @@ void setup_serve(SERVER *serve) {
        sa.sa_flags = SA_RESTART;
        if(sigaction(SIGTERM, &sa, NULL) == -1)
                err("sigaction: %m");
-       children=g_hash_table_new_full(g_int_hash, g_int_equal, NULL, destroy_pid_t);
 }
 
 /**
@@ -1371,6 +1412,7 @@ void setup_servers(GArray* servers) {
        for(i=0;i<servers->len;i++) {
                setup_serve(&(g_array_index(servers, SERVER, i)));
        }
+       children=g_hash_table_new_full(g_int_hash, g_int_equal, NULL, destroy_pid_t);
 }
 
 /**
@@ -1414,12 +1456,20 @@ int serveloop(GArray* servers) {
                        for(i=0;i<servers->len;i++) {
                                serve=&(g_array_index(servers, SERVER, i));
                                if(FD_ISSET(serve->socket, &rset)) {
+                                       int sock_flags;
+
                                        if ((net=accept(serve->socket, (struct sockaddr *) &addrin, &addrinlen)) < 0)
                                                err("accept: %m");
 
                                        client = g_malloc(sizeof(CLIENT));
                                        client->server=serve;
                                        client->exportsize=OFFT_MAX;
+                                       if ((sock_flags = fcntl(serve->socket, F_GETFL, 0)) == -1) {
+                                               err("fcntl F_GETFL");
+                                       }
+                                       if (fcntl(net, F_SETFL, sock_flags | O_NONBLOCK) == -1) {
+                                               err("fcntl F_SETFL O_NONBLOCK on client socket");
+                                       }
                                        client->net=net;
                                        set_peername(net, client);
                                        if (!authorized_client(client)) {