test: FakeUpstream threading fixes

test/extensions/filters/network/thrift_proxy/integration_test.cc

@@ -413,8 +413,7 @@ TEST_P(ThriftConnManagerIntegrationTest, OnewayEarlyClosePartialRequest) {
   ASSERT_TRUE(tcp_client->write(partial_request));
   tcp_client->close();
 
-  FakeRawConnectionPtr fake_upstream_connection;
-  ASSERT_TRUE(expected_upstream->waitForRawConnection(fake_upstream_connection));
+  ASSERT_TRUE(expected_upstream->waitForAndConsumeDisconnectedConnection());
 
   test_server_->waitForCounterGe("thrift.thrift_stats.cx_destroy_remote_with_active_rq", 1);
 

test/integration/fake_upstream.cc

@@ -20,6 +20,7 @@
 #include "test/test_common/utility.h"
 
 #include "absl/strings/str_cat.h"
+#include "absl/synchronization/notification.h"
 
 using namespace std::chrono_literals;
 
@@ -354,14 +355,24 @@ AssertionResult FakeConnectionBase::close(std::chrono::milliseconds timeout) {
 }
 
 AssertionResult FakeConnectionBase::readDisable(bool disable, std::chrono::milliseconds timeout) {
-  return shared_connection_.executeOnDispatcher(
-      [disable](Network::Connection& connection) { connection.readDisable(disable); }, timeout);
+  if (shared_connection_.connection().dispatcher().isThreadSafe()) {
+    shared_connection_.connection().readDisable(disable);
+    return AssertionSuccess();
+  } else {
+    return shared_connection_.executeOnDispatcher(
+        [disable](Network::Connection& connection) { connection.readDisable(disable); }, timeout);
+  }
 }
 
 AssertionResult FakeConnectionBase::enableHalfClose(bool enable,
                                                     std::chrono::milliseconds timeout) {
-  return shared_connection_.executeOnDispatcher(
-      [enable](Network::Connection& connection) { connection.enableHalfClose(enable); }, timeout);
+  if (shared_connection_.connection().dispatcher().isThreadSafe()) {
+    shared_connection_.connection().enableHalfClose(enable);
+    return AssertionSuccess();
+  } else {
+    return shared_connection_.executeOnDispatcher(
+        [enable](Network::Connection& connection) { connection.enableHalfClose(enable); }, timeout);
+  }
 }
 
 Http::RequestDecoder& FakeHttpConnection::newStream(Http::ResponseEncoder& encoder, bool) {
@@ -552,16 +563,28 @@ AssertionResult FakeUpstream::waitForHttpConnection(
             client_dispatcher, timeout)) {
       return AssertionFailure() << "Timed out waiting for new connection.";
     }
-
-    connection = std::make_unique<FakeHttpConnection>(
-        *this, consumeConnection(), http_type_, time_system_, max_request_headers_kb,
-        max_request_headers_count, headers_with_underscores_action);
   }
-  VERIFY_ASSERTION(connection->initialize());
-  if (read_disable_on_new_connection_) {
-    VERIFY_ASSERTION(connection->readDisable(false));
-  }
-  return AssertionSuccess();
+
+  AssertionResult result = AssertionSuccess();
+  absl::Notification done;
+  ASSERT(!dispatcher_->isThreadSafe());
+  dispatcher_->post([&]() {
+    {
+      absl::MutexLock lock(&lock_);
+      connection = std::make_unique<FakeHttpConnection>(
+          *this, consumeConnection(), http_type_, time_system_, max_request_headers_kb,
+          max_request_headers_count, headers_with_underscores_action);
+      result = connection->initialize();
+      if (result == AssertionSuccess()) {
+        if (read_disable_on_new_connection_) {
+          result = connection->readDisable(false);
+        }
+      }
+    }
+    done.Notify();
+  });
+  done.WaitForNotification();
+  return result;
 }
 
 AssertionResult
@@ -585,14 +608,27 @@ FakeUpstream::waitForHttpConnection(Event::Dispatcher& client_dispatcher,
                 client_dispatcher, 5ms)) {
           continue;
         }
