A few changes to connection migration code which is not being used:

Change QuicForceBlockablePacketWriter from a pure virtual interface to an implementation of QuicPacketWriterWrapper to avoid diamond inheritance in existing test packet writers which are already inheriting from QuicPacketWriterWrapper and will be used by QuicSpdyClientSessionMigration in tests in the future.

Implement TimeSinceLastStreamClose() instead of leaving it as a pure virtual interface.

Change QuicSpdyClientSessionWithMigration::OnServerPreferredAddressAvailable() also calls into base class OnServerPreferredAddressAvailable()

PiperOrigin-RevId: 813943031
diff --git a/build/source_list.bzl b/build/source_list.bzl
index f023778..8ddf213 100644
--- a/build/source_list.bzl
+++ b/build/source_list.bzl
@@ -672,6 +672,7 @@
     "quic/core/quic_dispatcher_stats.cc",
     "quic/core/quic_error_codes.cc",
     "quic/core/quic_flow_controller.cc",
+    "quic/core/quic_force_blockable_packet_writer.cc",
     "quic/core/quic_framer.cc",
     "quic/core/quic_generic_session.cc",
     "quic/core/quic_idle_network_detector.cc",
diff --git a/build/source_list.gni b/build/source_list.gni
index f0cbd28..4e627df 100644
--- a/build/source_list.gni
+++ b/build/source_list.gni
@@ -672,6 +672,7 @@
     "src/quiche/quic/core/quic_dispatcher_stats.cc",
     "src/quiche/quic/core/quic_error_codes.cc",
     "src/quiche/quic/core/quic_flow_controller.cc",
+    "src/quiche/quic/core/quic_force_blockable_packet_writer.cc",
     "src/quiche/quic/core/quic_framer.cc",
     "src/quiche/quic/core/quic_generic_session.cc",
     "src/quiche/quic/core/quic_idle_network_detector.cc",
diff --git a/build/source_list.json b/build/source_list.json
index 587bf76..8a20ecb 100644
--- a/build/source_list.json
+++ b/build/source_list.json
@@ -671,6 +671,7 @@
     "quiche/quic/core/quic_dispatcher_stats.cc",
     "quiche/quic/core/quic_error_codes.cc",
     "quiche/quic/core/quic_flow_controller.cc",
+    "quiche/quic/core/quic_force_blockable_packet_writer.cc",
     "quiche/quic/core/quic_framer.cc",
     "quiche/quic/core/quic_generic_session.cc",
     "quiche/quic/core/quic_idle_network_detector.cc",
diff --git a/quiche/quic/core/http/quic_connection_migration_manager_test.cc b/quiche/quic/core/http/quic_connection_migration_manager_test.cc
index 350f489..81a48ac 100644
--- a/quiche/quic/core/http/quic_connection_migration_manager_test.cc
+++ b/quiche/quic/core/http/quic_connection_migration_manager_test.cc
@@ -43,75 +43,6 @@
   }
 };
 
