libhashkit: fix #25

[awesomized/libmemcached] / libtest / port.cc

diff --git a/libtest/port.cc b/libtest/port.cc
index 9f9a5bfeede83d9053a1504b2f44cc70bc5e1ab5..d8e8c12d566beca74fc96836a9f0635b67342b5a 100644 (file)
--- a/libtest/port.cc
+++ b/libtest/port.cc
@@ -34,7 +34,7 @@
  *
  */
 
-#include <config.h>
+#include "libtest/yatlcon.h"
 #include <libtest/common.h>
 
 #include <cassert>
@@ -50,10 +50,25 @@
 #include <sys/wait.h>
 #include <unistd.h>
 
+#include <utility>
+#include <vector>
+
 #include <signal.h>
 
 #include <libtest/signal.h>
 
+#ifndef SOCK_CLOEXEC 
+#  define SOCK_CLOEXEC 0
+#endif
+
+#ifndef SOCK_NONBLOCK 
+#  define SOCK_NONBLOCK 0
+#endif
+
+#ifndef FD_CLOEXEC
+#  define FD_CLOEXEC 0
+#endif
+
 #ifndef __INTEL_COMPILER
 #pragma GCC diagnostic ignored "-Wold-style-cast"
 #endif
@@ -61,13 +76,36 @@
 using namespace libtest;
 
 struct socket_st {
-  std::vector<int> fd;
+  typedef std::vector< std::pair< int, in_port_t> > socket_port_t;
+  socket_port_t _pair;
+  in_port_t last_port;
+
+  socket_st():
+    last_port(0)
+  { }
+
+  void release(in_port_t _arg)
+  {
+    for (socket_port_t::iterator iter= _pair.begin();
+         iter != _pair.end();
+         ++iter)
+    {
+      if ((*iter).second == _arg)
+      {
+        shutdown((*iter).first, SHUT_RDWR);
+        close((*iter).first);
+      }
+    }
+  }
 
   ~socket_st()
   {
-    for(std::vector<int>::iterator iter= fd.begin(); iter != fd.end(); iter++)
+    for (socket_port_t::iterator iter= _pair.begin();
+         iter != _pair.end();
+         ++iter)
     {
-      close(*iter);
+      shutdown((*iter).first, SHUT_RDWR);
+      close((*iter).first);
     }
   }
 };
@@ -88,52 +126,105 @@ in_port_t default_port()
   return global_port;
 }
 
+void release_port(in_port_t arg)
+{
+  all_socket_fd.release(arg);
+}
+
 in_port_t get_free_port()
 {
-  in_port_t ret_port= in_port_t(0);
+  const in_port_t default_port= in_port_t(-1);
 
   int retries= 1024;
 
-  while (retries--)
+  in_port_t ret_port;
+  while (--retries)
   {
+    ret_port= default_port;
     int sd;
-    if ((sd= socket(AF_INET, SOCK_STREAM, 0)) != -1)
+    if ((sd= socket(AF_INET, SOCK_STREAM, 0)) != SOCKET_ERROR)
     {
       int optval= 1;
-      if (setsockopt(sd, SOL_SOCKET, SO_REUSEADDR, &optval, sizeof(optval)) != -1)
-      {
+      if (setsockopt(sd, SOL_SOCKET, SO_REUSEADDR, &optval, sizeof(optval)) != SOCKET_ERROR)
+      { 
         struct sockaddr_in sin;
         sin.sin_port= 0;
         sin.sin_addr.s_addr= 0;
         sin.sin_addr.s_addr= INADDR_ANY;
         sin.sin_family= AF_INET;
 
-        if (bind(sd, (struct sockaddr *)&sin,sizeof(struct sockaddr_in) ) != -1)
+        int bind_ret;
+        do
         {
-          socklen_t addrlen= sizeof(sin);
+          if ((bind_ret= bind(sd, (struct sockaddr *)&sin, sizeof(struct sockaddr_in) )) != SOCKET_ERROR)
+          {
+            socklen_t addrlen= sizeof(sin);
+
+            if (getsockname(sd, (struct sockaddr *)&sin, &addrlen) != -1)
+            {
+              ret_port= sin.sin_port;
+            }
+          }
+          else
+          {
+            if (errno != EADDRINUSE)
+            {
+              Error << strerror(errno);
+            }
+          }
 
-          if (getsockname(sd, (struct sockaddr *)&sin, &addrlen) != -1)
+          if (errno == EADDRINUSE)
           {
-            ret_port= sin.sin_port;
+            libtest::dream(2, 0);
           }
-        }
-      }
+        } while (bind_ret == -1 and errno == EADDRINUSE);
 
-      all_socket_fd.fd.push_back(sd);
+        all_socket_fd._pair.push_back(std::make_pair(sd, ret_port));
+      }
+      else
+      {
+        Error << strerror(errno);
+      }
+    }
+    else
+    {
+      Error << strerror(errno);
     }
 
-    if (ret_port > 1024)
+    if (ret_port == default_port)
+    {
+      Error << "no ret_port set:" << strerror(errno);
+    }
+    else if (ret_port > 1024 and ret_port != all_socket_fd.last_port)
     {
       break;
     }
   }
 
   // We handle the case where if we max out retries, we still abort.
+  if (retries == 0)
+  {
+    FATAL("No port could be found, exhausted retry");
+  }
+
+  if (ret_port == 0)
+  {
+    FATAL("No port could be found");
+  }
+
+  if (ret_port == default_port)
+  {
+    FATAL("No port could be found");
+  }
+
   if (ret_port <= 1024)
   {
-    fatal_message("No port could be found");
+    FATAL("No port could be found, though some where available below or at 1024");
   }
 
+  all_socket_fd.last_port= ret_port;
+  release_port(ret_port);
+
   return ret_port;
 }