Implement schema merging

Summary:
Merges deserialized schemas into one before storing.
This lets us benefit from value deduplication and means we don't need to build a separate index of definitions and values, as well as reducing the number of allocations needed.

Reviewed By: thedavekwon

Differential Revision: D53835988

fbshipit-source-id: 430222f41878c1a104108b5a819c37eb4eda208f

Author: iahs

thrift/lib/cpp2/runtime/SchemaRegistry.cpp

@@ -38,20 +38,44 @@ void SchemaRegistry::registerSchema(
 }
 
 #ifdef FBTHRIFT_HAS_SCHEMA
-folly::F14FastMap<type::ProgramId, type::Schema>& SchemaRegistry::getSchemas() {
-  static folly::Indestructible<folly::F14FastMap<type::ProgramId, type::Schema>>
-      schemas = [&] {
-        folly::F14FastMap<type::ProgramId, type::Schema> schemaMap;
-        accessed() = true;
-        for (auto& [name, data] : getRawSchemas()) {
-          auto schema = CompactSerializer::deserialize<type::Schema>(data.data);
-          auto id = *schema.programs()[0].id();
-          schemaMap[id] = std::move(schema);
+const type::Schema& SchemaRegistry::getMergedSchema() {
+  static const folly::Indestructible<type::Schema> merged = [&] {
+    accessed() = true;
+    type::Schema mergedSchema;
+    std::unordered_set<type::ProgramId> includedPrograms;
+
+    for (auto& [name, data] : getRawSchemas()) {
+      auto schema = CompactSerializer::deserialize<type::Schema>(data.data);
+
+      for (auto& program : *schema.programs()) {
+        auto id = *program.id();
+        if (!includedPrograms.insert(id).second) {
+          // We checked during registration that the program ids are unique,
+          // so this file was already included by another program bundle.
+          continue;
         }
-        return schemaMap;
-      }();
+        mergedSchema.programs()->push_back(std::move(program));
+      }
 
-  return *schemas;
+      mergedSchema.valuesMap()->insert(
+          std::make_move_iterator(schema.valuesMap()->begin()),
+          std::make_move_iterator(schema.valuesMap()->end()));
+      // This deduplicates common values.
+
+      auto ndefs = mergedSchema.definitionsMap()->size();
+      mergedSchema.definitionsMap()->insert(
+          std::make_move_iterator(schema.definitionsMap()->begin()),
+          std::make_move_iterator(schema.definitionsMap()->end()));
+      if (mergedSchema.definitionsMap()->size() !=
+          ndefs + schema.definitionsMap()->size()) {
+        throw std::runtime_error("DefinitionKey collision");
+      }
+    }
+
+    return mergedSchema;
+  }();
+
+  return *merged;
 }
 #endif
 

thrift/lib/cpp2/runtime/SchemaRegistry.h

@@ -42,14 +42,7 @@ class SchemaRegistry {
 
 #ifdef FBTHRIFT_HAS_SCHEMA
  public:
-  struct Iter {
-    auto begin() const { return SchemaRegistry::getSchemas().begin(); }
-    auto end() const { return SchemaRegistry::getSchemas().end(); }
-  };
-  static Iter iter() { return {}; }
-
- private:
-  static folly::F14FastMap<type::ProgramId, type::Schema>& getSchemas();
+  static const type::Schema& getMergedSchema();
 #endif
 };
 

thrift/test/SchemaTest.cpp

@@ -22,13 +22,31 @@
 
 // These deps pull in the corresponding schemas when enabled.
 #include <thrift/annotation/gen-cpp2/thrift_types.h>
+#include <thrift/test/gen-cpp2/TopologicallySortObjectsTest_types.h>
 #include <thrift/test/gen-cpp2/schema_types.h>
 
 using namespace apache::thrift;
 
 TEST(SchemaTest, not_linked) {
-  for (const auto& [id, schema] : SchemaRegistry::iter()) {
-    EXPECT_NE(schema.programs()[0].path(), "thrift/annotation/thrift.thrift");
+  const auto& schema = SchemaRegistry::getMergedSchema();
+  for (const auto& program : *schema.programs()) {
+    EXPECT_NE(
+        program.path(), "thrift/test/TopologicallySortObjectsTest.thrift");
+  }
+
+  // Use the types target
+  (void)apache::thrift::test::IncompleteMap{};
+}
+
+TEST(SchemaTest, not_linked_but_included) {
+  const auto& schema = SchemaRegistry::getMergedSchema();
+  for (const auto& program : *schema.programs()) {
+    if (program.path() != "thrift/annotation/thrift.thrift") {
+      continue;
+    }
+    for (const auto& def : *program.definitionKeys()) {
+      EXPECT_FALSE(schema.definitionsMap()->contains(def));
+    }
   }
 
   // Use the types target
@@ -37,8 +55,9 @@ TEST(SchemaTest, not_linked) {
 
 TEST(SchemaTest, linked) {
   bool found = false;
-  for (const auto& [id, schema] : SchemaRegistry::iter()) {
-    if (schema.programs()[0].path() == "thrift/test/schema.thrift") {
+  const auto& schema = SchemaRegistry::getMergedSchema();
+  for (const auto& program : *schema.programs()) {
+    if (program.path() == "thrift/test/schema.thrift") {
       found = true;
     } else {
       continue;
@@ -47,15 +66,14 @@ TEST(SchemaTest, linked) {
     // Includes definitions from root program
     EXPECT_EQ(
         schema.definitionsMap()
-            ->at(schema.programs()[0].definitionKeys()[0])
+            ->at(program.definitionKeys()[0])
             .structDef_ref()
             ->name(),
         "Empty");
 
     // Only includes definitions from root program
     EXPECT_EQ(
-        schema.definitionsMap()->size(),
-        schema.programs()[0].definitionKeys()->size());
+        schema.definitionsMap()->size(), program.definitionKeys()->size());
 
     // Transitive includes are listed
     EXPECT_GT(schema.programs()->size(), 4);