No public description PiperOrigin-RevId: 742788780

diff --git a/quiche/quic/core/quic_packet_writer_wrapper.cc b/quiche/quic/core/quic_packet_writer_wrapper.cc
index 6a61556..c22fea6 100644
--- a/quiche/quic/core/quic_packet_writer_wrapper.cc
+++ b/quiche/quic/core/quic_packet_writer_wrapper.cc

@@ -18,8 +18,12 @@
     const char* buffer, size_t buf_len, const QuicIpAddress& self_address,
     const QuicSocketAddress& peer_address, PerPacketOptions* options,
     const QuicPacketWriterParams& params) {
-  return writer_->WritePacket(buffer, buf_len, self_address, peer_address,
-                              options, params);
+  const WriteResult result = writer_->WritePacket(
+      buffer, buf_len, self_address, peer_address, options, params);
+  if (on_write_done_ != nullptr) {
+    on_write_done_(buf_len, result);
+  }
+  return result;
 }
 
 bool QuicPacketWriterWrapper::IsWriteBlocked() const {

diff --git a/quiche/quic/core/quic_packet_writer_wrapper.h b/quiche/quic/core/quic_packet_writer_wrapper.h
index 898a0e0..248aed5 100644
--- a/quiche/quic/core/quic_packet_writer_wrapper.h
+++ b/quiche/quic/core/quic_packet_writer_wrapper.h

@@ -7,8 +7,10 @@
 
 #include <cstddef>
 #include <memory>
+#include <utility>
 
 #include "quiche/quic/core/quic_packet_writer.h"
+#include "quiche/common/quiche_callbacks.h"
 
 namespace quic {
 
@@ -52,11 +54,23 @@
 
   QuicPacketWriter* writer() { return writer_; }
 
+  // First argument is the packet size. Second argument is the result of the
+  // write.
+  using OnWriteDoneCallback =
+      quiche::MultiUseCallback<void(size_t, const WriteResult&)>;
+
+  // If set, |on_write_done| will be called after each write.
+  void set_on_write_done(OnWriteDoneCallback on_write_done) {
+    on_write_done_ = std::move(on_write_done);
+  }
+
  private:
   void unset_writer();
 
   QuicPacketWriter* writer_ = nullptr;
   bool owns_writer_ = false;
+  // If not null, called after each write.
+  OnWriteDoneCallback on_write_done_;
 };
 
 }  // namespace quic