refactor: Add thread safety to connector registry (facebookincubator#16978)

Summary:
Pull Request resolved: facebookincubator#16978

Wrap the connector registry in folly::Synchronized to make all
operations thread-safe. Replace std::unordered_map with
folly::F14FastMap for the internal storage. Remove the deprecated
getAllConnectors() API.

Reviewed By: srsuryadev

Differential Revision: D98913848

fbshipit-source-id: b589ae001fef624682af13c0e294d3ca6e77cc2e

Author: mbasmanova

velox/connectors/Connector.cpp

@@ -20,34 +20,37 @@
 namespace facebook::velox::connector {
 
 bool registerConnector(std::shared_ptr<Connector> connector) {
-  bool ok = connectors().insert({connector->connectorId(), connector}).second;
-  VELOX_CHECK(
-      ok,
-      "Connector with ID is already registered: {}",
-      connector->connectorId());
-  return true;
+  return connectors().withWLock([&](auto& registry) {
+    bool ok = registry.insert({connector->connectorId(), connector}).second;
+    VELOX_CHECK(
+        ok,
+        "Connector with ID is already registered: {}",
+        connector->connectorId());
+    return true;
+  });
 }
 
 bool unregisterConnector(const std::string& connectorId) {
-  return connectors().erase(connectorId) == 1;
+  return connectors().withWLock(
+      [&](auto& registry) { return registry.erase(connectorId) == 1; });
 }
 
 std::shared_ptr<Connector> getConnector(const std::string& connectorId) {
-  auto it = connectors().find(connectorId);
-  VELOX_CHECK(
-      it != connectors().end(),
-      "Connector with ID is not registered: {}",
-      connectorId);
-  return it->second;
+  return connectors().withRLock(
+      [&](const auto& registry) -> std::shared_ptr<Connector> {
+        auto it = registry.find(connectorId);
+        VELOX_CHECK(
+            it != registry.end(),
+            "Connector with ID is not registered: {}",
+            connectorId);
+        return it->second;
+      });
 }
 
 bool hasConnector(const std::string& connectorId) {
-  return connectors().find(connectorId) != connectors().end();
-}
-
-const std::unordered_map<std::string, std::shared_ptr<Connector>>&
-getAllConnectors() {
-  return connectors();
+  return connectors().withRLock([&](const auto& registry) {
+    return registry.find(connectorId) != registry.end();
+  });
 }
 
 bool DataSink::Stats::empty() const {

velox/connectors/Connector.h

@@ -785,10 +785,4 @@ bool unregisterConnector(const std::string& connectorId);
 /// exist.
 std::shared_ptr<Connector> getConnector(const std::string& connectorId);
 
-/// Returns a map of all (connectorId -> connector) pairs currently
-/// registered.
-[[deprecated("Use ConnectorRegistry methods instead.")]] const std::
-    unordered_map<std::string, std::shared_ptr<Connector>>&
-    getAllConnectors();
-
 } // namespace facebook::velox::connector

velox/connectors/ConnectorRegistry.cpp

@@ -22,22 +22,30 @@ namespace facebook::velox::connector {
 // static
 std::shared_ptr<Connector> ConnectorRegistry::tryGet(
     const std::string& connectorId) {
-  auto it = connectors().find(connectorId);
-  if (it != connectors().end()) {
-    return it->second;
-  }
-  return nullptr;
+  return connectors().withRLock(
+      [&](const auto& registry) -> std::shared_ptr<Connector> {
+        auto it = registry.find(connectorId);
+        if (it != registry.end()) {
+          return it->second;
+        }
+        return nullptr;
+      });
 }
 
 // static
 void ConnectorRegistry::unregisterAll() {
-  connectors().clear();
+  folly::F14FastMap<std::string, std::shared_ptr<Connector>> entries;
+  connectors().withWLock([&](auto& registry) { entries.swap(registry); });
 }
 
 // static
-const std::unordered_map<std::string, std::shared_ptr<Connector>>&
-ConnectorRegistry::all() {
-  return connectors();
+void ConnectorRegistry::forEach(
+    std::function<void(const std::shared_ptr<Connector>&)> func) {
+  connectors().withRLock([&](const auto& registry) {
+    for (const auto& [_, connector] : registry) {
+      func(connector);
+    }
+  });
 }
 
 } // namespace facebook::velox::connector

velox/connectors/ConnectorRegistry.h

@@ -21,6 +21,7 @@
 namespace facebook::velox::connector {
 
 /// Provides static methods for connector registry operations.
+/// All methods are thread-safe.
 class ConnectorRegistry {
  public:
   /// Return the connector with the specified ID, or nullptr if not registered.
@@ -30,21 +31,21 @@ class ConnectorRegistry {
   template <typename T>
   static std::vector<std::shared_ptr<T>> findAll() {
     std::vector<std::shared_ptr<T>> result;
-    for (const auto& [_, connector] : all()) {
+    forEach([&](const std::shared_ptr<Connector>& connector) {
       if (auto casted = std::dynamic_pointer_cast<T>(connector)) {
         result.push_back(std::move(casted));
       }
-    }
+    });
     return result;
   }
 
   /// Unregister all connectors.
   static void unregisterAll();
 
  private:
-  // Return a reference to the internal connector map.
-  static const std::unordered_map<std::string, std::shared_ptr<Connector>>&
-  all();
+  // Invoke 'func' on each registered connector under the read lock.
+  static void forEach(
+      std::function<void(const std::shared_ptr<Connector>&)> func);
 };
 
 } // namespace facebook::velox::connector

velox/connectors/ConnectorRegistryInternal.h

@@ -18,17 +18,21 @@
 
 #include <memory>
 #include <string>
-#include <unordered_map>
+
+#include "folly/Synchronized.h"
+#include "folly/container/F14Map.h"
 
 namespace facebook::velox::connector {
 
 class Connector;
 
 // Internal helper shared by Connector.cpp and ConnectorRegistry.cpp.
 // Not part of the public API. Do not include from outside velox/connectors/.
-inline std::unordered_map<std::string, std::shared_ptr<Connector>>&
-connectors() {
-  static std::unordered_map<std::string, std::shared_ptr<Connector>> instance;
+using ConnectorMap = folly::Synchronized<
+    folly::F14FastMap<std::string, std::shared_ptr<Connector>>>;
+
+inline ConnectorMap& connectors() {
+  static ConnectorMap instance;
   return instance;
 }