-        connection = std::make_unique<FakeHttpConnection>(
-            upstream, upstream.consumeConnection(), upstream.http_type_, upstream.timeSystem(),
-            Http::DEFAULT_MAX_REQUEST_HEADERS_KB, Http::DEFAULT_MAX_HEADERS_COUNT,
-            envoy::config::core::v3::HttpProtocolOptions::ALLOW);
       }
-      VERIFY_ASSERTION(connection->initialize());
-      VERIFY_ASSERTION(connection->readDisable(false));
-      return AssertionSuccess();
+
+      AssertionResult result = AssertionSuccess();
+      absl::Notification done;
+      ASSERT(!upstream.dispatcher_->isThreadSafe());
+      upstream.dispatcher_->post([&]() {
+        {
+          absl::MutexLock lock(&upstream.lock_);
+          connection = std::make_unique<FakeHttpConnection>(
+              upstream, upstream.consumeConnection(), upstream.http_type_, upstream.timeSystem(),
+              Http::DEFAULT_MAX_REQUEST_HEADERS_KB, Http::DEFAULT_MAX_HEADERS_COUNT,
+              envoy::config::core::v3::HttpProtocolOptions::ALLOW);
+          result = connection->initialize();
+          if (result == AssertionSuccess()) {
+            result = connection->readDisable(false);
+          }
+        }
+        done.Notify();
+      });
+      done.WaitForNotification();
+      return result;
     }
   }
   return AssertionFailure() << "Timed out waiting for HTTP connection.";
@@ -610,11 +646,42 @@ AssertionResult FakeUpstream::waitForRawConnection(FakeRawConnectionPtr& connect
     if (!time_system_.waitFor(lock_, absl::Condition(&reached), timeout)) {
       return AssertionFailure() << "Timed out waiting for raw connection";
     }
-    connection = std::make_unique<FakeRawConnection>(consumeConnection(), timeSystem());
   }
-  VERIFY_ASSERTION(connection->initialize());
-  VERIFY_ASSERTION(connection->readDisable(false));
-  VERIFY_ASSERTION(connection->enableHalfClose(enable_half_close_));
+
+  AssertionResult result = AssertionSuccess();
+  absl::Notification done;
+  ASSERT(!dispatcher_->isThreadSafe());
+  dispatcher_->post([&]() {
+    {
+      absl::MutexLock lock(&lock_);
+      auto local_connection =
+          std::make_unique<FakeRawConnection>(consumeConnection(), timeSystem());
+      result = local_connection->initialize();
+      if (result == AssertionSuccess()) {
+        result = local_connection->readDisable(false);
+        if (result == AssertionSuccess()) {
+          result = local_connection->enableHalfClose(enable_half_close_);
+        }
+      }
+      connection = std::move(local_connection);
+    }
+    done.Notify();
+  });
+  done.WaitForNotification();
+  return result;
+}
+
+testing::AssertionResult
+FakeUpstream::waitForAndConsumeDisconnectedConnection(std::chrono::milliseconds timeout) {
+  absl::MutexLock lock(&lock_);
+  const auto reached = [this]() ABSL_EXCLUSIVE_LOCKS_REQUIRED(lock_) {
+    return !new_connections_.empty() && !new_connections_.front()->connected();
+  };
+
+  if (!time_system_.waitFor(lock_, absl::Condition(&reached), timeout)) {
+    return AssertionFailure() << "Timed out waiting for raw connection";
+  }
+  consumeConnection();
   return AssertionSuccess();
 }
 
