Refactoring: Create per-connection packet writers in QuicDispatcher.

To make porting the QUIC EndToEndTest to Chromium possible with fewer
Chromium-specific parts in shared code, I've made QuicDispatcher expose
and accept a QuicDispatcher::PacketWriterFactory which it uses to create
a new packet writer wrapper for every QuicConnection. I also changed
QuicConnection to accept a QuicConnection::PacketWriterFactory (a second
new type of factory) rather than the writer itself in its constructor,
since the per-connection packet writers need to be created with the
connection already existing.

Merge internal CL: 73064412

Written by Daniel Ziegler <dmziegler@chromium.org>
Original review URL: https://codereview.chromium.org/467963002/

R=rch@chromium.org,wtc@chromium.org
BUG=

Review URL: https://codereview.chromium.org/475113005

Cr-Commit-Position: refs/heads/master@{#291314}
git-svn-id: svn://svn.chromium.org/chrome/trunk/src@291314 0039d316-1c4b-4281-b951-d872f2087c98

17 files changed, 362 insertions, 42 deletions

diff --git a/net/tools/quic/end_to_end_test.cc b/net/tools/quic/end_to_end_test.cc
index 24ad3905..1f4103f 100644
--- a/net/tools/quic/end_to_end_test.cc
+++ b/net/tools/quic/end_to_end_test.cc
@@ -161,11 +161,17 @@ vector<TestParams> GetTestParams() {
 
 class ServerDelegate : public PacketDroppingTestWriter::Delegate {
  public:
-  explicit ServerDelegate(QuicDispatcher* dispatcher)
-      : dispatcher_(dispatcher) {}
+  ServerDelegate(TestWriterFactory* writer_factory,
+                 QuicDispatcher* dispatcher)
+      : writer_factory_(writer_factory),
+        dispatcher_(dispatcher) {}
   virtual ~ServerDelegate() {}
+  virtual void OnPacketSent(WriteResult result) override {
+    writer_factory_->OnPacketSent(result);
+  }
   virtual void OnCanWrite() OVERRIDE { dispatcher_->OnCanWrite(); }
  private:
+  TestWriterFactory* writer_factory_;
   QuicDispatcher* dispatcher_;
 };
 
@@ -173,6 +179,7 @@ class ClientDelegate : public PacketDroppingTestWriter::Delegate {
  public:
   explicit ClientDelegate(QuicClient* client) : client_(client) {}
   virtual ~ClientDelegate() {}
+  virtual void OnPacketSent(WriteResult result) OVERRIDE {}
   virtual void OnCanWrite() OVERRIDE {
     EpollEvent event(EPOLLOUT, false);
     client_->OnEvent(client_->fd(), &event);
@@ -326,7 +333,7 @@ class EndToEndTest : public ::testing::TestWithParam<TestParams> {
 
   virtual void SetUp() OVERRIDE {
     // The ownership of these gets transferred to the QuicPacketWriterWrapper
-    // and QuicDispatcher when Initialize() is executed.
+    // and TestWriterFactory when Initialize() is executed.
     client_writer_ = new PacketDroppingTestWriter();
     server_writer_ = new PacketDroppingTestWriter();
   }
@@ -346,10 +353,13 @@ class EndToEndTest : public ::testing::TestWithParam<TestParams> {
                                  server_thread_->GetPort());
     QuicDispatcher* dispatcher =
         QuicServerPeer::GetDispatcher(server_thread_->server());
+    TestWriterFactory* packet_writer_factory = new TestWriterFactory();
+    QuicDispatcherPeer::SetPacketWriterFactory(dispatcher,
+                                               packet_writer_factory);
     QuicDispatcherPeer::UseWriter(dispatcher, server_writer_);
     server_writer_->Initialize(
         QuicDispatcherPeer::GetHelper(dispatcher),
-        new ServerDelegate(dispatcher));
+        new ServerDelegate(packet_writer_factory, dispatcher));
     server_thread_->Start();
     server_started_ = true;
   }
@@ -1150,7 +1160,7 @@ TEST_P(EndToEndTest, ConnectionMigrationClientIPChanged) {
   writer->set_writer(new QuicDefaultPacketWriter(client_->client()->fd()));
   QuicConnectionPeer::SetWriter(client_->client()->session()->connection(),
                                 writer,
-                                true  /* owns_writer */);
+                                /* owns_writer= */ true);
 
   client_->SendSynchronousRequest("/bar");
 
diff --git a/net/tools/quic/quic_client.cc b/net/tools/quic/quic_client.cc
index ca87c2f..48f789f 100644
--- a/net/tools/quic/quic_client.cc
+++ b/net/tools/quic/quic_client.cc
@@ -97,6 +97,18 @@ bool QuicClient::Initialize() {
   return true;
 }
 
+QuicClient::DummyPacketWriterFactory::DummyPacketWriterFactory(
+    QuicPacketWriter* writer)
+    : writer_(writer) {}
+
+QuicClient::DummyPacketWriterFactory::~DummyPacketWriterFactory() {}
+
+QuicPacketWriter* QuicClient::DummyPacketWriterFactory::Create(
+    QuicConnection* /*connection*/) const {
+  return writer_;
+}
+
+
 bool QuicClient::CreateUDPSocket() {
   int address_family = server_address_.GetSockAddrFamily();
   fd_ = socket(address_family, SOCK_DGRAM | SOCK_NONBLOCK, IPPROTO_UDP);
@@ -179,15 +191,18 @@ bool QuicClient::StartConnect() {
 
   QuicPacketWriter* writer = CreateQuicPacketWriter();
 
+  DummyPacketWriterFactory factory(writer);
+
   session_.reset(new QuicClientSession(
       config_,
       new QuicConnection(GenerateConnectionId(),
                          server_address_,
                          helper_.get(),
-                         writer,
-                         false  /* owns_writer */,
-                         false  /* is_server */,
+                         factory,
+                         /* owns_writer= */ false,
+                         /* is_server= */ false,
                          supported_versions_)));
+
   // Reset |writer_| after |session_| so that the old writer outlives the old
   // session.
   if (writer_.get() != writer) {
diff --git a/net/tools/quic/quic_client.h b/net/tools/quic/quic_client.h
index a4fb6c3..8b0f1fe 100644
--- a/net/tools/quic/quic_client.h
+++ b/net/tools/quic/quic_client.h
@@ -188,6 +188,18 @@ class QuicClient : public EpollCallbackInterface,
  private:
   friend class net::tools::test::QuicClientPeer;
 
+  // A packet writer factory that always returns the same writer
+  class DummyPacketWriterFactory : public QuicConnection::PacketWriterFactory {
+   public:
+    DummyPacketWriterFactory(QuicPacketWriter* writer);
+    virtual ~DummyPacketWriterFactory();
+
+    virtual QuicPacketWriter* Create(QuicConnection* connection) const OVERRIDE;
+
+   private:
+    QuicPacketWriter* writer_;
+  };
+
   // Used during initialization: creates the UDP socket FD, sets socket options,
   // and binds the socket to our address.
   bool CreateUDPSocket();
diff --git a/net/tools/quic/quic_dispatcher.cc b/net/tools/quic/quic_dispatcher.cc
index cf3a46b..05cb39d 100644
--- a/net/tools/quic/quic_dispatcher.cc
+++ b/net/tools/quic/quic_dispatcher.cc
@@ -15,6 +15,7 @@
 #include "net/tools/epoll_server/epoll_server.h"
 #include "net/tools/quic/quic_default_packet_writer.h"
 #include "net/tools/quic/quic_epoll_connection_helper.h"
+#include "net/tools/quic/quic_per_connection_packet_writer.h"
 #include "net/tools/quic/quic_socket_utils.h"
 #include "net/tools/quic/quic_time_wait_list_manager.h"
 
@@ -159,15 +160,37 @@ class QuicDispatcher::QuicFramerVisitor : public QuicFramerVisitorInterface {
   QuicConnectionId connection_id_;
 };
 
+QuicPacketWriter* QuicDispatcher::DefaultPacketWriterFactory::Create(
+    QuicPacketWriter* writer,
+    QuicConnection* connection) {
+  return new QuicPerConnectionPacketWriter(writer, connection);
+}
+
+QuicDispatcher::PacketWriterFactoryAdapter::PacketWriterFactoryAdapter(
+    QuicDispatcher* dispatcher)
+    : dispatcher_(dispatcher) {}
+
+QuicDispatcher::PacketWriterFactoryAdapter::~PacketWriterFactoryAdapter() {}
+
+QuicPacketWriter* QuicDispatcher::PacketWriterFactoryAdapter::Create(
+    QuicConnection* connection) const {
+  return dispatcher_->packet_writer_factory_->Create(
+      dispatcher_->writer_.get(),
+      connection);
+}
+
 QuicDispatcher::QuicDispatcher(const QuicConfig& config,
                                const QuicCryptoServerConfig& crypto_config,
                                const QuicVersionVector& supported_versions,
+                               PacketWriterFactory* packet_writer_factory,
                                EpollServer* epoll_server)
     : config_(config),
       crypto_config_(crypto_config),
       delete_sessions_alarm_(new DeleteSessionsAlarm(this)),
       epoll_server_(epoll_server),
       helper_(new QuicEpollConnectionHelper(epoll_server_)),
+      packet_writer_factory_(packet_writer_factory),
+      connection_writer_factory_(this),
       supported_versions_(supported_versions),
       current_packet_(NULL),
       framer_(supported_versions, /*unused*/ QuicTime::Zero(), true),
@@ -366,9 +389,9 @@ QuicConnection* QuicDispatcher::CreateQuicConnection(
   return new QuicConnection(connection_id,
                             client_address,
                             helper_.get(),
-                            writer_.get(),
-                            false  /* owns_writer */,
-                            true   /* is_server */,
+                            connection_writer_factory_,
+                            /* owns_writer= */ true,
+                            /* is_server= */ true,
                             supported_versions_);
 }
 
diff --git a/net/tools/quic/quic_dispatcher.h b/net/tools/quic/quic_dispatcher.h
index 8d7b03b..ba6a7ad 100644
--- a/net/tools/quic/quic_dispatcher.h
+++ b/net/tools/quic/quic_dispatcher.h
@@ -20,18 +20,6 @@
 #include "net/tools/quic/quic_server_session.h"
 #include "net/tools/quic/quic_time_wait_list_manager.h"
 
-#if defined(COMPILER_GCC)
-namespace BASE_HASH_NAMESPACE {
-template<>
-struct hash<net::QuicBlockedWriterInterface*> {
-  std::size_t operator()(
-      const net::QuicBlockedWriterInterface* ptr) const {
-    return hash<size_t>()(reinterpret_cast<size_t>(ptr));
-  }
-};
-}
-#endif
-
 namespace net {
 
 class EpollServer;
@@ -61,15 +49,40 @@ class ProcessPacketInterface {
 class QuicDispatcher : public QuicServerSessionVisitor,
                        public ProcessPacketInterface {
  public:
+  // Creates per-connection packet writers out of the QuicDispatcher's shared
+  // QuicPacketWriter. The per-connection writers' IsWriteBlocked() state must
+  // always be the same as the shared writer's IsWriteBlocked(), or else the
+  // QuicDispatcher::OnCanWrite logic will not work. (This will hopefully be
+  // cleaned up for bug 16950226.)
+  class PacketWriterFactory {
+   public:
+    virtual ~PacketWriterFactory() {}
+
+    virtual QuicPacketWriter* Create(QuicPacketWriter* writer,
+                                     QuicConnection* connection) = 0;
+  };
+
+  // Creates ordinary QuicPerConnectionPacketWriter instances.
+  class DefaultPacketWriterFactory : public PacketWriterFactory {
+   public:
+    virtual ~DefaultPacketWriterFactory() {}
+
+    virtual QuicPacketWriter* Create(
+        QuicPacketWriter* writer,
+        QuicConnection* connection) OVERRIDE;
+  };
+
   // Ideally we'd have a linked_hash_set: the  boolean is unused.
   typedef linked_hash_map<QuicBlockedWriterInterface*, bool> WriteBlockedList;
 
-  // Due to the way delete_sessions_closure_ is registered, the Dispatcher
-  // must live until epoll_server Shutdown. |supported_versions| specifies the
-  // list of supported QUIC versions.
+  // Due to the way delete_sessions_closure_ is registered, the Dispatcher must
+  // live until epoll_server Shutdown. |supported_versions| specifies the list
+  // of supported QUIC versions. Takes ownership of |packet_writer_factory|,
+  // which is used to create per-connection writers.
   QuicDispatcher(const QuicConfig& config,
                  const QuicCryptoServerConfig& crypto_config,
                  const QuicVersionVector& supported_versions,
+                 PacketWriterFactory* packet_writer_factory,
                  EpollServer* epoll_server);
 
   virtual ~QuicDispatcher();
@@ -164,10 +177,29 @@ class QuicDispatcher : public QuicServerSessionVisitor,
 
   QuicPacketWriter* writer() { return writer_.get(); }
 
+  const QuicConnection::PacketWriterFactory& connection_writer_factory() {
+    return connection_writer_factory_;
+  }
+
  private:
   class QuicFramerVisitor;
   friend class net::tools::test::QuicDispatcherPeer;
 
+  // An adapter that creates packet writers using the dispatcher's
+  // PacketWriterFactory and shared writer. Essentially, it just curries the
+  // writer argument away from QuicDispatcher::PacketWriterFactory.
+  class PacketWriterFactoryAdapter :
+    public QuicConnection::PacketWriterFactory {
+   public:
+    PacketWriterFactoryAdapter(QuicDispatcher* dispatcher);
+    virtual ~PacketWriterFactoryAdapter ();
+
+    virtual QuicPacketWriter* Create(QuicConnection* connection) const OVERRIDE;
+
+   private:
+    QuicDispatcher* dispatcher_;
+  };
+
   // Called by |framer_visitor_| when the private header has been parsed
   // of a data packet that is destined for the time wait manager.
   void OnUnauthenticatedHeader(const QuicPacketHeader& header);
@@ -204,6 +236,12 @@ class QuicDispatcher : public QuicServerSessionVisitor,
   // The writer to write to the socket with.
   scoped_ptr<QuicPacketWriter> writer_;
 
+  // Used to create per-connection packet writers, not |writer_| itself.
+  scoped_ptr<PacketWriterFactory> packet_writer_factory_;
+
+  // Passed in to QuicConnection for it to create the per-connection writers
+  PacketWriterFactoryAdapter connection_writer_factory_;
+
   // This vector contains QUIC versions which we currently support.
   // This should be ordered such that the highest supported version is the first
   // element, with subsequent elements in descending order (versions can be
diff --git a/net/tools/quic/quic_dispatcher_test.cc b/net/tools/quic/quic_dispatcher_test.cc
index 4c778cb..33685c9 100644
--- a/net/tools/quic/quic_dispatcher_test.cc
+++ b/net/tools/quic/quic_dispatcher_test.cc
@@ -48,6 +48,7 @@ class TestDispatcher : public QuicDispatcher {
       : QuicDispatcher(config,
                        crypto_config,
                        QuicSupportedVersions(),
+                       new QuicDispatcher::DefaultPacketWriterFactory(),
                        eps) {
   }
 
@@ -271,12 +272,11 @@ class BlockingWriter : public QuicPacketWriterWrapper {
       size_t buf_len,
       const IPAddressNumber& self_client_address,
       const IPEndPoint& peer_client_address) OVERRIDE {
-    if (write_blocked_) {
-      return WriteResult(WRITE_STATUS_BLOCKED, EAGAIN);
-    } else {
-      return QuicPacketWriterWrapper::WritePacket(
-          buffer, buf_len, self_client_address, peer_client_address);
-    }
+    // It would be quite possible to actually implement this method here with
+    // the fake blocked status, but it would be significantly more work in
+    // Chromium, and since it's not called anyway, don't bother.
+    LOG(DFATAL) << "Not supported";
+    return WriteResult();
   }
 
   bool write_blocked_;
@@ -286,6 +286,8 @@ class QuicDispatcherWriteBlockedListTest : public QuicDispatcherTest {
  public:
   virtual void SetUp() {
     writer_ = new BlockingWriter;
+    QuicDispatcherPeer::SetPacketWriterFactory(&dispatcher_,
+                                               new TestWriterFactory());
     QuicDispatcherPeer::UseWriter(&dispatcher_, writer_);
 
     IPEndPoint client_address(net::test::Loopback4(), 1);
diff --git a/net/tools/quic/quic_per_connection_packet_writer.cc b/net/tools/quic/quic_per_connection_packet_writer.cc
new file mode 100644
index 0000000..508946b
--- /dev/null
+++ b/net/tools/quic/quic_per_connection_packet_writer.cc
@@ -0,0 +1,46 @@
+// Copyright 2014 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/tools/quic/quic_per_connection_packet_writer.h"
+
+namespace net {
+
+namespace tools {
+
+QuicPerConnectionPacketWriter::QuicPerConnectionPacketWriter(
+    QuicPacketWriter* shared_writer,
+    QuicConnection* connection)
+    : shared_writer_(shared_writer),
+      connection_(connection) {
+}
+
+QuicPerConnectionPacketWriter::~QuicPerConnectionPacketWriter() {
+}
+
+WriteResult QuicPerConnectionPacketWriter::WritePacket(
+    const char* buffer,
+    size_t buf_len,
+    const IPAddressNumber& self_address,
+    const IPEndPoint& peer_address) {
+  return shared_writer_->WritePacket(buffer,
+                                     buf_len,
+                                     self_address,
+                                     peer_address);
+}
+
+bool QuicPerConnectionPacketWriter::IsWriteBlockedDataBuffered() const {
+  return shared_writer_->IsWriteBlockedDataBuffered();
+}
+
+bool QuicPerConnectionPacketWriter::IsWriteBlocked() const {
+  return shared_writer_->IsWriteBlocked();
+}
+
+void QuicPerConnectionPacketWriter::SetWritable() {
+  shared_writer_->SetWritable();
+}
+
+}  // namespace tools
+
+}  // namespace net
diff --git a/net/tools/quic/quic_per_connection_packet_writer.h b/net/tools/quic/quic_per_connection_packet_writer.h
new file mode 100644
index 0000000..a442a9a
--- /dev/null
+++ b/net/tools/quic/quic_per_connection_packet_writer.h
@@ -0,0 +1,48 @@
+// Copyright 2014 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_TOOLS_QUIC_QUIC_PER_CONNECTION_PACKET_WRITER_H_
+#define NET_TOOLS_QUIC_QUIC_PER_CONNECTION_PACKET_WRITER_H_
+
+#include "net/quic/quic_connection.h"
+#include "net/quic/quic_packet_writer.h"
+
+namespace net {
+
+namespace tools {
+
+// A connection-specific packet writer that wraps a shared writer and keeps a
+// reference to the connection.
+class QuicPerConnectionPacketWriter : public QuicPacketWriter {
+ public:
+  // Does not take ownership of |shared_writer| or |connection|.
+  QuicPerConnectionPacketWriter(QuicPacketWriter* shared_writer,
+                                QuicConnection* connection);
+  virtual ~QuicPerConnectionPacketWriter();
+
+  QuicPacketWriter* shared_writer() const { return shared_writer_; }
+  QuicConnection* connection() const { return connection_; }
+
+  // Default implementation of the QuicPacketWriter interface: Passes everything
+  // to |shared_writer_|.
+  virtual WriteResult WritePacket(const char* buffer,
+                                  size_t buf_len,
+                                  const IPAddressNumber& self_address,
+                                  const IPEndPoint& peer_address) OVERRIDE;
+  virtual bool IsWriteBlockedDataBuffered() const OVERRIDE;
+  virtual bool IsWriteBlocked() const OVERRIDE;
+  virtual void SetWritable() OVERRIDE;
+
+ private:
+  QuicPacketWriter* shared_writer_;  // Not owned.
+  QuicConnection* connection_;  // Not owned.
+
+  DISALLOW_COPY_AND_ASSIGN(QuicPerConnectionPacketWriter);
+};
+
+}  // namespace tools
+
+}  // namespace net
+
+#endif  // NET_TOOLS_QUIC_QUIC_PER_CONNECTION_PACKET_WRITER_H_
diff --git a/net/tools/quic/quic_server.cc b/net/tools/quic/quic_server.cc
index 651f900..fa93678 100644
--- a/net/tools/quic/quic_server.cc
+++ b/net/tools/quic/quic_server.cc
@@ -166,6 +166,7 @@ QuicDispatcher* QuicServer::CreateQuicDispatcher() {
       config_,
       crypto_config_,
       supported_versions_,
+      new QuicDispatcher::DefaultPacketWriterFactory(),
       &epoll_server_);
 }
 
diff --git a/net/tools/quic/quic_server_test.cc b/net/tools/quic/quic_server_test.cc
index e2d7c4f..1107ded 100644
--- a/net/tools/quic/quic_server_test.cc
+++ b/net/tools/quic/quic_server_test.cc
@@ -21,7 +21,10 @@ class QuicServerDispatchPacketTest : public ::testing::Test {
  public:
   QuicServerDispatchPacketTest()
       : crypto_config_("blah", QuicRandom::GetInstance()),
-        dispatcher_(config_, crypto_config_, &eps_) {
+        dispatcher_(config_,
+                    crypto_config_,
+                    new QuicDispatcher::DefaultPacketWriterFactory(),
+                    &eps_) {
     dispatcher_.Initialize(1234);
   }
 
diff --git a/net/tools/quic/test_tools/mock_quic_dispatcher.cc b/net/tools/quic/test_tools/mock_quic_dispatcher.cc
index 13271ca..120c279 100644
--- a/net/tools/quic/test_tools/mock_quic_dispatcher.cc
+++ b/net/tools/quic/test_tools/mock_quic_dispatcher.cc
@@ -13,10 +13,12 @@ namespace test {
 MockQuicDispatcher::MockQuicDispatcher(
     const QuicConfig& config,
     const QuicCryptoServerConfig& crypto_config,
+    QuicDispatcher::PacketWriterFactory* packet_writer_factory,
     EpollServer* eps)
     : QuicDispatcher(config,
                      crypto_config,
                      QuicSupportedVersions(),
+                     packet_writer_factory,
                      eps) {}
 
 MockQuicDispatcher::~MockQuicDispatcher() {}
diff --git a/net/tools/quic/test_tools/mock_quic_dispatcher.h b/net/tools/quic/test_tools/mock_quic_dispatcher.h
index d155911..df32ce4 100644
--- a/net/tools/quic/test_tools/mock_quic_dispatcher.h
+++ b/net/tools/quic/test_tools/mock_quic_dispatcher.h
@@ -21,6 +21,7 @@ class MockQuicDispatcher : public QuicDispatcher {
  public:
   MockQuicDispatcher(const QuicConfig& config,
                      const QuicCryptoServerConfig& crypto_config,
+                     PacketWriterFactory* packet_writer_factory,
                      EpollServer* eps);
 
   virtual ~MockQuicDispatcher();
diff --git a/net/tools/quic/test_tools/packet_dropping_test_writer.h b/net/tools/quic/test_tools/packet_dropping_test_writer.h
index 3509722..b7babad 100644
--- a/net/tools/quic/test_tools/packet_dropping_test_writer.h
+++ b/net/tools/quic/test_tools/packet_dropping_test_writer.h
@@ -30,6 +30,7 @@ class PacketDroppingTestWriter : public QuicPacketWriterWrapper {
   class Delegate {
    public:
     virtual ~Delegate() {}
+    virtual void OnPacketSent(WriteResult result) = 0;
     virtual void OnCanWrite() = 0;
   };
 
diff --git a/net/tools/quic/test_tools/quic_dispatcher_peer.cc b/net/tools/quic/test_tools/quic_dispatcher_peer.cc
index 26ff490..1900420 100644
--- a/net/tools/quic/test_tools/quic_dispatcher_peer.cc
+++ b/net/tools/quic/test_tools/quic_dispatcher_peer.cc
@@ -31,6 +31,13 @@ QuicPacketWriter* QuicDispatcherPeer::GetWriter(QuicDispatcher* dispatcher) {
 }
 
 // static
+void QuicDispatcherPeer::SetPacketWriterFactory(
+    QuicDispatcher* dispatcher,
+    QuicDispatcher::PacketWriterFactory* packet_writer_factory) {
+  dispatcher->packet_writer_factory_.reset(packet_writer_factory);
+}
+
+// static
 QuicEpollConnectionHelper* QuicDispatcherPeer::GetHelper(
     QuicDispatcher* dispatcher) {
   return dispatcher->helper_.get();
diff --git a/net/tools/quic/test_tools/quic_dispatcher_peer.h b/net/tools/quic/test_tools/quic_dispatcher_peer.h
index 1d614d3..6271615 100644
--- a/net/tools/quic/test_tools/quic_dispatcher_peer.h
+++ b/net/tools/quic/test_tools/quic_dispatcher_peer.h
@@ -22,12 +22,16 @@ class QuicDispatcherPeer {
       QuicDispatcher* dispatcher,
       QuicTimeWaitListManager* time_wait_list_manager);
 
-  // Injects |writer| into |dispatcher| as the top level writer.
+  // Injects |writer| into |dispatcher| as the shared writer.
   static void UseWriter(QuicDispatcher* dispatcher,
                         QuicPacketWriterWrapper* writer);
 
   static QuicPacketWriter* GetWriter(QuicDispatcher* dispatcher);
 
+  static void SetPacketWriterFactory(
+      QuicDispatcher* dispatcher,
+      QuicDispatcher::PacketWriterFactory* packet_writer_factory);
+
   static QuicEpollConnectionHelper* GetHelper(QuicDispatcher* dispatcher);
 
   static QuicConnection* CreateQuicConnection(
diff --git a/net/tools/quic/test_tools/quic_test_utils.cc b/net/tools/quic/test_tools/quic_test_utils.cc
index 255fbc0..2e8c96d 100644
--- a/net/tools/quic/test_tools/quic_test_utils.cc
+++ b/net/tools/quic/test_tools/quic_test_utils.cc
@@ -18,12 +18,29 @@ namespace net {
 namespace tools {
 namespace test {
 
+namespace {
+class NiceMockPacketWriterFactory
+    : public QuicConnection::PacketWriterFactory {
+ public:
+  NiceMockPacketWriterFactory() {}
+  virtual ~NiceMockPacketWriterFactory() {}
+
+  virtual QuicPacketWriter* Create(
+      QuicConnection* /*connection*/) const override {
+    return new testing::NiceMock<MockPacketWriter>();
+  }
+
+ private:
+  DISALLOW_COPY_AND_ASSIGN(NiceMockPacketWriterFactory);
+};
+}  // namespace
+
 MockConnection::MockConnection(bool is_server)
     : QuicConnection(kTestConnectionId,
                      IPEndPoint(net::test::Loopback4(), kTestPort),
                      new testing::NiceMock<MockHelper>(),
-                     new testing::NiceMock<MockPacketWriter>(),
-                     true  /* owns_writer */,
+                     NiceMockPacketWriterFactory(),
+                     /* owns_writer= */ true,
                      is_server, QuicSupportedVersions()),
       helper_(helper()) {
 }
@@ -32,8 +49,8 @@ MockConnection::MockConnection(IPEndPoint address,
                                bool is_server)
     : QuicConnection(kTestConnectionId, address,
                      new testing::NiceMock<MockHelper>(),
-                     new testing::NiceMock<MockPacketWriter>(),
-                     true  /* owns_writer */,
+                     NiceMockPacketWriterFactory(),
+                     /* owns_writer= */ true,
                      is_server, QuicSupportedVersions()),
       helper_(helper()) {
 }
@@ -43,8 +60,8 @@ MockConnection::MockConnection(QuicConnectionId connection_id,
     : QuicConnection(connection_id,
                      IPEndPoint(net::test::Loopback4(), kTestPort),
                      new testing::NiceMock<MockHelper>(),
-                     new testing::NiceMock<MockPacketWriter>(),
-                     true  /* owns_writer */,
+                     NiceMockPacketWriterFactory(),
+                     /* owns_writer= */ true,
                      is_server, QuicSupportedVersions()),
       helper_(helper()) {
 }
@@ -54,8 +71,8 @@ MockConnection::MockConnection(bool is_server,
     : QuicConnection(kTestConnectionId,
                      IPEndPoint(net::test::Loopback4(), kTestPort),
                      new testing::NiceMock<MockHelper>(),
-                     new testing::NiceMock<MockPacketWriter>(),
-                     true  /* owns_writer */,
+                     NiceMockPacketWriterFactory(),
+                     /* owns_writer= */ true,
                      is_server, QuicSupportedVersions()),
       helper_(helper()) {
 }
@@ -112,6 +129,54 @@ MockAckNotifierDelegate::MockAckNotifierDelegate() {
 MockAckNotifierDelegate::~MockAckNotifierDelegate() {
 }
 
+TestWriterFactory::TestWriterFactory() : current_writer_(NULL) {}
+TestWriterFactory::~TestWriterFactory() {}
+
+QuicPacketWriter* TestWriterFactory::Create(QuicPacketWriter* writer,
+                                            QuicConnection* connection) {
+  return new PerConnectionPacketWriter(this, writer, connection);
+}
+
+void TestWriterFactory::OnPacketSent(WriteResult result) {
+  if (current_writer_ != NULL) {
+    current_writer_->connection()->OnPacketSent(result);
+    current_writer_ = NULL;
+  }
+}
+
+void TestWriterFactory::Unregister(PerConnectionPacketWriter* writer) {
+  if (current_writer_ == writer) {
+    current_writer_ = NULL;
+  }
+}
+
+TestWriterFactory::PerConnectionPacketWriter::PerConnectionPacketWriter(
+    TestWriterFactory* factory,
+    QuicPacketWriter* writer,
+    QuicConnection* connection)
+    : QuicPerConnectionPacketWriter(writer, connection),
+      factory_(factory) {
+}
+
+TestWriterFactory::PerConnectionPacketWriter::~PerConnectionPacketWriter() {
+  factory_->Unregister(this);
+}
+
+WriteResult TestWriterFactory::PerConnectionPacketWriter::WritePacket(
+    const char* buffer,
+    size_t buf_len,
+    const IPAddressNumber& self_address,
+    const IPEndPoint& peer_address) {
+  // A DCHECK(factory_current_writer_ == NULL) would be wrong here -- this class
+  // may be used in a setting where connection()->OnPacketSent() is called in a
+  // different way, so TestWriterFactory::OnPacketSent might never be called.
+  factory_->current_writer_ = this;
+  return QuicPerConnectionPacketWriter::WritePacket(buffer,
+                                                    buf_len,
+                                                    self_address,
+                                                    peer_address);
+}
+
 }  // namespace test
 }  // namespace tools
 }  // namespace net
diff --git a/net/tools/quic/test_tools/quic_test_utils.h b/net/tools/quic/test_tools/quic_test_utils.h
index b6449d3..1eca03b 100644
--- a/net/tools/quic/test_tools/quic_test_utils.h
+++ b/net/tools/quic/test_tools/quic_test_utils.h
@@ -14,6 +14,8 @@
 #include "net/quic/quic_packet_writer.h"
 #include "net/quic/quic_session.h"
 #include "net/spdy/spdy_framer.h"
+#include "net/tools/quic/quic_dispatcher.h"
+#include "net/tools/quic/quic_per_connection_packet_writer.h"
 #include "net/tools/quic/quic_server_session.h"
 #include "testing/gmock/include/gmock/gmock.h"
 
@@ -158,6 +160,46 @@ class MockAckNotifierDelegate : public QuicAckNotifier::DelegateInterface {
   DISALLOW_COPY_AND_ASSIGN(MockAckNotifierDelegate);
 };
 
+// Creates per-connection packet writers that register themselves with the
+// TestWriterFactory on each write so that TestWriterFactory::OnPacketSent can
+// be routed to the appropriate QuicConnection.
+class TestWriterFactory : public QuicDispatcher::PacketWriterFactory {
+ public:
+  TestWriterFactory();
+  virtual ~TestWriterFactory();
+
+  virtual QuicPacketWriter* Create(QuicPacketWriter* writer,
+                                   QuicConnection* connection) override;
+
+  // Calls OnPacketSent on the last QuicConnection to write through one of the
+  // packet writers created by this factory.
+  void OnPacketSent(WriteResult result);
+
+ private:
+  class PerConnectionPacketWriter : public QuicPerConnectionPacketWriter {
+   public:
+    PerConnectionPacketWriter(TestWriterFactory* factory,
+                              QuicPacketWriter* writer,
+                              QuicConnection* connection);
+    virtual ~PerConnectionPacketWriter();
+
+    virtual WriteResult WritePacket(
+        const char* buffer,
+        size_t buf_len,
+        const IPAddressNumber& self_address,
+        const IPEndPoint& peer_address) OVERRIDE;
+
+   private:
+    TestWriterFactory* factory_;
+  };
+
+  // If an asynchronous write is happening and |writer| gets deleted, this
+  // clears the pointer to it to prevent use-after-free.
+  void Unregister(PerConnectionPacketWriter* writer);
+
+  PerConnectionPacketWriter* current_writer_;
+};
+
 }  // namespace test
 }  // namespace tools
 }  // namespace net