Keep pool of pre-connected DNS sockets

net/dns/dns_session.cc

Index: net/dns/dns_session.cc
diff --git a/net/dns/dns_session.cc b/net/dns/dns_session.cc
index fb6c4b4c5ba269157bf7ca29cc7d5c06052ccda1..287303a1ec6dce0a6d21c0f95f7a80ca3cdc30b5 100644
--- a/net/dns/dns_session.cc
+++ b/net/dns/dns_session.cc
@@ -6,26 +6,111 @@
 
 #include "base/basictypes.h"
 #include "base/bind.h"
+#include "base/rand_util.h"
+#include "base/stl_util.h"
 #include "base/time.h"
 #include "net/base/ip_endpoint.h"
+#include "net/base/net_errors.h"
 #include "net/dns/dns_config_service.h"
 #include "net/socket/client_socket_factory.h"
 
 namespace net {
 
+DnsSession::SocketLease::SocketLease(scoped_refptr<DnsSession> session,
+                                     int server_index,
+                                     scoped_ptr<DatagramClientSocket> socket)
+    : session_(session), server_index_(server_index), socket_(socket.Pass()) {}
+
+DnsSession::SocketLease::~SocketLease() {
+  session_->FreeSocket(server_index_, socket_.Pass());
+}
+
+#if defined(OS_WIN)
+
+// Don't request specific (random) ports on Windows, since that will trigger
+// firewall prompts.
+static const DatagramSocket::BindType kBindType = DatagramSocket::DEFAULT_BIND;
+// Keep a buffer of the sockets the OS gives us to give some entropy in source
+// ports.
+static const size_t kMaxPoolSize = 256;
+// TODO(ttuttle): I want 1024, but we get ERR_INSUFFICIENT_RESOURCES if we
+//                try to allocate that many.
+
+#else
+
+static const DatagramSocket::BindType kBindType = DatagramSocket::RANDOM_BIND;
+static const size_t kMaxPoolSize = 0;
+
+#endif
+
 DnsSession::DnsSession(const DnsConfig& config,
                        ClientSocketFactory* factory,
                        const RandIntCallback& rand_int_callback,
                        NetLog* net_log)
     : config_(config),
       socket_factory_(factory),
-      rand_callback_(base::Bind(rand_int_callback, 0, kuint16max)),
+      rand_int_callback_(rand_int_callback),
       net_log_(net_log),
+      pools_(0),
       server_index_(0) {
+  CreateDefaultPools(kMaxPoolSize);
+}
+
+DnsSession::DnsSession(const DnsConfig& config,
+                       ClientSocketFactory* socket_factory,
+                       const RandIntCallback& rand_int_callback,
+                       int pool_size,
+                       NetLog* net_log)
+    : config_(config),
+      socket_factory_(socket_factory),
+      rand_int_callback_(rand_int_callback),
+      net_log_(net_log),
+      pools_(0),
+      server_index_(0) {
+  CreateDefaultPools(pool_size);
+}
+
+DnsSession::DnsSession(const DnsConfig& config,
+                       ClientSocketFactory* socket_factory,
+                       const RandIntCallback& rand_int_callback,
+                       const SocketPoolFactoryCallback& pool_factory_callback,
+                       NetLog* net_log)
+    : config_(config),
+      socket_factory_(socket_factory),
+      rand_int_callback_(rand_int_callback),
+      net_log_(net_log),
+      pools_(0),
+      server_index_(0) {
+  CreatePools(pool_factory_callback);
+}
+
+void DnsSession::CreatePools(
+    const SocketPoolFactoryCallback& pool_factory_callback) {
+  CHECK(pools_.empty());
+
+  for (unsigned i = 0; i < config_.nameservers.size(); i++) {
+    EndPointSocketFactoryCallback socket_factory =
+        base::Bind(&DnsSession::CreateConnectedSocket,
+                   this,
+                   config_.nameservers[i]);
+    SocketPool* pool = pool_factory_callback.Run(socket_factory);
+    pools_.push_back(pool);
+  }
+}
+
+void DnsSession::CreateDefaultPools(int pool_size) {
+  CreatePools(base::Bind(SocketPoolImpl::CreateSocketPool,
+                         pool_size,
+                         rand_int_callback_));
+}
+
+DnsSession::~DnsSession() {
+  STLDeleteElements(&pools_);
 }
 
 int DnsSession::NextQueryId() const {
-  return rand_callback_.Run();
+  /* NB: This callback is also used for selecting sockets from the pool. */

[szym, 2012/09/10 19:31:00]

I think the RNG used for selecting sockets from the pool should be internal to
SocketPool. If you intend to test if the order of selection matches the outputs
of the RNG, it would probably be best if you made SocketPool a proper class. It
can be an inner class (since it won't be used outside of DnsSession), but it
would be easier to test.

nit: code style forbids /* this style of comments */

+  return rand_int_callback_.Run(0, kuint16max);
 }
 
 int DnsSession::NextFirstServerIndex() {
@@ -41,7 +126,107 @@ base::TimeDelta DnsSession::NextTimeout(int attempt) {
   return config_.timeout * (1 << (attempt / config_.nameservers.size()));
 }
 
-DnsSession::~DnsSession() {}
+DatagramClientSocket* DnsSession::CreateConnectedSocket(
+    const IPEndPoint& endpoint) {
+  DatagramClientSocket* socket;
+
+  socket = socket_factory_->CreateDatagramClientSocket(
+      kBindType,
+      base::Bind(&base::RandInt),
+      net_log_,
+      NetLog::Source());
+  if (!socket)
+    return NULL;
+
+  int rv = socket->Connect(endpoint);
+  if (rv != OK) {
+    delete socket;
+    return NULL;
+  }
+
+  return socket;
+}
+
+DnsSession::SocketPoolImpl::SocketPoolImpl(
+    unsigned pool_size,
+    const RandIntCallback& rand_int_callback,
+    const EndPointSocketFactoryCallback& factory_callback)
+    : pool_(0),
+      pool_size_(pool_size),
+      rand_int_callback_(rand_int_callback),
+      factory_callback_(factory_callback) {
+  FillPool();
+}
+
+DnsSession::SocketPoolImpl::~SocketPoolImpl() {
+  STLDeleteElements(&pool_);
+}
+
+// static
+DnsSession::SocketPool* DnsSession::SocketPoolImpl::CreateSocketPool(
+    unsigned pool_size,
+    const RandIntCallback& rand_int_callback,
+    const EndPointSocketFactoryCallback& factory_callback) {
+  SocketPoolImpl* socket_pool = new SocketPoolImpl(pool_size,
+                                                   rand_int_callback,
+                                                   factory_callback);
+  return static_cast<SocketPool*>(socket_pool);
+}
+
+scoped_ptr<DatagramClientSocket> DnsSession::SocketPoolImpl::GetSocket() {
+  DatagramClientSocket* socket = NULL;
+
+  if (!pool_.empty()) {
+    unsigned index = rand_int_callback_.Run(0, pool_.size() - 1);
+    socket = pool_[index];
+    pool_[index] = pool_.back();
+    pool_.pop_back();
+  }
+
+  if (!socket)
+    socket = factory_callback_.Run();
+
+  return scoped_ptr<DatagramClientSocket>(socket);
+}
+
+void DnsSession::SocketPoolImpl::PutSocket(
+    scoped_ptr<DatagramClientSocket> socket) {
+  // Do nothing, so the socket is destroyed.
+}
+
+void DnsSession::SocketPoolImpl::FillPool() {
+  while (pool_.size() < pool_size_) {
+    DatagramClientSocket* socket = factory_callback_.Run();
+    if (!socket)
+      return;
+    pool_.push_back(socket);
+  }
+}
+
+// Allocate a socket, already connected to the server address.
+scoped_ptr<DnsSession::SocketLease> DnsSession::AllocateSocket(int index) {
+  CHECK(index >= 0);

[szym, 2012/09/10 19:31:00]

Consider making |index| unsigned.

[Deprecated (see juliatuttle), 2012/09/14 15:28:18]

Done.

+  CHECK(index < (int)config_.nameservers.size());
+  DCHECK(pools_.size() == config_.nameservers.size());
+
+  scoped_ptr<DatagramClientSocket> socket(pools_[index]->GetSocket());
+
+  if (socket.get())

[szym, 2012/09/10 19:31:00]

nit: needs braces.

+    return scoped_ptr<SocketLease>(new SocketLease(this, index, socket.Pass()));
+  else
+    return scoped_ptr<SocketLease>(NULL);
+}
+
+// Release a socket.
+void DnsSession::FreeSocket(int index,
+                            scoped_ptr<DatagramClientSocket> socket) {
+  CHECK(index >= 0);
+  CHECK(index < (int)config_.nameservers.size());
+  DCHECK(pools_.size() == config_.nameservers.size());
+
+  pools_[index]->PutSocket(socket.Pass());
+}
+
 
 }  // namespace net