tcp client: make close type configurable (envoyproxy#26078)

Signed-off-by: Boteng Yao <[email protected]>
pchaseh · Mar 29, 2023 · 80e5a3f · 80e5a3f
1 parent 8e58191
commit 80e5a3f
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 13 deletions.
diff --git a/envoy/tcp/async_tcp_client.h b/envoy/tcp/async_tcp_client.h
@@ -45,10 +45,12 @@ class AsyncTcpClient {
   virtual bool connect() PURE;
 
   /**
-   * Close the client. It abortively closes the connection discarding any unsent data.
+   * Close the client. It closes the connection based on close type.
    * The underlying connection will be defer deleted when a Close event is received.
+   * Abrt/NoFlush will abortively closes the connection discarding any unsent data.
+   * @param type the connection close type.
    */
-  virtual void close() PURE;
+  virtual void close(Network::ConnectionCloseType type) PURE;
 
   /**
    * Write data through the client.

diff --git a/source/common/tcp/async_tcp_client_impl.cc b/source/common/tcp/async_tcp_client_impl.cc
@@ -21,8 +21,6 @@ AsyncTcpClientImpl::AsyncTcpClientImpl(Event::Dispatcher& dispatcher,
     : dispatcher_(dispatcher), thread_local_cluster_(thread_local_cluster), context_(context),
       enable_half_close_(enable_half_close) {}
 
-AsyncTcpClientImpl::~AsyncTcpClientImpl() { close(); }
-
 bool AsyncTcpClientImpl::connect() {
   connection_ = std::move(thread_local_cluster_.tcpConn(context_).connection_);
   if (!connection_) {
@@ -36,9 +34,9 @@ bool AsyncTcpClientImpl::connect() {
   return true;
 }
 
-void AsyncTcpClientImpl::close() {
+void AsyncTcpClientImpl::close(Network::ConnectionCloseType type) {
   if (connection_) {
-    connection_->close(Network::ConnectionCloseType::NoFlush);
+    connection_->close(type);
   }
 }
 

diff --git a/source/common/tcp/async_tcp_client_impl.h b/source/common/tcp/async_tcp_client_impl.h
@@ -23,9 +23,7 @@ class AsyncTcpClientImpl : public AsyncTcpClient, public Network::ConnectionCall
                      Upstream::ThreadLocalCluster& thread_local_cluster,
                      Upstream::LoadBalancerContext* context, bool enable_half_close);
 
-  ~AsyncTcpClientImpl() override;
-
-  void close() override;
+  void close(Network::ConnectionCloseType type) override;
 
   /**
    * @return true means a host is successfully picked from a Cluster.

diff --git a/test/common/tcp/async_tcp_client_impl_test.cc b/test/common/tcp/async_tcp_client_impl_test.cc
@@ -60,7 +60,7 @@ TEST_F(AsyncTcpClientImplTest, BasicWrite) {
 
   EXPECT_CALL(callbacks_, onEvent(Network::ConnectionEvent::LocalClose));
   EXPECT_CALL(dispatcher_, deferredDelete_(_));
-  client_->close();
+  client_->close(Network::ConnectionCloseType::NoFlush);
   ASSERT_FALSE(client_->connected());
 }
 
@@ -74,7 +74,7 @@ TEST_F(AsyncTcpClientImplTest, WaterMark) {
 
   EXPECT_CALL(callbacks_, onEvent(Network::ConnectionEvent::LocalClose));
   EXPECT_CALL(dispatcher_, deferredDelete_(_));
-  client_->close();
+  client_->close(Network::ConnectionCloseType::NoFlush);
   ASSERT_FALSE(client_->connected());
 }
 
@@ -95,7 +95,16 @@ TEST_F(AsyncTcpClientImplTest, TestReadDisable) {
 
   EXPECT_CALL(callbacks_, onEvent(Network::ConnectionEvent::LocalClose));
   EXPECT_CALL(dispatcher_, deferredDelete_(_));
-  client_->close();
+  client_->close(Network::ConnectionCloseType::NoFlush);
+  ASSERT_FALSE(client_->connected());
+}
+
+TEST_F(AsyncTcpClientImplTest, TestCloseType) {
+  expectCreateConnection();
+  EXPECT_CALL(callbacks_, onEvent(Network::ConnectionEvent::LocalClose));
+  EXPECT_CALL(*connection_, close(Network::ConnectionCloseType::Abort));
+  EXPECT_CALL(dispatcher_, deferredDelete_(_));
+  client_->close(Network::ConnectionCloseType::Abort);
   ASSERT_FALSE(client_->connected());
 }
 

diff --git a/test/integration/filters/test_network_async_tcp_filter.cc b/test/integration/filters/test_network_async_tcp_filter.cc
@@ -84,7 +84,7 @@ class TestNetworkAsyncTcpFilter : public Network::ReadFilter {
                      data.length());
       parent_.read_callbacks_->connection().write(data, end_stream);
       if (end_stream) {
-        parent_.client_->close();
+        parent_.client_->close(Network::ConnectionCloseType::NoFlush);
       }
     }