From 3e29bcfb74ee493f02f101d19dd858c2d5e47bee Mon Sep 17 00:00:00 2001 From: Piotr Sikora Date: Tue, 3 Mar 2020 17:17:33 -0800 Subject: [PATCH] network: draining the buffer on close (#9870) (#110) (#179) Signed-off-by: Alyssa Wilk Co-authored-by: Lizan Zhou --- source/common/network/connection_impl.cc | 7 +++++++ test/common/network/connection_impl_test.cc | 8 ++++++-- test/extensions/transport_sockets/tls/ssl_socket_test.cc | 2 +- 3 files changed, 14 insertions(+), 3 deletions(-) diff --git a/source/common/network/connection_impl.cc b/source/common/network/connection_impl.cc index f0416b2feac7..f4cd73ab9689 100644 --- a/source/common/network/connection_impl.cc +++ b/source/common/network/connection_impl.cc @@ -195,6 +195,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 index ffb0ba7ecba1..b13f057cc577 100644 --- a/test/common/network/connection_impl_test.cc +++ b/test/common/network/connection_impl_test.cc @@ -141,6 +141,11 @@ class ConnectionImplTest : public testing::TestWithParam { } 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) { @@ -845,8 +850,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); @@ -924,6 +927,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 index d16fd5d9e7e9..6918791f1fcc 100644 --- a/test/extensions/transport_sockets/tls/ssl_socket_test.cc +++ b/test/extensions/transport_sockets/tls/ssl_socket_test.cc @@ -4126,7 +4126,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(); }));