-class TestQuicForceBlockablePacketWriter
-    : public QuicForceBlockablePacketWriter {
- public:
-  TestQuicForceBlockablePacketWriter() : QuicForceBlockablePacketWriter() {
-    ON_CALL(writer_, WritePacket(_, _, _, _, _, _))
-        .WillByDefault(Return(WriteResult(WRITE_STATUS_OK, 0)));
-    ON_CALL(writer_, GetMaxPacketSize(_))
-        .WillByDefault(Return(kMaxOutgoingPacketSize));
-    ON_CALL(writer_, IsBatchMode()).WillByDefault(Return(false));
-    ON_CALL(writer_, GetNextWriteLocation(_, _))
-        .WillByDefault(Return(QuicPacketBuffer()));
-    ON_CALL(writer_, Flush())
-        .WillByDefault(Return(WriteResult(WRITE_STATUS_OK, 0)));
-    ON_CALL(writer_, SupportsReleaseTime()).WillByDefault(Return(false));
-  }
-
-  // QuicPacketWriter.
-  WriteResult WritePacket(const char* buffer, size_t buf_len,
-                          const QuicIpAddress& self_address,
-                          const QuicSocketAddress& peer_address,
-                          PerPacketOptions* options,
-                          const QuicPacketWriterParams& params) override {
-    return writer_.WritePacket(buffer, buf_len, self_address, peer_address,
-                               options, params);
-  }
-
-  void SetWritable() override { writer_.SetWritable(); }
-
-  std::optional<int> MessageTooBigErrorCode() const override {
-    return kSocketErrorMsgSize;
-  }
-
-  QuicByteCount GetMaxPacketSize(
-      const QuicSocketAddress& peer_address) const override {
-    return writer_.GetMaxPacketSize(peer_address);
-  }
-
-  bool SupportsReleaseTime() const override {
-    return writer_.SupportsReleaseTime();
-  }
-
-  bool IsBatchMode() const override { return writer_.IsBatchMode(); }
-
-  bool SupportsEcn() const override { return writer_.SupportsEcn(); }
-
-  QuicPacketBuffer GetNextWriteLocation(
-      const QuicIpAddress& self_address,
-      const QuicSocketAddress& peer_address) override {
-    return writer_.GetNextWriteLocation(self_address, peer_address);
-  }
-
-  WriteResult Flush() override { return writer_.Flush(); }
-
-  bool IsWriteBlocked() const override {
-    return force_write_blocked_ || writer_.IsWriteBlocked();
-  }
-
-  // QuicForceBlockablePacketWriter
-  void ForceWriteBlocked(bool enforce_write_block) override {
-    force_write_blocked_ = enforce_write_block;
-  }
-
-  MockPacketWriter& GetMockWriter() { return writer_; }
-
- private:
-  NiceMock<MockPacketWriter> writer_;
-  bool force_write_blocked_ = false;
-};
-
 class TestQuicClientPathValidationContext
     : public QuicClientPathValidationContext {
  public:
@@ -119,7 +50,23 @@
       const quic::QuicSocketAddress& self_address,
       const quic::QuicSocketAddress& peer_address, QuicNetworkHandle network)
       : QuicClientPathValidationContext(self_address, peer_address, network),
