tcp: extending tcp integration test

test/integration/autonomous_upstream.cc

@@ -21,6 +21,7 @@ const char AutonomousStream::RESPONSE_SIZE_BYTES[] = "response_size_bytes";
 const char AutonomousStream::RESPONSE_DATA_BLOCKS[] = "response_data_blocks";
 const char AutonomousStream::EXPECT_REQUEST_SIZE_BYTES[] = "expect_request_size_bytes";
 const char AutonomousStream::RESET_AFTER_REQUEST[] = "reset_after_request";
+const char AutonomousStream::CLOSE_AFTER_RESPONSE[] = "close_after_response";
 const char AutonomousStream::NO_TRAILERS[] = "no_trailers";
 const char AutonomousStream::NO_END_STREAM[] = "no_end_stream";
 
@@ -84,6 +85,11 @@ void AutonomousStream::sendResponse() {
       encodeTrailers(upstream_.responseTrailers());
     }
   }
+  if (!headers.get_(CLOSE_AFTER_RESPONSE).empty()) {
+    parent_.connection().dispatcher().post(
+        [this]() -> void { parent_.connection().close(Network::ConnectionCloseType::FlushWrite); });
+    return;
+  }
 }
 
 AutonomousHttpConnection::AutonomousHttpConnection(AutonomousUpstream& autonomous_upstream,

test/integration/autonomous_upstream.h

@@ -25,6 +25,8 @@ class AutonomousStream : public FakeStream {
   static const char NO_TRAILERS[];
   // Prevents upstream from finishing response.
   static const char NO_END_STREAM[];
+  // Closes the underlying connection after a given response is sent.
+  static const char CLOSE_AFTER_RESPONSE[];
 
   AutonomousStream(FakeHttpConnection& parent, Http::ResponseEncoder& encoder,
                    AutonomousUpstream& upstream, bool allow_incomplete_streams);

test/integration/fake_upstream.h

@@ -222,9 +222,9 @@ class FakeStream : public Http::RequestDecoder,
   absl::Mutex lock_;
   Http::RequestHeaderMapPtr headers_ ABSL_GUARDED_BY(lock_);
   Buffer::OwnedImpl body_ ABSL_GUARDED_BY(lock_);
+  FakeHttpConnection& parent_;
 
 private:
-  FakeHttpConnection& parent_;
   Http::ResponseEncoder& encoder_;
   Http::RequestTrailerMapPtr trailers_ ABSL_GUARDED_BY(lock_);
   bool end_stream_ ABSL_GUARDED_BY(lock_){};

test/integration/integration_tcp_client.cc

@@ -48,7 +48,12 @@ IntegrationTcpClient::IntegrationTcpClient(
         client_write_buffer_ =
             new NiceMock<MockWatermarkBuffer>(below_low, above_high, above_overflow);
         return client_write_buffer_;
+      }))
+      .WillRepeatedly(Invoke([](std::function<void()> below_low, std::function<void()> above_high,
+                                std::function<void()> above_overflow) -> Buffer::Instance* {
+        return new Buffer::WatermarkBuffer(below_low, above_high, above_overflow);
       }));
+  ;
 
   connection_ = dispatcher.createClientConnection(
       Network::Utility::resolveUrl(

test/integration/tcp_proxy_integration_test.cc

@@ -25,6 +25,7 @@
 
 using testing::_;
 using testing::AtLeast;
+using testing::HasSubstr;
 using testing::Invoke;
 using testing::MatchesRegex;
 using testing::NiceMock;
@@ -150,6 +151,73 @@ TEST_P(TcpProxyIntegrationTest, TcpProxyDownstreamDisconnect) {
   tcp_client->waitForDisconnect();
 }
 
+TEST_P(TcpProxyIntegrationTest, TcpProxyManyConnections) {
+  autonomous_upstream_ = true;
+  initialize();
+  int num_connections = 50;
+  std::vector<IntegrationTcpClientPtr> clients(num_connections);
+
+  for (int i = 0; i < num_connections; ++i) {
+    clients[i] = makeTcpConnection(lookupPort("tcp_proxy"));
+  }
+  test_server_->waitForCounterGe("cluster.cluster_0.upstream_cx_total", num_connections);
+  for (int i = 0; i < num_connections; ++i) {
+    IntegrationTcpClientPtr& tcp_client = clients[i];
+    ASSERT_TRUE(tcp_client->write(
+        "GET / HTTP/1.1\r\nHost: foo\r\nclose_after_response: yes\r\ncontent-length: 0\r\n\r\n",
+        false));
+    tcp_client->waitForHalfClose();
+    tcp_client->close();
+    EXPECT_THAT(tcp_client->data(), HasSubstr("aaaaaaaaaa"));
+  }
+}
+
+TEST_P(TcpProxyIntegrationTest, TcpProxyRandomBehavior) {
+  autonomous_upstream_ = true;
+  initialize();
+  std::list<IntegrationTcpClientPtr> clients;
+
+  const char* basic_request = "GET / HTTP/1.1\r\nHost: foo\r\ncontent-length: 0\r\n\r\n";
+  const char* request_with_close =
+      "GET / HTTP/1.1\r\nHost: foo\r\nclose_after_response: yes\r\ncontent-length: 0\r\n\r\n";
+  TestRandomGenerator rand;
+
+  for (int i = 0; i < 50; ++i) {
+    int action = rand.random() % 3;
+
+    if (action == 0) {
+      // Add a new connection.
+      clients.push_back(makeTcpConnection(lookupPort("tcp_proxy")));
+    }
+    if (clients.empty()) {
+      break;
+    }
+    IntegrationTcpClientPtr& tcp_client = clients.front();
+    if (action == 1) {
+      // Write to the first connection.
+      ASSERT_TRUE(tcp_client->write(basic_request, false));
+      tcp_client->waitForData("\r\n\r\n", false);
+      tcp_client->clearData(tcp_client->data().size());
+    } else if (action == 2) {
+      // Close the first connection.
+      ASSERT_TRUE(tcp_client->write(request_with_close, false));
+      tcp_client->waitForData("\r\n\r\n", false);
+      tcp_client->waitForHalfClose();
+      tcp_client->close();
+      clients.pop_front();
+    }
+  }
+
+  while (!clients.empty()) {
+    IntegrationTcpClientPtr& tcp_client = clients.front();
+    ASSERT_TRUE(tcp_client->write(request_with_close, false));
+    tcp_client->waitForData("\r\n\r\n", false);
+    tcp_client->waitForHalfClose();
+    tcp_client->close();
+    clients.pop_front();
+  }
+}
+
 TEST_P(TcpProxyIntegrationTest, NoUpstream) {
   // Set the first upstream to have an invalid port, so connection will fail,
   // but it won't fail synchronously (as it would if there were simply no