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..0c30f76a92cb303906d1d2e4a2e4c760c54ca17d 100644
--- a/net/dns/dns_session.cc
+++ b/net/dns/dns_session.cc
@@ -6,26 +6,42 @@
 
 #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/dns/dns_socket_pool.h"
 #include "net/socket/client_socket_factory.h"
 
 namespace net {
 
+DnsSession::SocketLease::SocketLease(scoped_refptr<DnsSession> session,
+                                     unsigned server_index,
+                                     scoped_ptr<DatagramClientSocket> socket)
+    : session_(session), server_index_(server_index), socket_(socket.Pass()) {}
+
+DnsSession::SocketLease::~SocketLease() {
+  session_->FreeSocket(server_index_, socket_.Pass());
+}
+
 DnsSession::DnsSession(const DnsConfig& config,
-                       ClientSocketFactory* factory,
+                       ClientSocketFactory* socket_factory,
+                       scoped_ptr<DnsSocketPool> socket_pool,
                        const RandIntCallback& rand_int_callback,
                        NetLog* net_log)
     : config_(config),
-      socket_factory_(factory),
-      rand_callback_(base::Bind(rand_int_callback, 0, kuint16max)),
+      socket_pool_(socket_pool.Pass()),
+      rand_int_callback_(rand_int_callback),
       net_log_(net_log),
       server_index_(0) {
+  socket_pool_->Initialize(socket_factory, net_log, &config_.nameservers);
 }
 
 int DnsSession::NextQueryId() const {
-  return rand_callback_.Run();
+  /* NB: This callback is also used for selecting sockets from the pool. */
+  return rand_int_callback_.Run(0, kuint16max);
 }
 
 int DnsSession::NextFirstServerIndex() {
@@ -41,7 +57,34 @@ base::TimeDelta DnsSession::NextTimeout(int attempt) {
   return config_.timeout * (1 << (attempt / config_.nameservers.size()));
 }
 
-DnsSession::~DnsSession() {}
+// Allocate a socket, already connected to the server address.
+scoped_ptr<DnsSession::SocketLease> DnsSession::AllocateSocket(
+    unsigned server_index,
+    const NetLog::Source& source) {
+  scoped_ptr<DatagramClientSocket> socket;
 
-}  // namespace net
+  socket = socket_pool_->AllocateSocket(server_index).Pass();
+  if (!socket.get()) {
+    return scoped_ptr<SocketLease>(NULL);
+  }
+
+  socket->NetLog().BeginEvent(
+      NetLog::TYPE_SOCKET_IN_USE,
+      source.ToEventParametersCallback());
+
+  SocketLease* lease = new SocketLease(this, server_index, socket.Pass());
+  return scoped_ptr<SocketLease>(lease);
+}
+
+// Release a socket.
+void DnsSession::FreeSocket(
+    unsigned server_index,
+    scoped_ptr<DatagramClientSocket> socket) {
+  CHECK(socket.get());
+
+  socket->NetLog().EndEvent(NetLog::TYPE_SOCKET_IN_USE);
 
+  socket_pool_->FreeSocket(server_index, socket.Pass());
+}
+
+}  // namespace net