Add thisCap() method to capability servers which returns a capability to itself.…

Add thisCap() method to capability servers which returns a capability to itself. Remove CapabilityServerSet's weak pointers since this supersedes them.

c++/src/capnp/capability-test.c++

@@ -854,9 +854,7 @@ TEST(Capability, CapabilityServerSet) {
 
   auto ownServer2 = kj::heap<TestInterfaceImpl>(callCount);
   auto& server2 = *ownServer2;
-  auto client2AndWeak = set2.addWeak(kj::mv(ownServer2));
-  test::TestInterface::Client client2 = kj::mv(client2AndWeak.client);
-  kj::Own<WeakCapability<test::TestInterface>> client2Weak = kj::mv(client2AndWeak.weak);
+  test::TestInterface::Client client2 = set2.add(kj::mv(ownServer2));
 
   // Getting the local server using the correct set works.
   EXPECT_EQ(&server1, &KJ_ASSERT_NONNULL(set1.getLocalServer(client1).wait(waitScope)));
@@ -913,16 +911,53 @@ TEST(Capability, CapabilityServerSet) {
   EXPECT_TRUE(resolved1);
   EXPECT_TRUE(resolved2);
   EXPECT_TRUE(resolved3);
+}
 
-  // Check weak pointer.
-  {
-    auto restored = KJ_ASSERT_NONNULL(client2Weak->get());
-    EXPECT_EQ(&server2, &KJ_ASSERT_NONNULL(set2.getLocalServer(restored).wait(waitScope)));
+class TestThisCap final: public test::TestInterface::Server {
+public:
+  TestThisCap(int& callCount): callCount(callCount) {}
+  ~TestThisCap() noexcept(false) { callCount = -1; }
+
+  test::TestInterface::Client getSelf() {
+    return thisCap();
+  }
+
+protected:
+  kj::Promise<void> bar(BarContext context) {
+    ++callCount;
+    return kj::READY_NOW;
   }
 
+private:
+  int& callCount;
+};
+
+TEST(Capability, ThisCap) {
+  int callCount = 0;
+  kj::EventLoop loop;
+  kj::WaitScope waitScope(loop);
+
+  auto server = kj::heap<TestThisCap>(callCount);
+  TestThisCap* serverPtr = server;
+
+  test::TestInterface::Client client = kj::mv(server);
+  client.barRequest().send().wait(waitScope);
+  EXPECT_EQ(1, callCount);
+
+  test::TestInterface::Client client2 = serverPtr->getSelf();
+  EXPECT_EQ(1, callCount);
+  client2.barRequest().send().wait(waitScope);
+  EXPECT_EQ(2, callCount);
+
+  client = nullptr;
+
+  EXPECT_EQ(2, callCount);
+  client2.barRequest().send().wait(waitScope);
+  EXPECT_EQ(3, callCount);
+
   client2 = nullptr;
 
-  EXPECT_TRUE(client2Weak->get() == nullptr);
+  EXPECT_EQ(-1, callCount);
 }
 
 }  // namespace

c++/src/capnp/capability.c++

@@ -463,21 +463,18 @@ private:
 
 class LocalClient final: public ClientHook, public kj::Refcounted {
 public:
-  LocalClient(kj::Own<Capability::Server>&& server): server(kj::mv(server)) {}
-  LocalClient(kj::Own<Capability::Server>&& server,
+  LocalClient(kj::Own<Capability::Server>&& serverParam)
+      : server(kj::mv(serverParam)) {
+    server->thisHook = this;
+  }
+  LocalClient(kj::Own<Capability::Server>&& serverParam,
               _::CapabilityServerSetBase& capServerSet, void* ptr)
-      : server(kj::mv(server)), capServerSet(&capServerSet), ptr(ptr) {}
-
-  ~LocalClient() noexcept(false) {
-    KJ_IF_MAYBE(w, weak) {
-      w->client = nullptr;
-    }
+      : server(kj::mv(serverParam)), capServerSet(&capServerSet), ptr(ptr) {
+    server->thisHook = this;
   }
 
-  void setWeak(_::WeakCapabilityBase& weak) {
-    KJ_REQUIRE(this->weak == nullptr && weak.client == nullptr);
-    weak.client = *this;
-    this->weak = weak;
+  ~LocalClient() noexcept(false) {
+    server->thisHook = nullptr;
   }
 
   Request<AnyPointer, AnyPointer> newCall(
@@ -555,8 +552,6 @@ private:
   kj::Own<Capability::Server> server;
   _::CapabilityServerSetBase* capServerSet = nullptr;
   void* ptr = nullptr;
-  kj::Maybe<_::WeakCapabilityBase&> weak;
-  friend class _::WeakCapabilityBase;
 };
 
 kj::Own<ClientHook> Capability::Client::makeLocalClient(kj::Own<Capability::Server>&& server) {
@@ -680,30 +675,11 @@ Request<AnyPointer, AnyPointer> newBrokenRequest(
 
 namespace _ {  // private
 
-WeakCapabilityBase::~WeakCapabilityBase() noexcept(false) {
-  KJ_IF_MAYBE(c, client) {
-    c->weak = nullptr;
-  }
-}
-
-kj::Maybe<Capability::Client> WeakCapabilityBase::getInternal() {
-  return client.map([](LocalClient& client) {
-    return Capability::Client(client.addRef());
-  });
-}
-
 Capability::Client CapabilityServerSetBase::addInternal(
     kj::Own<Capability::Server>&& server, void* ptr) {
   return Capability::Client(kj::refcounted<LocalClient>(kj::mv(server), *this, ptr));
 }
 
-Capability::Client CapabilityServerSetBase::addWeakInternal(
-    kj::Own<Capability::Server>&& server, _::WeakCapabilityBase& weak, void* ptr) {
-  auto result = kj::refcounted<LocalClient>(kj::mv(server), *this, ptr);
-  result->setWeak(weak);
-  return Capability::Client(kj::mv(result));
-}
-
 kj::Promise<void*> CapabilityServerSetBase::getLocalServerInternal(Capability::Client& client) {
   ClientHook* hook = client.hook.get();
 

c++/src/capnp/capability.h

@@ -332,6 +332,17 @@ public:
   //   a proxy.
 
 protected:
+  inline Capability::Client thisCap();
+  // Get a capability pointing to this object, much like the `this` keyword.
+  //
+  // The effect of this method is undefined if:
+  // - No capability client has been created pointing to this object. (This is always the case in
+  //   the server's constructor.)
+  // - The capability client pointing at this object has been destoryed. (This is always the case
+  //   in the server's destructor.)
+  // - Multiple capability clients have been created around the same server (possible if the server
+  //   is refcounted, which is not recommended since the client itself provides refcounting).
+
   template <typename Params, typename Results>
   CallContext<Params, Results> internalGetTypedContext(
       CallContext<AnyPointer, AnyPointer> typeless);
@@ -341,52 +352,28 @@ protected:
                                           uint64_t typeId, uint16_t methodId);
   kj::Promise<void> internalUnimplemented(const char* interfaceName, const char* methodName,
                                           uint64_t typeId, uint16_t methodId);
+
+private:
+  ClientHook* thisHook = nullptr;
+  friend class LocalClient;
 };
 
 // =======================================================================================
 
 namespace _ {  // private
 
-class WeakCapabilityBase {
-public:
-  ~WeakCapabilityBase() noexcept(false);
-
-  kj::Maybe<Capability::Client> getInternal();
-
-private:
-  kj::Maybe<LocalClient&> client;
-  friend class capnp::LocalClient;
-};
-
 class CapabilityServerSetBase {
 public:
   Capability::Client addInternal(kj::Own<Capability::Server>&& server, void* ptr);
-  Capability::Client addWeakInternal(kj::Own<Capability::Server>&& server,
-                                     _::WeakCapabilityBase& weak, void* ptr);
   kj::Promise<void*> getLocalServerInternal(Capability::Client& client);
 };
 
 }  // namespace _ (private)
 
-template <typename T>
-class WeakCapability: private _::WeakCapabilityBase {
-public:
-  kj::Maybe<typename T::Client> get();
-  // If the server is still alive, get a live client to it.
-
-private:
-  template <typename>
-  friend class CapabilityServerSet;
-};
-
 template <typename T>
 class CapabilityServerSet: private _::CapabilityServerSetBase {
-  // Allows a server to:
-  // 1) Recognize its own capabilities when passed back to it, and obtain the underlying Server
-  //    objects associated with them.
-  // 2) Obtain "weak" versions of these capabilities, which do not prevent the underlying Server
-  //    from being destroyed but can be upgraded to normal Clients as long as the Server is still
-  //    alive.
+  // Allows a server to recognize its own capabilities when passed back to it, and obtain the
+  // underlying Server objects associated with them.
   //
   // All objects in the set must have the same interface type T. The objects may implement various
   // interfaces derived from T (and in fact T can be `capnp::Capability` to accept all objects),
@@ -404,14 +391,6 @@ public:
   typename T::Client add(kj::Own<typename T::Server>&& server);
   // Create a new capability Client for the given Server and also add this server to the set.
 
-  struct ClientAndWeak {
-    typename T::Client client;
-    kj::Own<WeakCapability<T>> weak;
-  };
-
-  ClientAndWeak addWeak(kj::Own<typename T::Server>&& server);
-  // Like add() but also creates a weak reference.
-
   kj::Promise<kj::Maybe<typename T::Server&>> getLocalServer(typename T::Client& client);
   // Given a Client pointing to a server previously passed to add(), return the corresponding
   // Server. This returns a promise because if the input client is itself a promise, this must
@@ -782,11 +761,8 @@ CallContext<Params, Results> Capability::Server::internalGetTypedContext(
   return CallContext<Params, Results>(*typeless.hook);
 }
 
-template <typename T>
-kj::Maybe<typename T::Client> WeakCapability<T>::get() {
-  return getInternal().map([](Capability::Client&& client) {
-    return client.castAs<T>();
-  });
+Capability::Client Capability::Server::thisCap() {
+  return Client(thisHook->addRef());
 }
 
 template <typename T>
@@ -797,17 +773,6 @@ typename T::Client CapabilityServerSet<T>::add(kj::Own<typename T::Server>&& ser
   return addInternal(kj::mv(server), ptr).template castAs<T>();
 }
 
-template <typename T>
-typename CapabilityServerSet<T>::ClientAndWeak CapabilityServerSet<T>::addWeak(
-    kj::Own<typename T::Server>&& server) {
-  void* ptr = reinterpret_cast<void*>(server.get());
-  auto weak = kj::heap<WeakCapability<T>>();
-  // Clang insists that `castAs` is a template-dependent member and therefore we need the
-  // `template` keyword here, but AFAICT this is wrong: addWeakImpl() is not a template.
-  auto client = addWeakInternal(kj::mv(server), *weak, ptr).template castAs<T>();
-  return { kj::mv(client), kj::mv(weak) };
-}
-
 template <typename T>
 kj::Promise<kj::Maybe<typename T::Server&>> CapabilityServerSet<T>::getLocalServer(
     typename T::Client& client) {

c++/src/capnp/compiler/capnpc-c++.c++

@@ -250,6 +250,10 @@ public:
     return hasInterfaces_;
   }
 
+  kj::StringTree strNoTypename() const & { return name.flatten(); }
+  kj::StringTree strNoTypename() && { return kj::mv(name); }
+  // Stringify but never prefix with `typename`. Use in contexts where `typename` is implicit.
+
 private:
   kj::StringTree name;
 
@@ -2164,7 +2168,9 @@ private:
       };
     }
 
-    CppTypeName clientName = cppFullName(schema, nullptr);
+    CppTypeName typeName = cppFullName(schema, nullptr);
+
+    CppTypeName clientName = typeName;
     clientName.addMemberType("Client");
 
     kj::String templates = kj::str(templateContext.allDecls());  // Ends with a newline
@@ -2227,7 +2233,7 @@ private:
           "class ", fullName, "::Client\n"
           "    : public virtual ::capnp::Capability::Client",
           KJ_MAP(s, superclasses) {
-            return kj::strTree(",\n      public virtual ", s.typeName, "::Client");
+            return kj::strTree(",\n      public virtual ", s.typeName.strNoTypename(), "::Client");
           }, " {\n"
           "public:\n"
           "  typedef ", fullName, " Calls;\n"
@@ -2261,7 +2267,7 @@ private:
           "class ", fullName, "::Server\n"
           "    : public virtual ::capnp::Capability::Server",
           KJ_MAP(s, superclasses) {
-            return kj::strTree(",\n      public virtual ", s.typeName, "::Server");
+            return kj::strTree(",\n      public virtual ", s.typeName.strNoTypename(), "::Server");
           }, " {\n"
           "public:\n",
           "  typedef ", fullName, " Serves;\n"
@@ -2273,6 +2279,11 @@ private:
           "protected:\n",
           KJ_MAP(m, methods) { return kj::mv(m.serverDecls); },
           "\n"
+          "  inline ", clientName, " thisCap() {\n"
+          "    return ::capnp::Capability::Server::thisCap()\n"
+          "        .template castAs<", typeName, ">();\n"
+          "  }\n"
+          "\n"
           "  ::kj::Promise<void> dispatchCallInternal(uint16_t methodId,\n"
           "      ::capnp::CallContext< ::capnp::AnyPointer, ::capnp::AnyPointer> context);\n"
           "};\n"
@@ -2326,7 +2337,8 @@ private:
           KJ_MAP(s, transitiveSuperclasses) {
             return kj::strTree(
               "    case 0x", kj::hex(s.id), "ull:\n"
-              "      return ", s.typeName, "::Server::dispatchCallInternal(methodId, context);\n");
+              "      return ", s.typeName.strNoTypename(),
+                         "::Server::dispatchCallInternal(methodId, context);\n");
           },
           "    default:\n"
           "      return internalUnimplemented(\"", proto.getDisplayName(), "\", interfaceId);\n"

c++/src/capnp/persistent.capnp.h

@@ -204,6 +204,11 @@ protected:
   typedef ::capnp::CallContext<typename  ::capnp::Persistent<SturdyRef, Owner>::SaveParams, typename  ::capnp::Persistent<SturdyRef, Owner>::SaveResults> SaveContext;
   virtual ::kj::Promise<void> save(SaveContext context);
 
+  inline typename  ::capnp::Persistent<SturdyRef, Owner>::Client thisCap() {
+    return ::capnp::Capability::Server::thisCap()
+        .template castAs< ::capnp::Persistent<SturdyRef, Owner>>();
+  }
+
   ::kj::Promise<void> dispatchCallInternal(uint16_t methodId,
       ::capnp::CallContext< ::capnp::AnyPointer, ::capnp::AnyPointer> context);
 };
@@ -449,6 +454,11 @@ protected:
   typedef ::capnp::CallContext<ExportParams, typename  ::capnp::Persistent<ExternalRef, ExternalOwner>::SaveResults> ExportContext;
   virtual ::kj::Promise<void> export_(ExportContext context);
 
+  inline typename  ::capnp::RealmGateway<InternalRef, ExternalRef, InternalOwner, ExternalOwner>::Client thisCap() {
+    return ::capnp::Capability::Server::thisCap()
+        .template castAs< ::capnp::RealmGateway<InternalRef, ExternalRef, InternalOwner, ExternalOwner>>();
+  }
+
   ::kj::Promise<void> dispatchCallInternal(uint16_t methodId,
       ::capnp::CallContext< ::capnp::AnyPointer, ::capnp::AnyPointer> context);
 };

c++/src/kj/thread.h

@@ -50,6 +50,9 @@ public:
 
   void detach();
   // Don't join the thread in ~Thread().
+  //
+  // TODO(soon): Currently broken: the thread uses the Thread objects during its execution; instead
+  //   the Thread object and the thread itself will need to share a refcounted object.
 
 private:
   Function<void()> func;