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,28 @@ 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);
+  // Do the work inline if called from the dispatcher thread, executeOnDispatcher can only be called
+  // from outside the dispatcher thread.
+  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);
+  // Do the work inline if called from the dispatcher thread, executeOnDispatcher can only be called
+  // from outside the dispatcher thread.
+  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 +567,19 @@ AssertionResult FakeUpstream::waitForHttpConnection(
             client_dispatcher, timeout)) {
       return AssertionFailure() << "Timed out waiting for new connection.";
     }
+  }
 
+  return runOnDispatcherThreadAndWait([&]() {
+    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);
-  }
-  VERIFY_ASSERTION(connection->initialize());
-  if (read_disable_on_new_connection_) {
-    VERIFY_ASSERTION(connection->readDisable(false));
-  }
-  return AssertionSuccess();
+    VERIFY_ASSERTION(connection->initialize());
+    if (read_disable_on_new_connection_) {
+      VERIFY_ASSERTION(connection->readDisable(false));
+    }
+    return AssertionSuccess();
+  });
 }
 
 AssertionResult
@@ -585,14 +603,18 @@ FakeUpstream::waitForHttpConnection(Event::Dispatcher& client_dispatcher,
                 client_dispatcher, 5ms)) {
           continue;
         }
+      }
+
+      return upstream.runOnDispatcherThreadAndWait([&]() {
+        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);
-      }
-      VERIFY_ASSERTION(connection->initialize());
-      VERIFY_ASSERTION(connection->readDisable(false));
-      return AssertionSuccess();
+        VERIFY_ASSERTION(connection->initialize());
+        VERIFY_ASSERTION(connection->readDisable(false));
+        return AssertionSuccess();
+      });
     }
   }
   return AssertionFailure() << "Timed out waiting for HTTP connection.";
@@ -610,11 +632,29 @@ AssertionResult FakeUpstream::waitForRawConnection(FakeRawConnectionPtr& connect
     if (!time_system_.waitFor(lock_, absl::Condition(&reached), timeout)) {
       return AssertionFailure() << "Timed out waiting for raw connection";
     }
+  }
+
+  return runOnDispatcherThreadAndWait([&]() {
+    absl::MutexLock lock(&lock_);
     connection = std::make_unique<FakeRawConnection>(consumeConnection(), timeSystem());
+    VERIFY_ASSERTION(connection->initialize());
+    VERIFY_ASSERTION(connection->readDisable(false));
+    VERIFY_ASSERTION(connection->enableHalfClose(enable_half_close_));
+    return AssertionSuccess();
+  });
+}
+
+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";
   }
-  VERIFY_ASSERTION(connection->initialize());
-  VERIFY_ASSERTION(connection->readDisable(false));
-  VERIFY_ASSERTION(connection->enableHalfClose(enable_half_close_));
+  consumeConnection();
   return AssertionSuccess();
 }
 
@@ -647,6 +687,18 @@ void FakeUpstream::onRecvDatagram(Network::UdpRecvData& data) {
   received_datagrams_.emplace_back(std::move(data));
 }
 
+AssertionResult FakeUpstream::runOnDispatcherThreadAndWait(std::function<AssertionResult()> cb) {
+  AssertionResult result = AssertionSuccess();
+  absl::Notification done;
+  ASSERT(!dispatcher_->isThreadSafe());
+  dispatcher_->post([&]() {
+    result = cb();
+    done.Notify();
+  });
+  done.WaitForNotification();
+  return result;
+}
+
 void FakeUpstream::sendUdpDatagram(const std::string& buffer,
                                    const Network::Address::InstanceConstSharedPtr& peer) {
   dispatcher_->post([this, buffer, peer] {
@@ -683,14 +735,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
@@ -734,6 +729,7 @@ class FakeUpstream : Logger::Loggable<Logger::Id::testing>,
   void threadRoutine();
   SharedConnectionWrapper& consumeConnection() ABSL_EXCLUSIVE_LOCKS_REQUIRED(lock_);
   void onRecvDatagram(Network::UdpRecvData& data);
+  AssertionResult runOnDispatcherThreadAndWait(std::function<AssertionResult()> cb);
 
   Network::SocketSharedPtr socket_;
   Network::ListenSocketFactorySharedPtr socket_factory_;

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());