From b11ac4971c1856d3cb5d08b2c26a37c756f44141 Mon Sep 17 00:00:00 2001
From: Ian Denhardt <ian@zenhack.net>
Date: Mon, 20 Mar 2023 20:35:31 -0400
Subject: [PATCH] Fix the ownership semantics for NewLocalPromise

These are much cleaner.
---
 localpromise.go          | 43 +++++++++++++++++++++++++---------------
 rpc/localpromise_test.go | 20 +++++++++++++++----
 2 files changed, 43 insertions(+), 20 deletions(-)

diff --git a/localpromise.go b/localpromise.go
index f366884f..1c131b5b 100644
--- a/localpromise.go
+++ b/localpromise.go
@@ -5,15 +5,15 @@ import (
 )
 
 // ClientHook for a promise that will be resolved to some other capability
-// at some point. Buffers calls in a queue until the promsie is fulfilled,
+// at some point. Buffers calls in a queue until the promise is fulfilled,
 // then forwards them.
 type localPromise struct {
 	aq *AnswerQueue
 }
 
 // NewLocalPromise returns a client that will eventually resolve to a capability,
-// supplied via the fulfiller.
-func NewLocalPromise[C ~ClientKind]() (C, Resolver[C]) {
+// supplied via resolver. resolver.Fulfill steals the reference to its argument.
+func NewLocalPromise[C ~ClientKind]() (promise C, resolver Resolver[C]) {
 	lp := newLocalPromise()
 	p, f := NewPromisedClient(lp)
 	return C(p), localResolver[C]{
@@ -44,27 +44,38 @@ func (lp localPromise) String() string {
 	return "localPromise{...}"
 }
 
-func (lp localPromise) Fulfill(c Client) {
-	msg, seg := NewSingleSegmentMessage(nil)
-	capID := msg.AddCap(c)
-	lp.aq.Fulfill(NewInterface(seg, capID).ToPtr())
-}
-
-func (lp localPromise) Reject(err error) {
-	lp.aq.Reject(err)
-}
-
 type localResolver[C ~ClientKind] struct {
 	lp             localPromise
 	clientResolver Resolver[Client]
 }
 
 func (lf localResolver[C]) Fulfill(c C) {
-	lf.lp.Fulfill(Client(c))
-	lf.clientResolver.Fulfill(Client(c))
+	// This is convoluted, for a few reasons:
+	//
+	// 1. AnswerQueue wants a Ptr, not a Client, so we have to construct
+	//    a message for this.
+	// 2. Message.AddCap steals the reference, so we have to get the client
+	//    back from the message instead of using the reference we already
+	//    have.
+	// 3. The semantics of NewPromisedClient differ from what we want, and
+	//    are kindof odd: when it is resolved it does not steal the reference,
+	//    nor does it borrow it -- instead it merges the refcounts, so that
+	//    the two clients point to the same place. So we have to drop our
+	//    reference to get the semantics we want.
+	//
+	//    TODO: We should probably push this part down into the implementation of
+	//    clientPromise. That requires auditing its uses and adjusting call sites
+	//    though.
+	msg, seg := NewSingleSegmentMessage(nil)
+	capID := msg.AddCap(Client(c))
+	iface := NewInterface(seg, capID)
+	client := iface.Client()
+	defer client.Release()
+	lf.lp.aq.Fulfill(iface.ToPtr())
+	lf.clientResolver.Fulfill(client)
 }
 
 func (lf localResolver[C]) Reject(err error) {
-	lf.lp.Reject(err)
+	lf.lp.aq.Reject(err)
 	lf.clientResolver.Reject(err)
 }
diff --git a/rpc/localpromise_test.go b/rpc/localpromise_test.go
index e67cbfaf..5501209e 100644
--- a/rpc/localpromise_test.go
+++ b/rpc/localpromise_test.go
@@ -45,12 +45,10 @@ func TestLocalPromiseFulfill(t *testing.T) {
 	})
 	defer rel2()
 
-	pp := testcapnp.PingPong_ServerToClient(&echoNumOrderChecker{
+	r.Fulfill(testcapnp.PingPong_ServerToClient(&echoNumOrderChecker{
 		t:       t,
 		nextNum: 1,
-	})
-	defer pp.Release()
-	r.Fulfill(pp)
+	}))
 
 	fut3, rel3 := p.EchoNum(ctx, func(p testcapnp.PingPong_echoNum_Params) error {
 		p.SetN(3)
@@ -104,3 +102,17 @@ func TestLocalPromiseReject(t *testing.T) {
 	_, err = fut3.Struct()
 	assert.NotNil(t, err)
 }
+
+// Test that the promise owns the capability it resolves to; no separate
+// release should be necessary.
+func TestLocalPromiseOwnsResult(t *testing.T) {
+	t.Parallel()
+
+	p, r := capnp.NewLocalPromise[testcapnp.PingPong]()
+	defer p.Release()
+
+	r.Fulfill(testcapnp.PingPong_ServerToClient(&echoNumOrderChecker{
+		t:       t,
+		nextNum: 1,
+	}))
+}