Keep pool of pre-connected DNS sockets

net/dns/dns_session.cc

 // Copyright (c) 2012 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "net/dns/dns_session.h"
 
 #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() {
   int index = server_index_;
   if (config_.rotate)
     server_index_ = (server_index_ + 1) % config_.nameservers.size();
   return index;
 }
 
 base::TimeDelta DnsSession::NextTimeout(int attempt) {
   // The timeout doubles every full round (each nameserver once).
   // TODO(szym): Adapt timeout to observed RTT. http://crbug.com/110197
   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