Refactor schema registries into singletons

Summary: Makes testing the next diff and reusing the logic for custom registries easier. Reviewed By: yoney Differential Revision: D67719826 fbshipit-source-id: e50b2375e7f10a496a0fbc7f2f245c820cec0e69
facebook · Dec 31, 2024 · 7553857 · 7553857
1 parent 9090004
commit 7553857
Show file tree

Hide file tree

Showing 7 changed files with 46 additions and 34 deletions.
diff --git a/thrift/compiler/generate/templates/cpp2/module_sinit.cpp.mustache b/thrift/compiler/generate/templates/cpp2/module_sinit.cpp.mustache
@@ -56,7 +56,7 @@ struct StaticInit {
     {{/program:any?}}
 
     {{#program:has_schema?}}
-    ::apache::thrift::BaseSchemaRegistry::registerSchema("{{program:schema_name}}", {{program:name}}_constants::{{program:schema_name}}(), "{{program:autogen_path}}");
+    ::apache::thrift::BaseSchemaRegistry::get().registerSchema("{{program:schema_name}}", {{program:name}}_constants::{{program:schema_name}}(), "{{program:autogen_path}}");
     {{/program:has_schema?}}
   }
 };

diff --git a/thrift/compiler/test/fixtures/service-schema/out/cpp2_sinit/gen-cpp2/module_sinit.cpp b/thrift/compiler/test/fixtures/service-schema/out/cpp2_sinit/gen-cpp2/module_sinit.cpp
@@ -24,7 +24,7 @@ namespace {
 struct StaticInit {
   StaticInit() {
 
-    ::apache::thrift::BaseSchemaRegistry::registerSchema("_fbthrift_schema_b747839c13cb3aa5", module_constants::_fbthrift_schema_b747839c13cb3aa5(), "thrift/compiler/test/fixtures/service-schema/src/module.thrift");
+    ::apache::thrift::BaseSchemaRegistry::get().registerSchema("_fbthrift_schema_b747839c13cb3aa5", module_constants::_fbthrift_schema_b747839c13cb3aa5(), "thrift/compiler/test/fixtures/service-schema/src/module.thrift");
   }
 };
 

diff --git a/thrift/lib/cpp2/runtime/BaseSchemaRegistry.cpp b/thrift/lib/cpp2/runtime/BaseSchemaRegistry.cpp
@@ -22,10 +22,10 @@ namespace apache::thrift {
 
 void BaseSchemaRegistry::registerSchema(
     std::string_view name, std::string_view data, std::string_view path) {
-  if (accessed()) {
+  if (accessed_) {
     throw std::runtime_error("Schemas accessed before registration complete.");
   }
-  if (auto it = getRawSchemas().find(name); it != getRawSchemas().end()) {
+  if (auto it = rawSchemas_.find(name); it != rawSchemas_.end()) {
     if (it->second.path != path) { // Needed to support dynamic linking
       throw std::runtime_error(fmt::format(
           "Checksum collision between {} and {}. Make any change to either file's content (e.g.. whitespace/comment) to fix it.",
@@ -34,20 +34,12 @@ void BaseSchemaRegistry::registerSchema(
     }
     return;
   }
-  getRawSchemas()[name] = {data, path};
+  rawSchemas_[name] = {data, path};
 }
 
-folly::F14FastMap<std::string_view, BaseSchemaRegistry::RawSchema>&
-BaseSchemaRegistry::getRawSchemas() {
-  static folly::Indestructible<
-      folly::F14FastMap<std::string_view, BaseSchemaRegistry::RawSchema>>
-      schemas;
-  return *schemas;
-}
-
-bool& BaseSchemaRegistry::accessed() {
-  static folly::Indestructible<bool> accessed = false;
-  return *accessed;
+BaseSchemaRegistry& BaseSchemaRegistry::get() {
+  static folly::Indestructible<BaseSchemaRegistry> self;
+  return *self;
 }
 
 } // namespace apache::thrift
diff --git a/thrift/lib/cpp2/runtime/BaseSchemaRegistry.h b/thrift/lib/cpp2/runtime/BaseSchemaRegistry.h
@@ -24,16 +24,19 @@ namespace apache::thrift {
 
 class BaseSchemaRegistry {
  public:
-  static void registerSchema(
+  // Access the global registry.
+  static BaseSchemaRegistry& get();
+
+  void registerSchema(
       std::string_view name, std::string_view data, std::string_view path);
 
  private:
   struct RawSchema {
     std::string_view data;
     std::string_view path;
   };
-  static folly::F14FastMap<std::string_view, RawSchema>& getRawSchemas();
-  static bool& accessed();
+  folly::F14FastMap<std::string_view, RawSchema> rawSchemas_;
+  bool accessed_;
 
   friend class SchemaRegistry;
 };

diff --git a/thrift/lib/cpp2/runtime/SchemaRegistry.cpp b/thrift/lib/cpp2/runtime/SchemaRegistry.cpp
@@ -19,22 +19,26 @@
 #include <folly/Indestructible.h>
 #include <folly/compression/Compression.h>
 #include <thrift/lib/cpp2/protocol/Serializer.h>
-#include <thrift/lib/cpp2/runtime/BaseSchemaRegistry.h>
 
 namespace apache::thrift {
 
+SchemaRegistry& SchemaRegistry::get() {
+  static folly::Indestructible<SchemaRegistry> self(BaseSchemaRegistry::get());
+  return *self;
+}
+
 const type::Schema& SchemaRegistry::getMergedSchema() {
-  static const folly::Indestructible<type::Schema> merged = [&] {
-    BaseSchemaRegistry::accessed() = true;
+  folly::call_once(mergedFlag_, [this]() {
+    base_.accessed_ = true;
     std::vector<std::string_view> schemas;
-    schemas.reserve(BaseSchemaRegistry::getRawSchemas().size());
-    for (auto& [name, data] : BaseSchemaRegistry::getRawSchemas()) {
+    schemas.reserve(base_.rawSchemas_.size());
+    for (auto& [name, data] : base_.rawSchemas_) {
       schemas.push_back(data.data);
     }
-    return mergeSchemas(folly::range(schemas));
-  }();
+    merged_ = mergeSchemas(folly::range(schemas));
+  });
 
-  return *merged;
+  return merged_;
 }
 
 namespace {

diff --git a/thrift/lib/cpp2/runtime/SchemaRegistry.h b/thrift/lib/cpp2/runtime/SchemaRegistry.h
@@ -16,16 +16,31 @@
 
 #pragma once
 
+#include <folly/synchronization/CallOnce.h>
+#include <thrift/lib/cpp2/runtime/BaseSchemaRegistry.h>
 #include <thrift/lib/thrift/gen-cpp2/schema_types.h>
 
 namespace apache::thrift {
 
 class SchemaRegistry {
  public:
-  static const type::Schema& getMergedSchema();
+  // Access the global registry.
+  static SchemaRegistry& get();
+
+  // Access all data registered
+  const type::Schema& getMergedSchema();
+
+  // Helpers for working with schemas.
   static type::Schema mergeSchemas(
       folly::Range<const std::string_view*> schemas);
   static type::Schema mergeSchemas(std::vector<type::Schema>&& schemas);
+
+  explicit SchemaRegistry(BaseSchemaRegistry& base) : base_(base) {}
+
+ private:
+  BaseSchemaRegistry& base_;
+  type::Schema merged_;
+  folly::once_flag mergedFlag_;
 };
 
 } // namespace apache::thrift
diff --git a/thrift/test/SchemaTest.cpp b/thrift/test/SchemaTest.cpp
@@ -14,8 +14,6 @@
  * limitations under the License.
  */
 
-#include <functional>
-#include <string>
 #include <vector>
 #include <folly/portability/GTest.h>
 #include <thrift/lib/cpp2/runtime/SchemaRegistry.h>
@@ -31,7 +29,7 @@
 using namespace apache::thrift;
 
 TEST(SchemaTest, not_linked) {
-  const auto& schema = SchemaRegistry::getMergedSchema();
+  const auto& schema = SchemaRegistry::get().getMergedSchema();
   for (const auto& program : *schema.programs()) {
     EXPECT_NE(
         program.path(), "thrift/test/TopologicallySortObjectsTest.thrift");
@@ -42,7 +40,7 @@ TEST(SchemaTest, not_linked) {
 }
 
 TEST(SchemaTest, not_linked_but_included) {
-  const auto& schema = SchemaRegistry::getMergedSchema();
+  const auto& schema = SchemaRegistry::get().getMergedSchema();
   for (const auto& program : *schema.programs()) {
     if (program.path() != "thrift/annotation/thrift.thrift") {
       continue;
@@ -58,7 +56,7 @@ TEST(SchemaTest, not_linked_but_included) {
 
 TEST(SchemaTest, linked) {
   bool found = false;
-  const auto& schema = SchemaRegistry::getMergedSchema();
+  const auto& schema = SchemaRegistry::get().getMergedSchema();
   for (const auto& program : *schema.programs()) {
     if (program.path() == "thrift/test/schema.thrift") {
       found = true;
@@ -99,7 +97,7 @@ TEST(SchemaTest, static_schema) {
   }
   ASSERT_TRUE(static_program);
 
-  const auto& dynamic_schema = SchemaRegistry::getMergedSchema();
+  const auto& dynamic_schema = SchemaRegistry::get().getMergedSchema();
   const type::Program* dynamic_program = nullptr;
   for (const auto& program : *dynamic_schema.programs()) {
     if (program.path() == "thrift/test/schema.thrift") {