@@ -683,14 +750,19 @@ FakeRawConnection::~FakeRawConnection() {
 }
 
 testing::AssertionResult FakeRawConnection::initialize() {
-  auto filter = Network::ReadFilterSharedPtr{new ReadFilter(*this)};
+  ASSERT(shared_connection_.connection().dispatcher().isThreadSafe());
+  Network::ReadFilterSharedPtr filter{new ReadFilter(*this)};
   read_filter_ = filter;
-  testing::AssertionResult result = shared_connection_.executeOnDispatcher(
-      [filter = std::move(filter)](Network::Connection& connection) {
-        connection.addReadFilter(filter);
-      });
-  if (!result) {
-    return result;
+  if (shared_connection_.connection().dispatcher().isThreadSafe()) {
+    shared_connection_.connection().addReadFilter(filter);
+  } else {
+    testing::AssertionResult result = shared_connection_.executeOnDispatcher(
+        [filter = std::move(filter)](Network::Connection& connection) {
+          connection.addReadFilter(filter);
+        });
+    if (!result) {
+      return result;
+    }
   }
   return FakeConnectionBase::initialize();
 }

test/integration/fake_upstream.h

@@ -257,17 +257,6 @@ class SharedConnectionWrapper : public Network::ConnectionCallbacks,
     connection_.addConnectionCallbacks(*this);
   }
 
-  Common::CallbackHandle* addDisconnectCallback(DisconnectCallback callback) {
-    absl::MutexLock lock(&lock_);
-    return disconnect_callback_manager_.add(callback);
-  }
-
-  // Avoid directly removing by caller, since CallbackManager is not thread safe.
-  void removeDisconnectCallback(Common::CallbackHandle* handle) {
-    absl::MutexLock lock(&lock_);
-    handle->remove();
-  }
-
   // Network::ConnectionCallbacks
   void onEvent(Network::ConnectionEvent event) override {
     // Throughout this entire function, we know that the connection_ cannot disappear, since this
@@ -278,7 +267,6 @@ class SharedConnectionWrapper : public Network::ConnectionCallbacks,
     if (event == Network::ConnectionEvent::RemoteClose ||
         event == Network::ConnectionEvent::LocalClose) {
       disconnected_ = true;
-      disconnect_callback_manager_.runCallbacks();
     }
   }
 
@@ -315,6 +303,8 @@ class SharedConnectionWrapper : public Network::ConnectionCallbacks,
     if (disconnected_) {
       return testing::AssertionSuccess();
     }
+    ASSERT(!connection_.dispatcher().isThreadSafe(),
+           "deadlock: executeOnDispatcher called from dispatcher thread.");
     bool callback_ready_event = false;
     bool unexpected_disconnect = false;
     connection_.dispatcher().post(
@@ -345,13 +335,13 @@ class SharedConnectionWrapper : public Network::ConnectionCallbacks,
 
   void setParented() {
     absl::MutexLock lock(&lock_);
+    ASSERT(!parented_);
     parented_ = true;
   }
 
 private:
   Network::Connection& connection_;
   absl::Mutex lock_;
-  Common::CallbackManager<> disconnect_callback_manager_ ABSL_GUARDED_BY(lock_);
   bool parented_ ABSL_GUARDED_BY(lock_){};
   bool disconnected_ ABSL_GUARDED_BY(lock_){};
 };
@@ -579,6 +569,11 @@ class FakeUpstream : Logger::Loggable<Logger::Id::testing>,
   testing::AssertionResult
   waitForRawConnection(FakeRawConnectionPtr& connection,
                        std::chrono::milliseconds timeout = TestUtility::DefaultTimeout);
+
+  ABSL_MUST_USE_RESULT
+  testing::AssertionResult waitForAndConsumeDisconnectedConnection(
+      std::chrono::milliseconds timeout = TestUtility::DefaultTimeout);
+
   Network::Address::InstanceConstSharedPtr localAddress() const { return socket_->localAddress(); }
 
   // Wait for one of the upstreams to receive a connection

test/integration/sds_dynamic_integration_test.cc

@@ -654,9 +654,7 @@ TEST_P(SdsDynamicUpstreamIntegrationTest, WrongSecretFirst) {
   EXPECT_EQ("503", response->headers().getStatusValue());
 
   // To flush out the reset connection from the first request in upstream.
-  FakeRawConnectionPtr fake_upstream_connection;
-  ASSERT_TRUE(fake_upstreams_[0]->waitForRawConnection(fake_upstream_connection));
-  ASSERT_TRUE(fake_upstream_connection->waitForDisconnect());
+  ASSERT_TRUE(fake_upstreams_[0]->waitForAndConsumeDisconnectedConnection());
 
   // Failure
   EXPECT_EQ(0, test_server_->counter("sds.client_cert.update_success")->value());