Add incremental resolver for SyntaxGraph

Summary:
Enables building up SyntaxGraph as subgraphs are accessed instead of paying the full cost upfront.
Internally synchronizes accesses to the data structures to provide thread safety.

Reviewed By: praihan

Differential Revision: D72751923

fbshipit-source-id: 5d6fda62969f3fc5a79cf65f1982efa91bf72772

Author: iahs

third-party/thrift/src/thrift/lib/cpp2/runtime/SchemaRegistry.cpp

@@ -24,6 +24,13 @@
 
 namespace apache::thrift {
 
+SchemaRegistry::SchemaRegistry(BaseSchemaRegistry& base) : base_(base) {
+  auto resolver = std::make_unique<schema::detail::IncrementalResolver>();
+  resolver_ = resolver.get();
+  syntaxGraph_ = std::make_unique<schema::SyntaxGraph>(std::move(resolver));
+}
+SchemaRegistry::~SchemaRegistry() = default;
+
 SchemaRegistry& SchemaRegistry::get() {
   static folly::Indestructible<SchemaRegistry> self(BaseSchemaRegistry::get());
   return *self;

third-party/thrift/src/thrift/lib/cpp2/runtime/SchemaRegistry.h

@@ -23,6 +23,9 @@
 #include <unordered_set>
 #include <folly/synchronization/RelaxedAtomic.h>
 #include <thrift/lib/cpp2/runtime/BaseSchemaRegistry.h>
+#include <thrift/lib/cpp2/schema/SyntaxGraph.h>
+#include <thrift/lib/cpp2/schema/detail/SchemaBackedResolver.h>
+#include <thrift/lib/thrift/gen-cpp2/schema_types.h>
 
 namespace apache::thrift {
 
@@ -36,13 +39,25 @@ class SchemaRegistry {
   // Access all data registered
   Ptr getMergedSchema();
 
-  explicit SchemaRegistry(BaseSchemaRegistry& base) : base_(base) {}
+  /**
+   * Gets node for given definition, or throws `std::out_of_range` if not
+   * present in schema.
+   */
+  template <typename T>
+  const schema::DefinitionNode& getDefinitionNode() const {
+    return resolver_->getDefinitionNode<T>();
+  }
+
+  explicit SchemaRegistry(BaseSchemaRegistry& base);
+  ~SchemaRegistry();
 
  private:
   BaseSchemaRegistry& base_;
   Ptr mergedSchema_;
   folly::relaxed_atomic<bool> mergedSchemaAccessed_{false};
   std::unordered_set<type::ProgramId> includedPrograms_;
+  std::unique_ptr<schema::SyntaxGraph> syntaxGraph_;
+  schema::detail::IncrementalResolver* resolver_;
 };
 
 } // namespace apache::thrift

third-party/thrift/src/thrift/lib/cpp2/schema/SyntaxGraph.h

@@ -1449,11 +1449,11 @@ class SyntaxGraph final {
    */
   ProgramNode::IncludesList programs() const;
 
+  explicit SyntaxGraph(std::unique_ptr<detail::Resolver> resolver);
+
  private:
   folly::not_null_unique_ptr<const detail::Resolver> resolver_;
 
-  explicit SyntaxGraph(std::unique_ptr<detail::Resolver> resolver);
-
   friend const DefinitionNode& detail::lookUpDefinition(
       const SyntaxGraph&, const apache::thrift::type::DefinitionKey&);
 };

third-party/thrift/src/thrift/lib/cpp2/schema/detail/SchemaBackedResolver.cpp

@@ -17,6 +17,7 @@
 #include <thrift/lib/cpp2/schema/detail/SchemaBackedResolver.h>
 
 #include <thrift/lib/cpp2/schema/SyntaxGraph.h>
+#include <thrift/lib/cpp2/schema/detail/Merge.h>
 
 #ifdef THRIFT_SCHEMA_AVAILABLE
 
@@ -69,12 +70,12 @@ class DefinitionKeyEqual : public std::equal_to<type::DefinitionKey> {
     return Delegate::operator()(lhs.get(), rhs);
   }
 };
+} // namespace
 
 class SchemaIndex {
  public:
   explicit SchemaIndex(Resolver& resolver) : resolver_(resolver) {}
 
- private:
   // Annotations are referred to in schema.thrift using DefinitionKey, not
   // TypeStruct. Therefore, we use this map to keep their resolved TypeStruct
   // instances alive.
@@ -784,8 +785,6 @@ const type::Type* SchemaIndex::tryGetAnnotationType(
   return std::addressof(element->second);
 }
 
-} // namespace
-
 folly::not_null_unique_ptr<Resolver> createResolverfromSchema(
     type::Schema&& schema) {
   return std::make_unique<FullyResolvedSchemaBackedResolver>(std::move(schema));
@@ -795,5 +794,82 @@ folly::not_null_unique_ptr<Resolver> createResolverfromSchemaRef(
   return std::make_unique<FullyResolvedSchemaRefBackedResolver>(schema);
 }
 
+IncrementalResolver::IncrementalResolver()
+    : index_(std::make_unique<SchemaIndex>(*this)) {}
+IncrementalResolver::~IncrementalResolver() = default;
+
+const ProgramNode& IncrementalResolver::programOf(
+    const type::ProgramId& id) const {
+  return index_->programOf(id);
+}
+const protocol::Value& IncrementalResolver::valueOf(
+    const type::ValueId& id) const {
+  return index_->valueOf(id);
+}
+const DefinitionNode* IncrementalResolver::definitionOf(
+    const type::DefinitionKey& key) const {
+  return index_->definitionOf(key);
+}
+ProgramNode::IncludesList IncrementalResolver::programs() const {
+  return index_->programs();
+}
+
+const DefinitionNode& IncrementalResolver::getDefinitionNode(
+    const type::DefinitionKey& key,
+    type::ProgramId programId,
+    std::string_view name,
+    ::folly::Range<const ::std::string_view*> (*bundle)()) const {
+  {
+    auto schemaReadGuard = schema_.rlock();
+    if (auto* def = index_->definitionOf(key)) {
+      return *def;
+    }
+  }
+
+  if (!bundle) {
+    folly::throw_exception<std::out_of_range>(
+        fmt::format("Definition `{}` does not have bundled schema.", name));
+  }
+
+  auto schemaWriteGuard = schema_.wlock();
+  if (auto* def = index_->definitionOf(key)) {
+    return *def;
+  }
+
+  auto src = mergeSchemas((*bundle)());
+  auto& dst = *schemaWriteGuard;
+
+  // Merge new schema data in
+  // TODO: avoid deserializing shared deps
+  std::copy_if(
+      std::make_move_iterator(src.programs()->begin()),
+      std::make_move_iterator(src.programs()->end()),
+      std::back_inserter(*dst.programs()),
+      [this](const auto& program) {
+        return index_->programsById_.count(*program.id()) == 0;
+      });
+  dst.valuesMap()->insert(
+      std::make_move_iterator(src.valuesMap()->begin()),
+      std::make_move_iterator(src.valuesMap()->end()));
+  dst.definitionsMap()->insert(
+      std::make_move_iterator(src.definitionsMap()->begin()),
+      std::make_move_iterator(src.definitionsMap()->end()));
+
+  // Rebuild indices
+  // TODO: incremental rebuild
+  index_->init(dst);
+
+  if (auto* def = index_->definitionOf(key)) {
+    return *def;
+  }
+
+  if (index_->programsById_.count(programId)) {
+    folly::throw_exception<InvalidSyntaxGraphError>(fmt::format(
+        "Definition `{}` not found in its program's schema.", name));
+  }
+  folly::throw_exception<std::out_of_range>(
+      fmt::format("Definition `{}` does not have bundled schema.", name));
+}
+
 } // namespace apache::thrift::schema::detail
 #endif

