Cleanup unnecessary socket abstractions in DnsSession

Bug: 1116579
Change-Id: Ifd4958aa53824914335715df58bd8a8d89913ad9
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2363401
Commit-Queue: Eric Orth <ericorth@chromium.org>
Reviewed-by: Dan McArdle <dmcardle@chromium.org>
Cr-Commit-Position: refs/heads/master@{#800253}

Author: Eric Orth

net/dns/dns_session.cc

@@ -10,34 +10,15 @@
 #include <utility>
 
 #include "base/bind.h"
-#include "base/macros.h"
 #include "base/metrics/histogram_macros.h"
 #include "base/rand_util.h"
 #include "base/stl_util.h"
-#include "net/base/ip_endpoint.h"
-#include "net/base/net_errors.h"
 #include "net/dns/dns_config.h"
 #include "net/dns/dns_socket_pool.h"
-#include "net/log/net_log_event_type.h"
-#include "net/log/net_log_source.h"
-#include "net/log/net_log_with_source.h"
-#include "net/socket/datagram_client_socket.h"
-#include "net/socket/stream_socket.h"
+#include "net/log/net_log.h"
 
 namespace net {
 
-DnsSession::SocketLease::SocketLease(
-    scoped_refptr<DnsSession> session,
-    size_t server_index,
-    std::unique_ptr<DatagramClientSocket> socket)
-    : session_(session),
-      server_index_(server_index),
-      socket_(std::move(socket)) {}
-
-DnsSession::SocketLease::~SocketLease() {
-  session_->FreeSocket(server_index_, std::move(socket_));
-}
-
 DnsSession::DnsSession(const DnsConfig& config,
                        std::unique_ptr<DnsSocketPool> socket_pool,
                        const RandIntCallback& rand_int_callback,
@@ -58,39 +39,8 @@ uint16_t DnsSession::NextQueryId() const {
   return static_cast<uint16_t>(rand_callback_.Run());
 }
 
-// Allocate a socket, already connected to the server address.
-std::unique_ptr<DnsSession::SocketLease> DnsSession::AllocateSocket(
-    size_t server_index,
-    const NetLogSource& source) {
-  std::unique_ptr<DatagramClientSocket> socket;
-
-  socket = socket_pool_->CreateConnectedUdpSocket(server_index);
-  if (!socket.get())
-    return std::unique_ptr<SocketLease>();
-
-  socket->NetLog().BeginEventReferencingSource(NetLogEventType::SOCKET_IN_USE,
-                                               source);
-
-  SocketLease* lease = new SocketLease(this, server_index, std::move(socket));
-  return std::unique_ptr<SocketLease>(lease);
-}
-
-std::unique_ptr<StreamSocket> DnsSession::CreateTCPSocket(
-    size_t server_index,
-    const NetLogSource& source) {
-  return socket_pool_->CreateTcpSocket(server_index, source);
-}
-
 void DnsSession::InvalidateWeakPtrsForTesting() {
   weak_ptr_factory_.InvalidateWeakPtrs();
 }
 
-// Release a socket.
-void DnsSession::FreeSocket(size_t server_index,
-                            std::unique_ptr<DatagramClientSocket> socket) {
-  DCHECK(socket.get());
-
-  socket->NetLog().EndEvent(NetLogEventType::SOCKET_IN_USE);
-}
-
 }  // namespace net

net/dns/dns_session.h

@@ -15,15 +15,12 @@
 #include "net/base/net_export.h"
 #include "net/base/rand_callback.h"
 #include "net/dns/dns_config.h"
-#include "net/dns/dns_socket_pool.h"
 #include "net/dns/dns_udp_tracker.h"
 
 namespace net {
 
-class DatagramClientSocket;
+class DnsSocketPool;
 class NetLog;
-struct NetLogSource;
-class StreamSocket;
 
 // Session parameters and state shared between DnsTransactions for a specific
 // instance/version of a DnsConfig. Also may be used as a key handle for
@@ -38,51 +35,19 @@ class NET_EXPORT_PRIVATE DnsSession : public base::RefCounted<DnsSession> {
  public:
   typedef base::RepeatingCallback<int()> RandCallback;
 
-  // TODO(crbug.com/1116579): Remove this unnecessary abstraction now that DNS
-  // sockets are not pooled and do not need a special release signal.
-  class NET_EXPORT_PRIVATE SocketLease {
-   public:
-    SocketLease(scoped_refptr<DnsSession> session,
-                size_t server_index,
-                std::unique_ptr<DatagramClientSocket> socket);
-    ~SocketLease();
-
-    size_t server_index() const { return server_index_; }
-
-    DatagramClientSocket* socket() { return socket_.get(); }
-
-   private:
-    scoped_refptr<DnsSession> session_;
-    size_t server_index_;
-    std::unique_ptr<DatagramClientSocket> socket_;
-
-    DISALLOW_COPY_AND_ASSIGN(SocketLease);
-  };
-
   DnsSession(const DnsConfig& config,
              std::unique_ptr<DnsSocketPool> socket_pool,
              const RandIntCallback& rand_int_callback,
              NetLog* net_log);
 
   const DnsConfig& config() const { return config_; }
+  DnsSocketPool* socket_pool() { return socket_pool_.get(); }
   DnsUdpTracker* udp_tracker() { return &udp_tracker_; }
   NetLog* net_log() const { return net_log_; }
 
   // Return the next random query ID.
   uint16_t NextQueryId() const;
 
-  // Allocate a socket, already connected to the server address.
-  // When the SocketLease is destroyed, the socket will be freed.
-  // TODO(crbug.com/1116579): Remove this and directly call into DnsSocketPool.
-  std::unique_ptr<SocketLease> AllocateSocket(size_t server_index,
-                                              const NetLogSource& source);
-
-  // Creates a StreamSocket from the factory for a transaction over TCP. These
-  // sockets are not pooled.
-  // TODO(crbug.com/1116579): Remove this and directly call into DnsSocketPool.
-  std::unique_ptr<StreamSocket> CreateTCPSocket(size_t server_index,
-                                                const NetLogSource& source);
-
   base::WeakPtr<DnsSession> GetWeakPtr() {
     return weak_ptr_factory_.GetWeakPtr();
   }
@@ -98,10 +63,6 @@ class NET_EXPORT_PRIVATE DnsSession : public base::RefCounted<DnsSession> {
 
   ~DnsSession();
 
-  // Release a socket.
-  void FreeSocket(size_t server_index,
-                  std::unique_ptr<DatagramClientSocket> socket);
-
   const DnsConfig config_;
   std::unique_ptr<DnsSocketPool> socket_pool_;
   DnsUdpTracker udp_tracker_;

net/dns/dns_test_util.cc

@@ -18,6 +18,7 @@
 #include "net/dns/dns_hosts.h"
 #include "net/dns/dns_query.h"
 #include "net/dns/dns_session.h"
+#include "net/dns/dns_socket_pool.h"
 #include "net/dns/dns_util.h"
 #include "net/dns/resolve_context.h"
 #include "testing/gtest/include/gtest/gtest.h"

net/dns/dns_transaction.cc

@@ -48,6 +48,7 @@
 #include "net/dns/dns_response.h"
 #include "net/dns/dns_server_iterator.h"
 #include "net/dns/dns_session.h"
+#include "net/dns/dns_socket_pool.h"
 #include "net/dns/dns_udp_tracker.h"
 #include "net/dns/dns_util.h"
 #include "net/dns/public/dns_over_https_server_config.h"
@@ -187,12 +188,12 @@ class DnsAttempt {
 class DnsUDPAttempt : public DnsAttempt {
  public:
   DnsUDPAttempt(size_t server_index,
-                std::unique_ptr<DnsSession::SocketLease> socket_lease,
+                std::unique_ptr<DatagramClientSocket> socket,
                 std::unique_ptr<DnsQuery> query,
                 DnsUdpTracker* udp_tracker)
       : DnsAttempt(server_index),
         next_state_(STATE_NONE),
-        socket_lease_(std::move(socket_lease)),
+        socket_(std::move(socket)),
         query_(std::move(query)),
         udp_tracker_(udp_tracker) {}
 
@@ -205,7 +206,7 @@ class DnsUDPAttempt : public DnsAttempt {
     next_state_ = STATE_SEND_QUERY;
 
     IPEndPoint local_address;
-    if (socket_lease_->socket()->GetLocalAddress(&local_address) == OK)
+    if (socket_->GetLocalAddress(&local_address) == OK)
       udp_tracker_->RecordQuery(local_address.port(), query_->id());
 
     return DoLoop(OK);
@@ -219,7 +220,7 @@ class DnsUDPAttempt : public DnsAttempt {
   }
 
   const NetLogWithSource& GetSocketNetLog() const override {
-    return socket_lease_->socket()->NetLog();
+    return socket_->NetLog();
   }
 
  private:
@@ -231,8 +232,6 @@ class DnsUDPAttempt : public DnsAttempt {
     STATE_NONE,
   };
 
-  DatagramClientSocket* socket() { return socket_lease_->socket(); }
-
   int DoLoop(int result) {
     CHECK_NE(STATE_NONE, next_state_);
     int rv = result;
@@ -270,7 +269,7 @@ class DnsUDPAttempt : public DnsAttempt {
 
   int DoSendQuery() {
     next_state_ = STATE_SEND_QUERY_COMPLETE;
-    return socket()->Write(
+    return socket_->Write(
         query_->io_buffer(), query_->io_buffer()->size(),
         base::BindOnce(&DnsUDPAttempt::OnIOComplete, base::Unretained(this)),
         kTrafficAnnotation);
@@ -292,7 +291,7 @@ class DnsUDPAttempt : public DnsAttempt {
   int DoReadResponse() {
     next_state_ = STATE_READ_RESPONSE_COMPLETE;
     response_ = std::make_unique<DnsResponse>();
-    return socket()->Read(
+    return socket_->Read(
         response_->io_buffer(), response_->io_buffer_size(),
         base::BindOnce(&DnsUDPAttempt::OnIOComplete, base::Unretained(this)));
   }
@@ -328,7 +327,7 @@ class DnsUDPAttempt : public DnsAttempt {
   State next_state_;
   base::TimeTicks start_time_;
 
-  std::unique_ptr<DnsSession::SocketLease> socket_lease_;
+  std::unique_ptr<DatagramClientSocket> socket_;
   std::unique_ptr<DnsQuery> query_;
 
   // Should be owned by the DnsSession, to which the transaction should own a
@@ -1258,13 +1257,13 @@ class DnsTransactionImpl : public DnsTransaction,
     DCHECK(!secure_);
     size_t attempt_number = attempts_.size();
 
-    std::unique_ptr<DnsSession::SocketLease> lease =
-        session_->AllocateSocket(server_index, net_log_.source());
+    std::unique_ptr<DatagramClientSocket> socket =
+        session_->socket_pool()->CreateConnectedUdpSocket(server_index);
 
-    bool got_socket = !!lease.get();
+    bool got_socket = !!socket.get();
 
     DnsUDPAttempt* attempt =
-        new DnsUDPAttempt(server_index, std::move(lease), std::move(query),
+        new DnsUDPAttempt(server_index, std::move(socket), std::move(query),
                           session_->udp_tracker());
 
     attempts_.push_back(base::WrapUnique(attempt));
@@ -1339,7 +1338,8 @@ class DnsTransactionImpl : public DnsTransaction,
     DCHECK(!secure_);
 
     std::unique_ptr<StreamSocket> socket(
-        session_->CreateTCPSocket(server_index, net_log_.source()));
+        session_->socket_pool()->CreateTcpSocket(server_index,
+                                                 net_log_.source()));
 
     unsigned attempt_number = attempts_.size();
 

net/dns/dns_transaction_unittest.cc

@@ -35,6 +35,7 @@
 #include "net/dns/dns_response.h"
 #include "net/dns/dns_server_iterator.h"
 #include "net/dns/dns_session.h"
+#include "net/dns/dns_socket_pool.h"
 #include "net/dns/dns_test_util.h"
 #include "net/dns/dns_util.h"
 #include "net/dns/public/dns_over_https_server_config.h"