-        writer_(std::make_unique<TestQuicForceBlockablePacketWriter>()) {}
+        writer_(std::make_unique<QuicForceBlockablePacketWriter>()) {
+    auto* writer = new NiceMock<MockPacketWriter>();
+    // Owns writer.
+    writer_->set_writer(writer);
+    ON_CALL(*writer, WritePacket(_, _, _, _, _, _))
+        .WillByDefault(Return(WriteResult(WRITE_STATUS_OK, 0)));
+    ON_CALL(*writer, GetMaxPacketSize(_))
+        .WillByDefault(Return(kMaxOutgoingPacketSize));
+    ON_CALL(*writer, IsBatchMode()).WillByDefault(Return(false));
+    ON_CALL(*writer, GetNextWriteLocation(_, _))
+        .WillByDefault(Return(QuicPacketBuffer()));
+    ON_CALL(*writer, Flush())
+        .WillByDefault(Return(WriteResult(WRITE_STATUS_OK, 0)));
+    ON_CALL(*writer, SupportsReleaseTime()).WillByDefault(Return(false));
+    ON_CALL(*writer, MessageTooBigErrorCode())
+        .WillByDefault(Return(kSocketErrorMsgSize));
+  }
 
   QuicForceBlockablePacketWriter* ForceBlockableWriterToUse() override {
     return writer_.get();
@@ -130,7 +77,7 @@
   void ReleasePacketWriter() { writer_.release(); }
 
  private:
-  std::unique_ptr<TestQuicForceBlockablePacketWriter> writer_;
+  std::unique_ptr<QuicForceBlockablePacketWriter> writer_;
 };
 
 class TestQuicPathContextFactory : public QuicPathContextFactory {
@@ -372,7 +319,7 @@
       : QuicSpdyClientSessionWithMigration(
             connection, writer, visitor, config, supported_versions,
             default_network, current_network, std::move(path_context_factory),
-            migration_config),
+            migration_config, QuicPriorityType::kHttp),
         crypto_stream_(this) {
     ON_CALL(*this, IsSessionProxied()).WillByDefault(Return(false));
     ON_CALL(*this, OnMigrationToPathDone(_, _))
@@ -391,7 +338,6 @@
   MOCK_METHOD(void, PrepareForProbingOnPath,
               (QuicPathValidationContext & context), (override));
   MOCK_METHOD(bool, IsSessionProxied, (), (const));
-  MOCK_METHOD(QuicTimeDelta, TimeSinceLastStreamClose, (), (override));
   MOCK_METHOD(bool, PrepareForMigrationToPath,
               (QuicClientPathValidationContext&));
   MOCK_METHOD(void, OnMigrationToPathDone,
@@ -488,7 +434,22 @@
         connection_(new StrictMock<test::MockQuicConnection>(
             &connection_helper_, &alarm_factory_, Perspective::IS_CLIENT,
             versions_)),
-        default_writer_(new NiceMock<TestQuicForceBlockablePacketWriter>()) {
+        default_writer_(new QuicForceBlockablePacketWriter()) {
+    auto* writer = new NiceMock<MockPacketWriter>();
+    // Owns writer.
+    default_writer_->set_writer(writer);
+    ON_CALL(*writer, WritePacket(_, _, _, _, _, _))
+        .WillByDefault(Return(WriteResult(WRITE_STATUS_OK, 0)));
+    ON_CALL(*writer, GetMaxPacketSize(_))
+        .WillByDefault(Return(kMaxOutgoingPacketSize));
+    ON_CALL(*writer, IsBatchMode()).WillByDefault(Return(false));
+    ON_CALL(*writer, GetNextWriteLocation(_, _))
+        .WillByDefault(Return(QuicPacketBuffer()));
+    ON_CALL(*writer, Flush())
+        .WillByDefault(Return(WriteResult(WRITE_STATUS_OK, 0)));
+    ON_CALL(*writer, SupportsReleaseTime()).WillByDefault(Return(false));
+    ON_CALL(*writer, MessageTooBigErrorCode())
+        .WillByDefault(Return(kSocketErrorMsgSize));
     connection_->SetQuicPacketWriter(default_writer_, true);
   }
 
@@ -576,7 +537,7 @@
   TestQuicPathContextFactory* path_context_factory_ = nullptr;
   // Owned by |session_|
   StrictMock<test::MockQuicConnection>* connection_;
-  NiceMock<TestQuicForceBlockablePacketWriter>* default_writer_;
+  QuicForceBlockablePacketWriter* default_writer_;
   std::unique_ptr<TestQuicSpdyClientSessionWithMigration> session_;
   QuicConnectionMigrationManager* migration_manager_ = nullptr;
   bool connection_migration_on_path_degrading_ = true;
@@ -661,8 +622,8 @@
   path_context_factory_->SetSelfAddressForNetwork(alternate_network,
                                                   alternate_self_address);
 
-  EXPECT_CALL(*session_, TimeSinceLastStreamClose())
-      .WillOnce(Return(QuicTimeDelta::Zero()));
+  EXPECT_EQ(session_->TimeSinceLastStreamClose(),
+            QuicTimeDelta::FromSeconds(1));
   EXPECT_CALL(*session_, ResetNonMigratableStreams());
   EXPECT_CALL(*session_, PrepareForMigrationToPath(_)).WillOnce(Return(true));
   EXPECT_CALL(*session_, OnMigrationToPathDone(_, true));
@@ -691,8 +652,8 @@
                                                   alternate_self_address);
   QuicSocketAddress self_address = connection_->self_address();
 
-  EXPECT_CALL(*session_, TimeSinceLastStreamClose())
-      .WillOnce(Return(QuicTimeDelta::Zero()));
+  EXPECT_EQ(session_->TimeSinceLastStreamClose(),
+            QuicTimeDelta::FromSeconds(1));
   EXPECT_CALL(*session_, ResetNonMigratableStreams());
   EXPECT_CALL(*session_, PrepareForMigrationToPath(_)).WillOnce(Return(true));
   EXPECT_CALL(*session_, OnMigrationToPathDone(_, true));
@@ -746,8 +707,8 @@
   EXPECT_EQ(migrate_back_alarm->deadline(),
             connection_helper_.GetClock()->Now());
   // Fire the alarm to migrate back to default network, starting with probing.
-  EXPECT_CALL(*session_, TimeSinceLastStreamClose())
-      .WillOnce(Return(QuicTimeDelta::Zero()));
+  EXPECT_EQ(session_->TimeSinceLastStreamClose(),
+            QuicTimeDelta::FromSeconds(2));
   QuicPathFrameBuffer path_frame_payload;
   EXPECT_CALL(*session_, PrepareForProbingOnPath(_));
   EXPECT_CALL(*connection_, SendPathChallenge(_, _, _, _, _))
@@ -773,8 +734,8 @@
                                                       self_address2);
   const QuicPathResponseFrame path_response(0, path_frame_payload);
   EXPECT_CALL(*session_, ResetNonMigratableStreams());
-  EXPECT_CALL(*session_, TimeSinceLastStreamClose())
-      .WillOnce(Return(QuicTimeDelta::Zero()));
+  EXPECT_EQ(session_->TimeSinceLastStreamClose(),
+            QuicTimeDelta::FromSeconds(2));
   EXPECT_CALL(*session_, PrepareForMigrationToPath(_)).WillOnce(Return(true));
   EXPECT_CALL(*session_, OnMigrationToPathDone(_, true));
   connection_->ReallyOnPathResponseFrame(path_response);
@@ -803,8 +764,8 @@
 
   // Receive a network disconnected signal, migration should be attempted
   // immediately.
-  EXPECT_CALL(*session_, TimeSinceLastStreamClose())
-      .WillOnce(Return(QuicTimeDelta::Zero()));
+  EXPECT_EQ(session_->TimeSinceLastStreamClose(),
+            QuicTimeDelta::FromSeconds(1));
   EXPECT_CALL(*session_, ResetNonMigratableStreams());
   migration_manager_->OnNetworkDisconnected(initial_network_);
   EXPECT_EQ(path_context_factory_->num_creation_attempts(), 1u);
@@ -852,8 +813,8 @@
   session_->set_alternate_network(alternate_network);
   EXPECT_NE(alternate_network, migration_manager_->current_network());
 
-  EXPECT_CALL(*session_, TimeSinceLastStreamClose())
-      .WillOnce(Return(migration_config_.idle_migration_period));
+  connection_helper_.GetClock()->AdvanceTime(
+      migration_config_.idle_migration_period);
   EXPECT_CALL(
       *connection_,
       CloseConnection(QUIC_NETWORK_IDLE_TIMEOUT,
@@ -889,8 +850,8 @@
   EXPECT_NE(alternate_network, migration_manager_->current_network());
   EXPECT_TRUE(session_->config()->DisableConnectionMigration());
 
-  EXPECT_CALL(*session_, TimeSinceLastStreamClose())
-      .WillOnce(Return(QuicTimeDelta::Zero()));
+  EXPECT_EQ(session_->TimeSinceLastStreamClose(),
+            QuicTimeDelta::FromSeconds(1));
   EXPECT_CALL(*connection_,
               CloseConnection(QUIC_CONNECTION_MIGRATION_DISABLED_BY_CONFIG,
                               "Migration disabled by config",
@@ -905,8 +866,8 @@
   Initialize();
   EXPECT_TRUE(session_->config()->DisableConnectionMigration());
 
-  EXPECT_CALL(*session_, TimeSinceLastStreamClose())
-      .WillOnce(Return(QuicTimeDelta::Zero()));
+  EXPECT_EQ(session_->TimeSinceLastStreamClose(),
+            QuicTimeDelta::FromSeconds(1));
   migration_manager_->MaybeStartMigrateSessionOnWriteError(/*error_code=*/111);
   // An alarm should have been scheduled to run pending callbacks.
   QuicAlarm* pending_callbacks_alarm =
@@ -940,8 +901,8 @@
   path_context_factory_->SetSelfAddressForNetwork(alternate_network,
                                                   alternate_self_address);
 
-  EXPECT_CALL(*session_, TimeSinceLastStreamClose())
-      .WillOnce(Return(QuicTimeDelta::Zero()));
+  EXPECT_EQ(session_->TimeSinceLastStreamClose(),
+            QuicTimeDelta::FromSeconds(1));
   EXPECT_CALL(*session_, ResetNonMigratableStreams());
   EXPECT_CALL(*session_, PrepareForMigrationToPath(_)).WillOnce(Return(true));
   EXPECT_CALL(*session_, OnMigrationToPathDone(_, true));
@@ -963,7 +924,6 @@
   // The migrate back timer will fire. Due to default network being
   // disconnected, no attempt will be exercised to migrate back.
   connection_helper_.GetClock()->AdvanceTime(QuicTimeDelta::FromSeconds(1));
-  QuicTimeDelta time_since_last_stream_close = QuicTimeDelta::FromSeconds(1);
   EXPECT_CALL(*session_, PrepareForProbingOnPath(_)).Times(0);
   alarm_factory_.FireAlarm(migrate_back_alarm);
   EXPECT_EQ(path_context_factory_->num_creation_attempts(), 1u);
@@ -978,8 +938,6 @@
     path_context_factory_->SetSelfAddressForNetwork(
         initial_network_,
         QuicSocketAddress(QuicIpAddress::Loopback4(), kTestPort + i));
-    EXPECT_CALL(*session_, TimeSinceLastStreamClose())
-        .WillOnce(Return(time_since_last_stream_close));
     // Update CIDs.
     QuicConnectionPeer::RetirePeerIssuedConnectionIdsNoLongerOnPath(
         connection_);
@@ -1020,15 +978,12 @@
     EXPECT_EQ(migrate_back_alarm->deadline(),
               connection_helper_.GetClock()->Now() + next_delay);
     connection_helper_.GetClock()->AdvanceTime(next_delay);
-    time_since_last_stream_close = time_since_last_stream_close + next_delay;
   }
 
   // The connection should have been idle for longer than the idle migration
   // period. Next attempt to migrate back will close the connection.
-  EXPECT_GT(time_since_last_stream_close,
+  EXPECT_GT(session_->TimeSinceLastStreamClose(),
             migration_config_.idle_migration_period);
-  EXPECT_CALL(*session_, TimeSinceLastStreamClose())
-      .WillOnce(Return(time_since_last_stream_close));
   //  The connection should be closed instead of attempting to migrate back.
   EXPECT_CALL(
       *connection_,
@@ -1970,8 +1925,8 @@
   path_context_factory_->SetSelfAddressForNetwork(alternate_network,
                                                   alternate_self_address);
 
-  EXPECT_CALL(*session_, TimeSinceLastStreamClose())
-      .WillOnce(Return(QuicTimeDelta::Zero()));
+  EXPECT_EQ(session_->TimeSinceLastStreamClose(),
+            QuicTimeDelta::FromSeconds(1));
   EXPECT_CALL(*session_, ResetNonMigratableStreams());
   // Session failed to prepare for migration. Migration should not be attempted.
   EXPECT_CALL(*session_, PrepareForMigrationToPath(_)).WillOnce(Return(false));
diff --git a/quiche/quic/core/http/quic_spdy_client_session_with_migration.cc b/quiche/quic/core/http/quic_spdy_client_session_with_migration.cc
index 4c9939b..a9f2c99 100644
--- a/quiche/quic/core/http/quic_spdy_client_session_with_migration.cc
+++ b/quiche/quic/core/http/quic_spdy_client_session_with_migration.cc
@@ -14,16 +14,24 @@
     const ParsedQuicVersionVector& supported_versions,
     QuicNetworkHandle default_network, QuicNetworkHandle current_network,
     std::unique_ptr<QuicPathContextFactory> path_context_factory,
-    QuicConnectionMigrationConfig migration_config)
-    : QuicSpdyClientSessionBase(connection, visitor, config,
-                                supported_versions),
+    QuicConnectionMigrationConfig migration_config,
+    QuicPriorityType priority_type)
+    : QuicSpdyClientSessionBase(connection, visitor, config, supported_versions,
+                                priority_type),
       path_context_factory_(std::move(path_context_factory)),
       migration_manager_(this, connection->clock(), default_network,
                          current_network, path_context_factory_.get(),
                          migration_config),
-      writer_(writer) {
-  QUICHE_DCHECK_EQ(writer_, connection->writer())
-      << "Writer is not the connection writer";
+      writer_(writer),
+      most_recent_stream_close_time_(connection->clock()->ApproximateNow()) {
+  QUICHE_DCHECK(writer_ == nullptr || writer_ == connection->writer())
+      << "Writer is should be either null or the connection writer";
+  if (migration_config.migrate_session_on_network_change ||
+      migration_config.allow_port_migration ||
+      migration_config.allow_server_preferred_address) {
+    QUICHE_DCHECK_EQ(writer_, connection->writer())
+        << "Writer is not the connection writer";
+  }
 }
 
 QuicSpdyClientSessionWithMigration::~QuicSpdyClientSessionWithMigration() =
@@ -73,6 +81,8 @@
 void QuicSpdyClientSessionWithMigration::OnServerPreferredAddressAvailable(
     const QuicSocketAddress& server_preferred_address) {
   QUICHE_DCHECK(version().HasIetfQuicFrames());
+  QuicSpdyClientSessionBase::OnServerPreferredAddressAvailable(
+      server_preferred_address);
   migration_manager_.MaybeStartMigrateSessionToServerPreferredAddress(
       server_preferred_address);
 }
@@ -87,4 +97,15 @@
   return migration_manager_.config();
 }
 
+void QuicSpdyClientSessionWithMigration::OnStreamClosed(
+    QuicStreamId stream_id) {
+  most_recent_stream_close_time_ = connection()->clock()->ApproximateNow();
+  QuicSpdyClientSessionBase::OnStreamClosed(stream_id);
+}
+
+QuicTimeDelta QuicSpdyClientSessionWithMigration::TimeSinceLastStreamClose() {
+  return connection()->clock()->ApproximateNow() -
+         most_recent_stream_close_time_;
+}
+
 }  // namespace quic
diff --git a/quiche/quic/core/http/quic_spdy_client_session_with_migration.h b/quiche/quic/core/http/quic_spdy_client_session_with_migration.h
index e434a3c..da219c5 100644
--- a/quiche/quic/core/http/quic_spdy_client_session_with_migration.h
+++ b/quiche/quic/core/http/quic_spdy_client_session_with_migration.h
@@ -33,13 +33,16 @@
 class QUICHE_EXPORT QuicSpdyClientSessionWithMigration
     : public QuicSpdyClientSessionBase {
  public:
+  // `writer` must be the same as the connection's writer if any type of
+  // migration is enabled. Otherwise, it can also be nullptr.
   QuicSpdyClientSessionWithMigration(
       QuicConnection* connection, QuicForceBlockablePacketWriter* writer,
       QuicSession::Visitor* visitor, const QuicConfig& config,
       const ParsedQuicVersionVector& supported_versions,
       QuicNetworkHandle default_network, QuicNetworkHandle current_network,
       std::unique_ptr<QuicPathContextFactory> path_context_factory,
-      QuicConnectionMigrationConfig migration_config);
+      QuicConnectionMigrationConfig migration_config,
+      QuicPriorityType priority_type);
 
   ~QuicSpdyClientSessionWithMigration() override;
 
@@ -68,7 +71,7 @@
   virtual bool IsSessionProxied() const = 0;
 
   // Returns the time elapsed since the latest stream closure.
-  virtual quic::QuicTimeDelta TimeSinceLastStreamClose() = 0;
+  quic::QuicTimeDelta TimeSinceLastStreamClose();
 
   // QuicSpdyClientSessionBase
   void OnPathDegrading() override;
@@ -76,6 +79,7 @@
   void SetDefaultEncryptionLevel(EncryptionLevel level) override;
   void OnServerPreferredAddressAvailable(
       const quic::QuicSocketAddress& server_preferred_address) override;
+  void OnStreamClosed(QuicStreamId stream_id) override;
 
   // Migrates session onto the new path, i.e. changing the default writer and
   // network.
@@ -113,6 +117,7 @@
   std::unique_ptr<QuicPathContextFactory> path_context_factory_;
   QuicConnectionMigrationManager migration_manager_;
   QuicForceBlockablePacketWriter* absl_nonnull writer_;
+  QuicTime most_recent_stream_close_time_;
 };
 
 }  // namespace quic
diff --git a/quiche/quic/core/quic_force_blockable_packet_writer.cc b/quiche/quic/core/quic_force_blockable_packet_writer.cc
new file mode 100644
index 0000000..10655d3
--- /dev/null
+++ b/quiche/quic/core/quic_force_blockable_packet_writer.cc
@@ -0,0 +1,18 @@
+// Copyright 2025 The Chromium Authors
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "quiche/quic/core/quic_force_blockable_packet_writer.h"
+
+namespace quic {
+
+void QuicForceBlockablePacketWriter::ForceWriteBlocked(
+    bool enforce_write_block) {
+  enforce_write_block_ = enforce_write_block;
+}
+
+bool QuicForceBlockablePacketWriter::IsWriteBlocked() const {
+  return enforce_write_block_ || QuicPacketWriterWrapper::IsWriteBlocked();
+}
+
+}  // namespace quic
diff --git a/quiche/quic/core/quic_force_blockable_packet_writer.h b/quiche/quic/core/quic_force_blockable_packet_writer.h
index a487885..677490f 100644
--- a/quiche/quic/core/quic_force_blockable_packet_writer.h
+++ b/quiche/quic/core/quic_force_blockable_packet_writer.h
@@ -5,20 +5,26 @@
 #ifndef QUICHE_QUIC_CORE_QUIC_FORCE_BLOCKABLE_PACKET_WRITER_H_
 #define QUICHE_QUIC_CORE_QUIC_FORCE_BLOCKABLE_PACKET_WRITER_H_
 
-#include "quiche/quic/core/quic_packet_writer.h"
+#include "quiche/quic/core/quic_packet_writer_wrapper.h"
 
 namespace quic {
 
-// A extended interface of QuicPacketWriter that can be forced to be write
+// A QuicPacketWriterWrapper implementation that can be forced to be write
 // blocked.
-class QUICHE_EXPORT QuicForceBlockablePacketWriter : public QuicPacketWriter {
+class QUICHE_EXPORT QuicForceBlockablePacketWriter
+    : public QuicPacketWriterWrapper {
  public:
   // If `enforce_write_block` is true, IsWriteBlocked() will always return true
   // regardless of whether SetWritable() is called or not until
-  // this method is called again with |enforce_write_block| false.
-  // If |enforce_write_block| is false, SetWritable() may still be needed to
+  // this method is called again with `enforce_write_block` false.
+  // If `enforce_write_block` is false, SetWritable() may still be needed to
   // make IsWriteBlocked() to return true.
-  virtual void ForceWriteBlocked(bool enforce_write_block) = 0;
+  void ForceWriteBlocked(bool enforce_write_block);
+
+  bool IsWriteBlocked() const override;
+
+ private:
+  bool enforce_write_block_ = false;
 };
 
 }  // namespace quic