network: draining the buffer on close (envoyproxy#9870) (#110) (#124)

Signed-off-by: Alyssa Wilk <alyssar@chromium.org> Signed-off-by: Jianfei Hu <jianfeih@google.com> Co-authored-by: Lizan Zhou <lizan@tetrate.io>
istio · Mar 3, 2020 · b2fc186 · b2fc186
1 parent 92bd1e9
commit b2fc186
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 3 deletions.
diff --git a/source/common/network/connection_impl.cc b/source/common/network/connection_impl.cc
@@ -196,6 +196,13 @@ void ConnectionImpl::closeSocket(ConnectionEvent close_type) {
   // Drain input and output buffers.
   updateReadBufferStats(0, 0);
   updateWriteBufferStats(0, 0);
+
+  // As the socket closes, drain any remaining data.
+  // The data won't be written out at this point, and where there are reference
+  // counted buffer fragments, it helps avoid lifetime issues with the
+  // connection outlasting the subscriber.
+  write_buffer_->drain(write_buffer_->length());
+
   connection_stats_.reset();
 
   file_event_.reset();

diff --git a/test/common/network/connection_impl_test.cc b/test/common/network/connection_impl_test.cc
@@ -136,6 +136,11 @@ class ConnectionImplTest : public testing::TestWithParam<Address::IpVersion> {
   }
 
   void disconnect(bool wait_for_remote_close) {
+    if (client_write_buffer_) {
+      EXPECT_CALL(*client_write_buffer_, drain(_))
+          .Times(AnyNumber())
+          .WillOnce(Invoke([&](uint64_t size) -> void { client_write_buffer_->baseDrain(size); }));
+    }
     EXPECT_CALL(client_callbacks_, onEvent(ConnectionEvent::LocalClose));
     client_connection_->close(ConnectionCloseType::NoFlush);
     if (wait_for_remote_close) {
@@ -810,8 +815,6 @@ TEST_P(ConnectionImplTest, WriteWithWatermarks) {
   // call to write() will succeed, bringing the connection back under the low watermark.
   EXPECT_CALL(*client_write_buffer_, write(_))
       .WillOnce(Invoke(client_write_buffer_, &MockWatermarkBuffer::trackWrites));
-  EXPECT_CALL(*client_write_buffer_, drain(_))
-      .WillOnce(Invoke(client_write_buffer_, &MockWatermarkBuffer::trackDrains));
   EXPECT_CALL(client_callbacks_, onBelowWriteBufferLowWatermark()).Times(1);
 
   disconnect(true);
@@ -889,6 +892,7 @@ TEST_P(ConnectionImplTest, WatermarkFuzzing) {
     dispatcher_->run(Event::Dispatcher::RunType::NonBlock);
   }
 
+  EXPECT_CALL(client_callbacks_, onBelowWriteBufferLowWatermark()).Times(AnyNumber());
   disconnect(true);
 }
 

diff --git a/test/extensions/transport_sockets/tls/ssl_socket_test.cc b/test/extensions/transport_sockets/tls/ssl_socket_test.cc
@@ -3986,7 +3986,7 @@ class SslReadBufferLimitTest : public SslSocketTest {
     EXPECT_CALL(*client_write_buffer, move(_))
         .WillRepeatedly(DoAll(AddBufferToStringWithoutDraining(&data_written),
                               Invoke(client_write_buffer, &MockWatermarkBuffer::baseMove)));
-    EXPECT_CALL(*client_write_buffer, drain(_)).WillOnce(Invoke([&](uint64_t n) -> void {
+    EXPECT_CALL(*client_write_buffer, drain(_)).Times(2).WillOnce(Invoke([&](uint64_t n) -> void {
       client_write_buffer->baseDrain(n);
       dispatcher_->exit();
     }));