third-party/thrift/src/thrift/lib/cpp2/schema/detail/SchemaBackedResolver.h

@@ -19,14 +19,56 @@
 #include <thrift/lib/cpp2/schema/SchemaV1.h>
 #include <thrift/lib/cpp2/schema/detail/Resolver.h>
 
+#include <folly/Synchronized.h>
+
 #ifdef THRIFT_SCHEMA_AVAILABLE
 
 namespace apache::thrift::schema::detail {
 namespace type = apache::thrift::type;
+class SchemaIndex;
 
 folly::not_null_unique_ptr<Resolver> createResolverfromSchema(type::Schema&&);
 folly::not_null_unique_ptr<Resolver> createResolverfromSchemaRef(
     const type::Schema&);
 
+class IncrementalResolver : public Resolver {
+ public:
+  IncrementalResolver();
+  ~IncrementalResolver() override;
+
+  /**
+   * Gets node for given definition, or throws `std::out_of_range` if not
+   * present in schema.
+   */
+  template <typename T>
+  const DefinitionNode& getDefinitionNode() const;
+
+  const ProgramNode& programOf(const type::ProgramId& id) const override;
+  const protocol::Value& valueOf(const type::ValueId& id) const override;
+  const DefinitionNode* definitionOf(
+      const type::DefinitionKey& key) const override;
+  std::vector<folly::not_null<const ProgramNode*>> programs() const override;
+
+ private:
+  const DefinitionNode& getDefinitionNode(
+      const type::DefinitionKey& key,
+      type::ProgramId programId,
+      std::string_view name,
+      ::folly::Range<const ::std::string_view*> (*bundle)()) const;
+
+  mutable folly::Synchronized<type::Schema> schema_;
+  folly::not_null_unique_ptr<SchemaIndex> index_;
+};
+
+template <typename T>
+const DefinitionNode& IncrementalResolver::getDefinitionNode() const {
+  using ::apache::thrift::detail::TSchemaAssociation;
+  return getDefinitionNode(
+      type::DefinitionKey{TSchemaAssociation<T>::definitionKey},
+      type::ProgramId{TSchemaAssociation<T>::programId},
+      folly::pretty_name<T>(),
+      TSchemaAssociation<T>::bundle);
+}
+
 } // namespace apache::thrift::schema::detail
 #endif

