tcp_proxy: explicitly supporting upstream-write-first semantics in the new pool (envoyproxy#14919)

Additional Description:
Both the old and the new TCP proxy pools register for "unexpected data" and disconnect when receiving it. It appears that generally the timing works out in the old pool that the connection gets assigned right away, but the new pool doesn't pass the race deterministically. Fixing this by read-disabling, so data will be associated with the downstream connection as it is associated.
Risk Level: Medium (new pool is still off by default, but affects ALPN as well)
Testing: new integration tests
Docs Changes: n/a
Release Notes: n/a

Signed-off-by: Alyssa Wilk <[email protected]>

Author: alyssawilk

source/common/conn_pool/conn_pool_base.h

@@ -197,7 +197,8 @@ class ConnPoolImplBase : protected Logger::Loggable<Logger::Id::pool> {
     const char* spaces = spacesForLevel(indent_level);
     os << spaces << "ConnPoolImplBase " << this << DUMP_MEMBER(ready_clients_.size())
        << DUMP_MEMBER(busy_clients_.size()) << DUMP_MEMBER(connecting_clients_.size())
-       << DUMP_MEMBER(connecting_stream_capacity_) << DUMP_MEMBER(num_active_streams_);
+       << DUMP_MEMBER(connecting_stream_capacity_) << DUMP_MEMBER(num_active_streams_)
+       << DUMP_MEMBER(pending_streams_.size());
   }
 
   friend std::ostream& operator<<(std::ostream& os, const ConnPoolImplBase& s) {
@@ -206,7 +207,6 @@ class ConnPoolImplBase : protected Logger::Loggable<Logger::Id::pool> {
   }
 
 protected:
-  // Creates up to 3 connections, based on the preconnect ratio.
   virtual void onConnected(Envoy::ConnectionPool::ActiveClient&) {}
 
   enum class ConnectionResult {

source/common/http/mixed_conn_pool.cc

@@ -50,6 +50,10 @@ void HttpConnPoolImplMixed::onConnected(Envoy::ConnectionPool::ActiveClient& cli
 
   Upstream::Host::CreateConnectionData data{std::move(tcp_client->connection_),
                                             client.real_host_description_};
+  // As this connection comes from the tcp connection pool, it will be
+  // read-disabled to handle TCP traffic where upstream sends data first. Undo
+  // this as it is not necessary for HTTP/HTTPS.
+  data.connection_->readDisable(false);
   data.connection_->removeConnectionCallbacks(*tcp_client);
   data.connection_->removeReadFilter(tcp_client->read_filter_handle_);
   dispatcher_.deferredDelete(client.removeFromList(owningList(client.state_)));

source/common/tcp/conn_pool.cc

@@ -24,7 +24,6 @@ ActiveTcpClient::ActiveTcpClient(Envoy::ConnectionPool::ConnPoolImplBase& parent
   real_host_description_ = data.host_description_;
   connection_ = std::move(data.connection_);
   connection_->addConnectionCallbacks(*this);
-  connection_->detectEarlyCloseWhenReadDisabled(false);
   read_filter_handle_ = std::make_shared<ConnReadFilter>(*this);
   connection_->addReadFilter(read_filter_handle_);
   connection_->setConnectionStats({host->cluster().stats().upstream_cx_rx_bytes_total_,
@@ -63,6 +62,13 @@ void ActiveTcpClient::clearCallbacks() {
 }
 
 void ActiveTcpClient::onEvent(Network::ConnectionEvent event) {
+  // If this is a newly established TCP connection, readDisable. This is to handle a race condition
+  // for TCP for protocols like MySQL where the upstream writes first, and the data needs to be
+  // preserved until a downstream connection is associated.
+  // This is also necessary for prefetch to be used with such protocols.
+  if (event == Network::ConnectionEvent::Connected) {
+    connection_->readDisable(true);
+  }
   Envoy::ConnectionPool::ActiveClient::onEvent(event);
   if (callbacks_) {
     // Do not pass the Connected event to any session which registered during onEvent above.

source/common/tcp/conn_pool.h

@@ -94,6 +94,21 @@ class ActiveTcpClient : public Envoy::ConnectionPool::ActiveClient {
   void onAboveWriteBufferHighWatermark() override { callbacks_->onAboveWriteBufferHighWatermark(); }
   void onBelowWriteBufferLowWatermark() override { callbacks_->onBelowWriteBufferLowWatermark(); }
 
+  // Undo the readDisable done in onEvent(Connected) - now that there is an associated connection,
+  // drain any data.
+  void readEnableIfNew() {
+    // It is expected for Envoy use of ActiveTcpClient this function only be
+    // called once. Other users of the TcpConnPool may recycle Tcp connections,
+    // and this safeguards them against read-enabling too many times.
+    if (!associated_before_) {
+      associated_before_ = true;
+      connection_->readDisable(false);
+      // Also while we're at it, make sure the connection will proxy all TCP
+      // data before picking up a FIN.
+      connection_->detectEarlyCloseWhenReadDisabled(false);
+    }
+  }
+
   absl::optional<Http::Protocol> protocol() const override { return {}; }
   void close() override { connection_->close(Network::ConnectionCloseType::NoFlush); }
   uint32_t numActiveStreams() const override { return callbacks_ ? 1 : 0; }
@@ -115,6 +130,7 @@ class ActiveTcpClient : public Envoy::ConnectionPool::ActiveClient {
   Network::ClientConnectionPtr connection_;
   ConnectionPool::ConnectionStatePtr connection_state_;
   TcpConnectionData* tcp_connection_data_{};
+  bool associated_before_{};
 };
 
 class ConnPoolImpl : public Envoy::ConnectionPool::ConnPoolImplBase,
@@ -180,6 +196,7 @@ class ConnPoolImpl : public Envoy::ConnectionPool::ConnPoolImplBase,
   void onPoolReady(Envoy::ConnectionPool::ActiveClient& client,
                    Envoy::ConnectionPool::AttachContext& context) override {
     ActiveTcpClient* tcp_client = static_cast<ActiveTcpClient*>(&client);
+    tcp_client->readEnableIfNew();
     auto* callbacks = typedContext<TcpAttachContext>(context).callbacks_;
     std::unique_ptr<Envoy::Tcp::ConnectionPool::ConnectionData> connection_data =
         std::make_unique<ActiveTcpClient::TcpConnectionData>(*tcp_client, *tcp_client->connection_);

test/common/tcp/conn_pool_test.cc

@@ -321,6 +321,7 @@ class TcpConnPoolImplDestructorTest : public Event::TestUsingSimulatedTime,
     EXPECT_CALL(*connection_, setConnectionStats(_));
     EXPECT_CALL(*connection_, streamInfo()).Times(2);
     EXPECT_CALL(*connection_, id()).Times(AnyNumber());
+    EXPECT_CALL(*connection_, readDisable(_)).Times(AnyNumber());
 
     connect_timer_ = new NiceMock<Event::MockTimer>(&dispatcher_);
     EXPECT_CALL(dispatcher_, createClientConnection_(_, _, _, _)).WillOnce(Return(connection_));

test/integration/fake_upstream.cc

@@ -633,7 +633,7 @@ AssertionResult FakeUpstream::waitForRawConnection(FakeRawConnectionPtr& connect
 
   return runOnDispatcherThreadAndWait([&]() {
     absl::MutexLock lock(&lock_);
-    connection = std::make_unique<FakeRawConnection>(consumeConnection(), timeSystem());
+    connection = makeRawConnection(consumeConnection(), timeSystem());
     connection->initialize();
     // Skip enableHalfClose if the connection is already disconnected.
     if (connection->connected()) {

test/integration/fake_upstream.h

@@ -579,6 +579,12 @@ class FakeUpstream : Logger::Loggable<Logger::Id::testing>,
     return socket_->addressProvider().localAddress();
   }
 
+  virtual std::unique_ptr<FakeRawConnection>
+  makeRawConnection(SharedConnectionWrapper& shared_connection,
+                    Event::TestTimeSystem& time_system) {
+    return std::make_unique<FakeRawConnection>(shared_connection, time_system);
+  }
+
   // Wait for one of the upstreams to receive a connection
   ABSL_MUST_USE_RESULT
   static testing::AssertionResult

test/integration/tcp_proxy_integration_test.cc

@@ -32,6 +32,15 @@ using testing::NiceMock;
 
 namespace Envoy {
 
+std::vector<TcpProxyIntegrationTestParams> newPoolTestParams() {
+  std::vector<TcpProxyIntegrationTestParams> ret;
+
+  for (auto ip_version : TestEnvironment::getIpVersionsForTest()) {
+    ret.push_back(TcpProxyIntegrationTestParams{ip_version, false});
+  }
+  return ret;
+}
+
 std::vector<TcpProxyIntegrationTestParams> getProtocolTestParams() {
   std::vector<TcpProxyIntegrationTestParams> ret;
 
@@ -1297,4 +1306,179 @@ TEST_P(TcpProxySslIntegrationTest, UpstreamHalfClose) {
   ASSERT_TRUE(fake_upstream_connection_->waitForHalfClose());
 }
 
+// Integration test a Mysql upstream, where the upstream sends data immediately
+// after a connection is established.
+class FakeMysqlUpstream : public FakeUpstream {
+  using FakeUpstream::FakeUpstream;
+
+  bool createNetworkFilterChain(Network::Connection& connection,
+                                const std::vector<Network::FilterFactoryCb>& cb) override {
+    Buffer::OwnedImpl to_write("P");
+    connection.write(to_write, false);
+    return FakeUpstream::createNetworkFilterChain(connection, cb);
+  }
+};
+
+class MysqlIntegrationTest : public TcpProxyIntegrationTest {
+public:
+  void createUpstreams() override {
+    for (uint32_t i = 0; i < fake_upstreams_count_; ++i) {
+      Network::TransportSocketFactoryPtr factory =
+          upstream_tls_ ? createUpstreamTlsContext()
+                        : Network::Test::createRawBufferSocketFactory();
+      auto endpoint = upstream_address_fn_(i);
+      fake_upstreams_.emplace_back(
+          new FakeMysqlUpstream(std::move(factory), endpoint, upstreamConfig()));
+    }
+  }
+
+  absl::optional<uint64_t>
+  waitForNextUpstreamConnection(const std::vector<uint64_t>& upstream_indices,
+                                std::chrono::milliseconds connection_wait_timeout,
+                                FakeRawConnectionPtr& fake_upstream_connection) {
+    AssertionResult result = AssertionFailure();
+    int upstream_index = 0;
+    Event::TestTimeSystem::RealTimeBound bound(connection_wait_timeout);
+    // Loop over the upstreams until the call times out or an upstream request is received.
+    while (!result) {
+      upstream_index = upstream_index % upstream_indices.size();
+      result = fake_upstreams_[upstream_indices[upstream_index]]->waitForRawConnection(
+          fake_upstream_connection, std::chrono::milliseconds(5));
+      if (result) {
+        return upstream_index;
+      } else if (!bound.withinBound()) {
+        RELEASE_ASSERT(0, "Timed out waiting for new connection.");
+        break;
+      }
+      ++upstream_index;
+    }
+    RELEASE_ASSERT(result, result.message());
+    return {};
+  }
+
+  void globalPreconnect() {
+    config_helper_.addConfigModifier(
+        [&](envoy::config::bootstrap::v3::Bootstrap& bootstrap) -> void {
+          bootstrap.mutable_static_resources()
+              ->mutable_clusters(0)
+              ->mutable_preconnect_policy()
+              ->mutable_predictive_preconnect_ratio()
+              ->set_value(1.5);
+        });
+  }
+
+  void perUpstreamPreconnect() {
+    config_helper_.addConfigModifier(
+        [&](envoy::config::bootstrap::v3::Bootstrap& bootstrap) -> void {
+          bootstrap.mutable_static_resources()
+              ->mutable_clusters(0)
+              ->mutable_preconnect_policy()
+              ->mutable_per_upstream_preconnect_ratio()
+              ->set_value(1.5);
+        });
+  }
+
+  void testPreconnect();
+};
+
+// This just verifies that FakeMysqlUpstream works as advertised, and the early data
+// makes it to the client.
+TEST_P(MysqlIntegrationTest, UpstreamWritesFirst) {
+  globalPreconnect();
+  initialize();
+  IntegrationTcpClientPtr tcp_client = makeTcpConnection(lookupPort("tcp_proxy"));
+  FakeRawConnectionPtr fake_upstream_connection;
+  ASSERT_TRUE(fake_upstreams_[0]->waitForRawConnection(fake_upstream_connection));
+
+  tcp_client->waitForData("P", false);
+
+  ASSERT_TRUE(tcp_client->write("F"));
+  ASSERT_TRUE(fake_upstream_connection->waitForData(1));
+
+  ASSERT_TRUE(fake_upstream_connection->write("", true));
+  tcp_client->waitForHalfClose();
+  ASSERT_TRUE(tcp_client->write("", true));
+  ASSERT_TRUE(fake_upstream_connection->waitForHalfClose());
+  ASSERT_TRUE(fake_upstream_connection->waitForDisconnect());
+}
+
+// Make sure that with the connection read disabled, that disconnect detection
+// works.
+// Early close notification does not work for OSX
+#if !defined(__APPLE__)
+TEST_P(MysqlIntegrationTest, DisconnectDetected) {
+  // Switch to per-upstream preconnect.
+  perUpstreamPreconnect();
+  initialize();
+  IntegrationTcpClientPtr tcp_client = makeTcpConnection(lookupPort("tcp_proxy"));
+  FakeRawConnectionPtr fake_upstream_connection;
+  FakeRawConnectionPtr fake_upstream_connection1;
+  // The needed connection.
+  ASSERT_TRUE(fake_upstreams_[0]->waitForRawConnection(fake_upstream_connection));
+  // The prefetched connection.
+  ASSERT_TRUE(fake_upstreams_[0]->waitForRawConnection(fake_upstream_connection1));
+
+  // Close the prefetched connection.
+  ASSERT_TRUE(fake_upstream_connection1->close());
+  test_server_->waitForCounterGe("cluster.cluster_0.upstream_cx_destroy", 1);
+
+  tcp_client->close();
+}
+#endif
+
+void MysqlIntegrationTest::testPreconnect() {
+  globalPreconnect();
+  enableHalfClose(false);
+  config_helper_.addConfigModifier([this](envoy::config::bootstrap::v3::Bootstrap& bootstrap) {
+    auto* cluster = bootstrap.mutable_static_resources()->mutable_clusters(0);
+    auto* load_assignment = cluster->mutable_load_assignment();
+    load_assignment->clear_endpoints();
+    for (int i = 0; i < 5; ++i) {
+      auto locality = load_assignment->add_endpoints();
+      locality->add_lb_endpoints()->mutable_endpoint()->MergeFrom(
+          ConfigHelper::buildEndpoint(Network::Test::getLoopbackAddressString(version_)));
+    }
+  });
+
+  setUpstreamCount(5);
+  initialize();
+  int num_clients = 10;
+
+  std::vector<IntegrationTcpClientPtr> clients{10};
+  std::vector<FakeRawConnectionPtr> fake_connections{15};
+
+  int upstream_index = 0;
+  for (int i = 0; i < num_clients; ++i) {
+    // Start a new request.
+    clients[i] = makeTcpConnection(lookupPort("tcp_proxy"));
+    waitForNextUpstreamConnection(std::vector<uint64_t>({0, 1, 2, 3, 4}),
+                                  TestUtility::DefaultTimeout, fake_connections[upstream_index]);
+    ++upstream_index;
+
+    // For every other connection, an extra connection should be preconnected.
+    if (i % 2 == 0) {
+      waitForNextUpstreamConnection(std::vector<uint64_t>({0, 1, 2, 3, 4}),
+                                    TestUtility::DefaultTimeout, fake_connections[upstream_index]);
+      ++upstream_index;
+    }
+  }
+  EXPECT_EQ(0, test_server_->counter("cluster.cluster_0.upstream_cx_destroy")->value());
+  // Clean up.
+  for (int i = 0; i < num_clients; ++i) {
+    clients[i]->close();
+  }
+}
+
+TEST_P(MysqlIntegrationTest, Preconnect) { testPreconnect(); }
+
+TEST_P(MysqlIntegrationTest, PreconnectWithTls) {
+  upstream_tls_ = true;
+  setUpstreamProtocol(FakeHttpConnection::Type::HTTP1);
+  config_helper_.configureUpstreamTls();
+  testPreconnect();
+}
+
+INSTANTIATE_TEST_SUITE_P(TcpProxyIntegrationTestParams, MysqlIntegrationTest,
+                         testing::ValuesIn(newPoolTestParams()), protocolTestParamsToString);
+
 } // namespace Envoy