Allow SchemaLoader to use a callback to lazily load schema nodes, including…

Allow SchemaLoader to use a callback to lazily load schema nodes, including dependencies of loaded nodes.

c++/src/capnp/generated-header-support.h

@@ -147,11 +147,16 @@ struct RawSchema {
   //
   // This is an internal structure which could change in the future.
 
+  uint64_t id;
+
   const word* encodedNode;
   // Encoded SchemaNode, readable via readMessageUnchecked<schema::Node>(encodedNode).
 
   const RawSchema* const* dependencies;
-  // Pointers to other types on which this one depends, sorted by ID.
+  // Pointers to other types on which this one depends, sorted by ID.  The schemas in this table
+  // may be uninitialized -- you must call ensureInitialized() on the one you wish to use before
+  // using it.
+  //
   // TODO(someday):  Make this a hashtable.
 
   struct MemberInfo {
@@ -171,6 +176,23 @@ struct RawSchema {
   // Points to the RawSchema of a compiled-in type to which it is safe to cast any DynamicValue
   // with this schema.  This is null for all compiled-in types; it is only set by SchemaLoader on
   // dynamically-loaded types.
+
+  class Initializer {
+  public:
+    virtual void init(const RawSchema* schema) const = 0;
+  };
+
+  const Initializer* lazyInitializer;
+  // Lazy initializer, invoked by ensureInitialized().
+
+  inline void ensureInitialized() const {
+    // Lazy initialization support.  Invoke to ensure that initialization has taken place.  This
+    // is required in particular when traversing the dependency list.  RawSchemas for compiled-in
+    // types are always initialized; only dynamically-loaded schemas may be lazy.
+
+    const Initializer* i = __atomic_load_n(&lazyInitializer, __ATOMIC_ACQUIRE);
+    if (i != nullptr) i->init(this);
+  }
 };
 
 template <typename T>

c++/src/capnp/schema-loader-test.c++

@@ -208,6 +208,8 @@ TEST(SchemaLoader, Incompatible) {
       loadUnderAlternateTypeId<test::TestAllTypes>(loader, typeId<test::TestListDefaults>()));
 }
 
+// TODO(test):  More extensively test upgrade/downgrade checks.
+
 TEST(SchemaLoader, Enumerate) {
   SchemaLoader loader;
   loader.loadCompiledTypeAndDependencies<TestAllTypes>();
@@ -222,7 +224,88 @@ TEST(SchemaLoader, Enumerate) {
   }
 }
 
-// TODO(test):  More extensively test upgrade/downgrade checks.
+TEST(SchemaLoader, EnumerateNoPlaceholders) {
+  SchemaLoader loader;
+  Schema schema = loader.load(Schema::from<TestDefaults>().getProto());
+
+  {
+    auto list = loader.getAllLoaded();
+    ASSERT_EQ(1u, list.size());
+    EXPECT_TRUE(list[0] == schema);
+  }
+
+  Schema dep = schema.getDependency(typeId<TestAllTypes>());
+
+  {
+    auto list = loader.getAllLoaded();
+    ASSERT_EQ(2u, list.size());
+    if (list[0] == schema) {
+      EXPECT_TRUE(list[1] == dep);
+    } else {
+      EXPECT_TRUE(list[0] == dep);
+      EXPECT_TRUE(list[1] == schema);
+    }
+  }
+}
+
+class FakeLoaderCallback: public SchemaLoader::LazyLoadCallback {
+public:
+  FakeLoaderCallback(const schema::Node::Reader node): node(node), loaded(false) {}
+
+  bool isLoaded() { return loaded; }
+
+  void load(const SchemaLoader& loader, uint64_t id) const override {
+    if (id == 1234) {
+      // Magic "not found" ID.
+      return;
+    }
+
+    EXPECT_EQ(node.getId(), id);
+    EXPECT_FALSE(loaded);
+    loaded = true;
+    loader.loadIfNew(node);
+  }
+
+private:
+  const schema::Node::Reader node;
+  mutable bool loaded = false;
+};
+
+TEST(SchemaLoader, LazyLoad) {
+  FakeLoaderCallback callback(Schema::from<TestAllTypes>().getProto());
+  SchemaLoader loader(callback);
+
+  EXPECT_TRUE(loader.tryGet(1234) == nullptr);
+
+  EXPECT_FALSE(callback.isLoaded());
+  Schema schema = loader.get(typeId<TestAllTypes>());
+  EXPECT_TRUE(callback.isLoaded());
+
+  EXPECT_EQ(schema.getProto().getDisplayName(),
+            Schema::from<TestAllTypes>().getProto().getDisplayName());
+
+  EXPECT_EQ(schema, schema.getDependency(typeId<TestAllTypes>()));
+  EXPECT_EQ(schema, loader.get(typeId<TestAllTypes>()));
+}
+
+TEST(SchemaLoader, LazyLoadGetDependency) {
+  FakeLoaderCallback callback(Schema::from<TestAllTypes>().getProto());
+  SchemaLoader loader(callback);
+
+  Schema schema = loader.load(Schema::from<TestDefaults>().getProto());
+
+  EXPECT_FALSE(callback.isLoaded());
+
+  Schema dep = schema.getDependency(typeId<TestAllTypes>());
+
+  EXPECT_TRUE(callback.isLoaded());
+
+  EXPECT_EQ(dep.getProto().getDisplayName(),
+            Schema::from<TestAllTypes>().getProto().getDisplayName());
+
+  EXPECT_EQ(dep, schema.getDependency(typeId<TestAllTypes>()));
+  EXPECT_EQ(dep, loader.get(typeId<TestAllTypes>()));
+}
 
 }  // namespace
 }  // namespace _ (private)