third-party/thrift/src/thrift/lib/cpp2/schema/test/SyntaxGraphTest.cpp

@@ -19,8 +19,10 @@
 
 #include <folly/Utility.h>
 
+#include <thrift/lib/cpp2/runtime/SchemaRegistry.h>
 #include <thrift/lib/cpp2/schema/SyntaxGraph.h>
 
+#include <thrift/annotation/gen-cpp2/thrift_types.h>
 #include <thrift/lib/cpp2/schema/test/gen-cpp2/syntax_graph_2_handlers.h>
 #include <thrift/lib/cpp2/schema/test/gen-cpp2/syntax_graph_handlers.h>
 #include <thrift/lib/cpp2/schema/test/gen-cpp2/syntax_graph_types.h>
@@ -461,4 +463,26 @@ TEST_F(ServiceSchemaTest, LookupByDefinitionKeys) {
   }
 }
 
+TEST(SyntaxGraphTest, getDefinitionNode) {
+  auto& registry = SchemaRegistry::get();
+
+  const DefinitionNode& testStruct =
+      registry.getDefinitionNode<test::TestStruct>();
+  const StructNode& stct = testStruct.asStruct();
+  EXPECT_EQ(stct.definition().name(), "TestStruct");
+
+  const DefinitionNode& testEnum = registry.getDefinitionNode<test::TestEnum>();
+  const EnumNode& enm = testEnum.asEnum();
+  EXPECT_EQ(enm.definition().name(), "TestEnum");
+
+  const DefinitionNode& testService =
+      registry.getDefinitionNode<test::TestService>();
+  const ServiceNode& serv = testService.asService();
+  EXPECT_EQ(serv.definition().name(), "TestService");
+
+  EXPECT_THROW(
+      registry.getDefinitionNode<facebook::thrift::annotation::Experimental>(),
+      std::out_of_range);
+}
+
 } // namespace apache::thrift::schema