Keep pool of pre-connected DNS sockets

net/dns/dns_socket_pool.cc

Index: net/dns/dns_socket_pool.cc
diff --git a/net/dns/dns_socket_pool.cc b/net/dns/dns_socket_pool.cc
new file mode 100644
index 0000000000000000000000000000000000000000..e57262bfd53977678f4273a55432db932c830ce9
--- /dev/null
+++ b/net/dns/dns_socket_pool.cc
@@ -0,0 +1,161 @@
+// 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_socket_pool.h"
+
+#include "base/logging.h"
+#include "base/rand_util.h"
+#include "net/base/ip_endpoint.h"
+#include "net/base/net_errors.h"
+#include "net/base/rand_callback.h"
+#include "net/socket/client_socket_factory.h"
+#include "net/udp/datagram_client_socket.h"
+
+namespace net {
+
+DnsSocketPool::DnsSocketPool()
+    : socket_factory_(NULL),
+      net_log_(NULL),
+      nameservers_(NULL),
+      initialized_(false) {
+}
+
+void DnsSocketPool::Initialize(
+    ClientSocketFactory* socket_factory,
+    NetLog* net_log,
+    const std::vector<IPEndPoint>* nameservers) {
+  // net_log can be NULL if we don't want to log
+  CHECK(socket_factory && nameservers);

[szym, 2012/09/14 20:49:27]

use separate CHECKs, should probably use DCHECKs instead.

[Deprecated (see juliatuttle), 2012/09/18 20:28:00]

Done.

+
+  CHECK(!initialized_);
+
+  socket_factory_ = socket_factory;
+  net_log_ = net_log;
+  nameservers_ = nameservers;
+  initialized_ = true;
+}
+
+// On Windows, we can't request specific (random) ports, since that will
+// trigger firewall prompts, so request default ones, but keep a pile of
+// them.  Everywhere else, request fresh, random ports each time.
+#if defined(OS_WIN)
+static const DatagramSocket::BindType kBindType = DatagramSocket::DEFAULT_BIND;

[szym, 2012/09/14 20:49:27]

suggest anonymous namespace at the top of this file instead of static

[Deprecated (see juliatuttle), 2012/09/18 20:28:00]

Done.

+static const unsigned kInitialPoolSize = 256;
+static const unsigned kAllocateMinSize = 256;
+static const unsigned kRetainMaxSize = 0;
+#else
+static const DatagramSocket::BindType kBindType = DatagramSocket::RANDOM_BIND;
+static const unsigned kInitialPoolSize = 0;
+static const unsigned kAllocateMinSize = 0;
+static const unsigned kRetainMaxSize = 0;
+#endif
+
+scoped_ptr<DatagramClientSocket> DnsSocketPool::CreateConnectedSocket(
+    unsigned server_index) {
+  CHECK(server_index < nameservers_->size());

[szym, 2012/09/14 20:49:27]

Should probably use DCHECK. Use DCHECK_LT

[Deprecated (see juliatuttle), 2012/09/18 20:28:00]

Done.

+
+  DatagramClientSocket* socket;

[szym, 2012/09/14 20:49:27]

Use scoped_ptr here so that you don't need to worry about wrapping it or
cleaning it up.

[Deprecated (see juliatuttle), 2012/09/18 20:28:00]

Done.

+
+  NetLog::Source no_source;
+  socket = socket_factory_->CreateDatagramClientSocket(
+      kBindType, base::Bind(&base::RandInt), net_log_, no_source);
+
+  if (!socket) {
+    LOG(WARNING) << "Failed to create socket.";
+     return scoped_ptr<DatagramClientSocket>(NULL);

[szym, 2012/09/14 20:49:27]

nit: indent

[Deprecated (see juliatuttle), 2012/09/18 20:28:00]

Done.

+  }
+
+  int rv = socket->Connect((*nameservers_)[server_index]);
+  if (rv != OK) {
+    LOG(WARNING) << "Failed to connect socket: " << rv;
+    delete socket;
+    return scoped_ptr<DatagramClientSocket>(NULL);
+  }
+
+  return scoped_ptr<DatagramClientSocket>(socket);
+}
+
+class DefaultDnsSocketPool : public DnsSocketPool {
+ public:
+  virtual void Initialize(
+      ClientSocketFactory* socket_factory,
+      NetLog* net_log,
+      const std::vector<IPEndPoint>* nameservers);

[szym, 2012/09/14 20:49:27]

All three need OVERRIDE.

[Deprecated (see juliatuttle), 2012/09/18 20:28:00]

Done.

+  virtual scoped_ptr<DatagramClientSocket> AllocateSocket(
+      unsigned server_index);
+  virtual void FreeSocket(
+      unsigned server_index,
+      scoped_ptr<DatagramClientSocket> socket);
+ private:

[szym, 2012/09/14 20:49:27]

Add empty line above.

[Deprecated (see juliatuttle), 2012/09/18 20:28:00]

Done.

+  std::vector<std::vector<DatagramClientSocket*> > pools_;

[szym, 2012/09/14 20:49:27]

DISALLOW_COPY_AND_ASSIGN

[szym, 2012/09/14 20:49:27]

All of the sockets leak out in ~DefaultDnsSocketPool.

I wonder if it wouldn't be better to use
scoped_array<ScopedVector<DatagramClientSocket> >
You are already doing bound checking yourself.

[Deprecated (see juliatuttle), 2012/09/18 20:28:00]

Done.

[szym, 2012/09/26 11:16:12]

Missed the one about leak.

[Deprecated (see juliatuttle), 2012/09/26 20:44:50]

Done.

+};
+
+// static
+scoped_ptr<DnsSocketPool> DnsSocketPool::MakeDefault() {
+  DefaultDnsSocketPool* default_pool = new DefaultDnsSocketPool();
+  DnsSocketPool* pool = static_cast<DnsSocketPool*>(default_pool);

[szym, 2012/09/14 20:49:27]

no need for static_cast to base class pointer

[Deprecated (see juliatuttle), 2012/09/18 20:28:00]

Done.

+  return scoped_ptr<DnsSocketPool>(pool);

[szym, 2012/09/14 20:49:27]

scoped_ptr<DnsSocketPool>(new DefaultDnsSocketPool()) should just work

[Deprecated (see juliatuttle), 2012/09/18 20:28:00]

Done.

+}
+
+void DefaultDnsSocketPool::Initialize(
+    ClientSocketFactory* socket_factory,
+    NetLog* net_log,
+    const std::vector<IPEndPoint>* nameservers) OVERRIDE {

[szym, 2012/09/14 20:49:27]

I suggest you DCHECK that |pools_| is empty to enforce you can Initialize only
once.

[Deprecated (see juliatuttle), 2012/09/18 20:28:00]

InitializeInternal already keeps a flag and enforces that we can only initialize
once.

+  DnsSocketPool::Initialize(socket_factory, net_log, nameservers);
+
+  const unsigned num_servers = nameservers->size();
+  pools_.resize(num_servers);
+  for (unsigned server_index = 0; server_index < num_servers; ++server_index) {
+    std::vector<DatagramClientSocket*>& pool = pools_[server_index];
+    for (unsigned pool_index = 0; pool_index < kInitialPoolSize; ++pool_index) {
+      DatagramClientSocket* socket =
+          CreateConnectedSocket(server_index).release();
+      if (!socket) {

[szym, 2012/09/14 20:49:27]

nit: no need for {}

[Deprecated (see juliatuttle), 2012/09/18 20:28:00]

Done.

+        break;
+      }
+      pool.push_back(socket);
+    }
+  }
+}
+
+scoped_ptr<DatagramClientSocket> DefaultDnsSocketPool::AllocateSocket(
+    unsigned server_index) {
+  CHECK(server_index < pools_.size());

[szym, 2012/09/14 20:49:27]

DCHECK_LT

[Deprecated (see juliatuttle), 2012/09/18 20:28:00]

Done.

+  std::vector<DatagramClientSocket*>& pool = pools_[server_index];
+
+  // allocate one extra socket, since we're about to remove one

[szym, 2012/09/14 20:49:27]

You say "one extra" but this is a loop.

[Deprecated (see juliatuttle), 2012/09/18 20:28:00]

The comment is to explain "kAllocateMinSize + 1".

+  while (pool.size() < kAllocateMinSize + 1) {
+    DatagramClientSocket* socket =
+        CreateConnectedSocket(server_index).release();
+    if (!socket) {

[szym, 2012/09/14 20:49:27]

nit: no need for {}

[Deprecated (see juliatuttle), 2012/09/26 20:44:50]

Done.

+      break;
+    }
+    pool.push_back(socket);
+  }
+
+  if (pool.size() < kAllocateMinSize) {
+    LOG(WARNING) << "Low DNS port entropy: " << "wanted " << kAllocateMinSize

[szym, 2012/09/14 20:49:27]

nit: no need for the second <<

[Deprecated (see juliatuttle), 2012/09/18 20:28:00]

Done.

+                 << " sockets to choose from, but only have " << pool.size();
+  }
+
+  unsigned socket_index = base::RandInt(0, pool.size() - 1);
+  DatagramClientSocket* socket = pool[socket_index];
+  pool[socket_index] = pool.back();
+  pool.pop_back();
+
+  return scoped_ptr<DatagramClientSocket>(socket);
+}
+
+void DefaultDnsSocketPool::FreeSocket(
+    unsigned server_index,
+    scoped_ptr<DatagramClientSocket> socket) {
+  CHECK(server_index < pools_.size());

[szym, 2012/09/14 20:49:27]

DCHECK_LT

[Deprecated (see juliatuttle), 2012/09/18 20:28:00]

Done.

+  std::vector<DatagramClientSocket*>& pool = pools_[server_index];
+
+  if (pool.size() < kRetainMaxSize) {

[szym, 2012/09/14 20:49:27]

nit: no need for {}

[Deprecated (see juliatuttle), 2012/09/18 20:28:00]

Done.

+    pool.push_back(socket.release());
+  }
+}
+
+} // namespace net