c++/src/capnp/schema-loader.h

@@ -26,12 +26,34 @@
 
 #include "schema.h"
 #include <kj/memory.h>
+#include <kj/mutex.h>
 
 namespace capnp {
 
 class SchemaLoader {
 public:
+  class LazyLoadCallback {
+  public:
+    virtual void load(const SchemaLoader& loader, uint64_t id) const = 0;
+    // Request that the schema node with the given ID be loaded into the given SchemaLoader.  If
+    // the callback is able to find a schema for this ID, it should invoke `loadIfNew()` on
+    // `loader` to load it.  If no such node exists, it should simply do nothing and return.
+    //
+    // The callback is allowed to load schema nodes other than the one requested, e.g. because it
+    // expects they will be needed soon.
+    //
+    // If the `SchemaLoader` is used from multiple threads, the callback must be thread-safe.
+    // In particular, it's possible for multiple threads to invoke `load()` with the same ID.
+    // If the callback performs a large amount of work to look up IDs, it should be sure to
+    // de-dup these requests.
+  };
+
   SchemaLoader();
+
+  SchemaLoader(const LazyLoadCallback& callback);
+  // Construct a SchemaLoader which will invoke the given callback when a schema node is requested
+  // that isn't already loaded.
+
   ~SchemaLoader() noexcept(false);
   KJ_DISALLOW_COPY(SchemaLoader);
 
@@ -79,6 +101,11 @@ public:
   // Also note that unknown types are not considered invalid.  Instead, the dynamic API returns
   // a DynamicValue with type UNKNOWN for these.
 
+  Schema loadIfNew(const schema::Node::Reader& reader) const;
+  // Like `load()` but does nothing if a schema with the same ID is already loaded.  In contrast,
+  // `load()` would attempt to compare the schemas and take the newer one.  `loadIfNew()` is safe
+  // to call even while concurrently using schemas from this loader.
+
   template <typename T>
   void loadCompiledTypeAndDependencies();
   // Load the schema for the given compiled-in type and all of its dependencies.
@@ -96,7 +123,8 @@ private:
   class Validator;
   class CompatibilityChecker;
   class Impl;
-  kj::Own<Impl> impl;
+  class InitializerImpl;
+  kj::MutexGuarded<kj::Own<Impl>> impl;
 
   void loadNative(const _::RawSchema* nativeSchema);
 };

c++/src/capnp/schema.c++

@@ -38,11 +38,12 @@ Schema Schema::getDependency(uint64_t id) const {
   while (lower < upper) {
     uint mid = (lower + upper) / 2;
 
-    Schema candidate(raw->dependencies[mid]);
+    const _::RawSchema* candidate = raw->dependencies[mid];
 
-    uint64_t candidateId = candidate.getProto().getId();
+    uint64_t candidateId = candidate->id;
     if (candidateId == id) {
-      return candidate;
+      candidate->ensureInitialized();
+      return Schema(candidate);
     } else if (candidateId < id) {
       lower = mid + 1;
     } else {

c++/src/capnp/schema.h

@@ -86,7 +86,10 @@ public:
 private:
   const _::RawSchema* raw;
 
-  inline explicit Schema(const _::RawSchema* raw): raw(raw) {}
+  inline explicit Schema(const _::RawSchema* raw): raw(raw) {
+    KJ_IREQUIRE(raw->lazyInitializer == nullptr,
+        "Must call ensureInitialized() on RawSchema before constructing Schema.");
+  }
 
   template <typename T> static inline Schema fromImpl() {
     return Schema(&_::rawSchema<T>());

c++/src/kj/mutex.c++

@@ -50,15 +50,15 @@ Mutex::~Mutex() {
   KJ_PTHREAD_CLEANUP(pthread_rwlock_destroy(&mutex));
 }
 
-void Mutex::lock() noexcept {
+void Mutex::lock() {
   KJ_PTHREAD_CALL(pthread_rwlock_wrlock(&mutex));
 }
 
-void Mutex::readLock() noexcept {
+void Mutex::readLock() {
   KJ_PTHREAD_CALL(pthread_rwlock_rdlock(&mutex));
 }
 
-void Mutex::unlock(bool lockedForRead) noexcept {
+void Mutex::unlock(bool lockedForRead) {
   KJ_PTHREAD_CALL(pthread_rwlock_unlock(&mutex));
 }
 

c++/src/kj/mutex.h

@@ -47,9 +47,9 @@ public:
   ~Mutex();
   KJ_DISALLOW_COPY(Mutex);
 
-  void lock() noexcept;
-  void readLock() noexcept;
-  void unlock(bool lockedForRead) noexcept;
+  void lock();
+  void readLock();
+  void unlock(bool lockedForRead);
 
 private:
   mutable pthread_rwlock_t mutex;

compiler/src/c++-source.mustache

@@ -69,8 +69,8 @@ static const ::capnp::_::RawSchema::MemberInfo m_{{schemaId}}[] = {
 {{/schemaMembersByName}}
 };
 const ::capnp::_::RawSchema s_{{schemaId}} = {
-  b_{{schemaId}}.words, d_{{schemaId}}, m_{{schemaId}},
-  {{schemaDependencyCount}}, {{schemaMemberCount}}, nullptr
+  0x{{schemaId}}, b_{{schemaId}}.words, d_{{schemaId}}, m_{{schemaId}},
+  {{schemaDependencyCount}}, {{schemaMemberCount}}, nullptr, nullptr
 };
 {{/typeSchema}}
 {{/fileTypes}}