Keep pool of pre-connected DNS sockets

net/dns/dns_session.h

 // 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.
 
 #ifndef NET_DNS_DNS_SESSION_H_
 #define NET_DNS_DNS_SESSION_H_
 
+#include <vector>
+
 #include "base/memory/ref_counted.h"
+#include "base/memory/scoped_ptr.h"
 #include "base/time.h"
 #include "net/base/net_export.h"
 #include "net/base/rand_callback.h"
 #include "net/dns/dns_config_service.h"
+#include "net/udp/datagram_client_socket.h"
 
 namespace net {
 
 class ClientSocketFactory;
 class NetLog;
 
 // Session parameters and state shared between DNS transactions.
 // Ref-counted so that DnsClient::Request can keep working in absence of
 // DnsClient. A DnsSession must be recreated when DnsConfig changes.
 class NET_EXPORT_PRIVATE DnsSession
     : NON_EXPORTED_BASE(public base::RefCounted<DnsSession>) {
  public:
+  class SocketPool;
+
   typedef base::Callback<int()> RandCallback;
 
+  //typedef DatagramClientSocket* SocketFactoryFunction(const IPEndPoint&);
+  //typedef base::Callback<SocketFactoryFunction> SocketFactoryCallback;

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

What are these comments for?

+
+  typedef DatagramClientSocket* EndPointSocketFactoryFunction(void);

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

Do you really need this typedef? Is it used only in the one typedef below?

+  typedef base::Callback<EndPointSocketFactoryFunction>
+      EndPointSocketFactoryCallback;
+
+  typedef SocketPool* SocketPoolFactoryFunction(
+      const EndPointSocketFactoryCallback&);
+  typedef base::Callback<SocketPoolFactoryFunction> SocketPoolFactoryCallback;

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

A factory function is an unusual pattern. I think you should stick to ABC. This
also needs some comment explaining why SocketPoolFactoryFunction takes
EndPointSocketFactoryCallback.

+
+  class SocketLease {
+   public:
+    SocketLease(scoped_refptr<DnsSession> session,
+                int server_index,
+                scoped_ptr<DatagramClientSocket> socket);
+    ~SocketLease();
+
+    DatagramClientSocket* socket() { return socket_.get(); }
+
+   private:
+    scoped_refptr<DnsSession> session_;
+    int server_index_;
+    scoped_ptr<DatagramClientSocket> socket_;
+  };
+
+  class SocketPool {
+   public:
+    virtual scoped_ptr<DatagramClientSocket> GetSocket() = 0;
+    virtual void PutSocket(scoped_ptr<DatagramClientSocket> socket) = 0;

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

I think AllocateSocket/ReleaseSocket would be more clear.

+  };
+
+  // Use the built-in socket pool class with the default size

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

nit: End sentences with periods.

   DnsSession(const DnsConfig& config,
-             ClientSocketFactory* factory,
+             ClientSocketFactory* socket_factory,
              const RandIntCallback& rand_int_callback,
              NetLog* net_log);
 
+  // Use the built-in socket pool class, but with a custom size
+  DnsSession(const DnsConfig& config,
+             ClientSocketFactory* socket_factory,
+             const RandIntCallback& rand_int_callback,
+             int pool_size,
+             NetLog* net_log);

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

What do you need this constructor for? If only for dns_transaction_unittest.cc
then I suggest creating a MockSocketPoolFactoryCallback.

+
+  // Use a custom socket pool factory
+  DnsSession(const DnsConfig& config,
+             ClientSocketFactory* socket_factory,
+             const RandIntCallback& rand_int_callback,
+             const SocketPoolFactoryCallback& pool_factory,
+             NetLog* net_log);
+
   const DnsConfig& config() const { return config_; }
   NetLog* net_log() const { return net_log_; }
 
-  ClientSocketFactory* socket_factory() { return socket_factory_; }
-
   // Return the next random query ID.
   int NextQueryId() const;
 
   // Return the index of the first configured server to use on first attempt.
   int NextFirstServerIndex();
 
   // Return the timeout for the next query.
   base::TimeDelta NextTimeout(int attempt);
 
+  // Allocate a socket, already connected to the server address.
+  // TODO(szym): add NetLog param?
+  scoped_ptr<SocketLease> AllocateSocket(int server_index);
+
  private:
   friend class base::RefCounted<DnsSession>;
   ~DnsSession();
 
+  DatagramClientSocket* CreateConnectedSocket(const IPEndPoint& endpoint);
+
+  class SocketPoolImpl : public SocketPool {
+   public:
+    static SocketPool* CreateSocketPool(
+        unsigned pool_size,
+        const RandIntCallback& rand_int_callback,
+        const EndPointSocketFactoryCallback& factory_callback);
+    ~SocketPoolImpl();
+    virtual scoped_ptr<DatagramClientSocket> GetSocket() OVERRIDE;
+    virtual void PutSocket(scoped_ptr<DatagramClientSocket> socket) OVERRIDE;
+
+   private:
+    SocketPoolImpl(
+        unsigned pool_size,
+        const RandIntCallback& rand_int_callback,
+        const EndPointSocketFactoryCallback& factory_callback);
+    void FillPool();
+
+    std::vector<DatagramClientSocket*> pool_;
+    unsigned pool_size_;
+    RandIntCallback rand_int_callback_;
+    EndPointSocketFactoryCallback factory_callback_;
+  };
+
+  void CreatePools(const SocketPoolFactoryCallback& pool_factory_callback);
+  void CreateDefaultPools(int pool_size);
+
+  // Release a socket.
+  void FreeSocket(int server_index, scoped_ptr<DatagramClientSocket> socket);
+
   const DnsConfig config_;
   ClientSocketFactory* socket_factory_;
-  RandCallback rand_callback_;
+  RandIntCallback rand_int_callback_;
   NetLog* net_log_;
 
+  std::vector<SocketPool*> pools_;
+
   // Current index into |config_.nameservers| to begin resolution with.
   int server_index_;
 
   // TODO(szym): Add current RTT estimate.
   // TODO(szym): Add TCP connection pool to support DNS over TCP.
-  // TODO(szym): Add UDP port pool to avoid NAT table overload.
 
   DISALLOW_COPY_AND_ASSIGN(DnsSession);
 };
 
 }  // namespace net
 
 #endif  // NET_DNS_DNS_SESSION_H_