Merge clientHook refcounts, fix #348

capability.go

@@ -138,9 +138,6 @@ type clientHook struct {
 	// Place for callers to attach arbitrary metadata to the client.
 	metadata Metadata
 
-	// done is closed when refs == 0 and calls == 0.
-	done chan struct{}
-
 	state mutex.Mutex[clientHookState]
 }
 
@@ -149,7 +146,6 @@ type clientHookState struct {
 	resolved chan struct{}
 
 	refs         int         // how many open Clients reference this clientHook
-	calls        int         // number of outstanding ClientHook accesses
 	resolvedHook *clientHook // valid only if resolved is closed
 }
 
@@ -164,7 +160,6 @@ func NewClient(hook ClientHook) Client {
 	}
 	h := &clientHook{
 		ClientHook: hook,
-		done:       make(chan struct{}),
 		metadata:   *NewMetadata(),
 		state: mutex.New(clientHookState{
 			resolved: closedSignal,
@@ -199,7 +194,6 @@ func newPromisedClient(hook ClientHook) (Client, *clientPromise) {
 	}
 	h := &clientHook{
 		ClientHook: hook,
-		done:       make(chan struct{}),
 		metadata:   *NewMetadata(),
 		state: mutex.New(clientHookState{
 			resolved: make(chan struct{}),
@@ -228,17 +222,19 @@ func (c Client) startCall() (hook ClientHook, resolved, released bool, finish fu
 		if c.h == nil {
 			return nil, true, false, func() {}
 		}
-		l.Value().calls++
+		l.Value().refs++
 		isResolved := l.Value().isResolved()
 		l.Unlock()
 		savedHook := c.h
 		return savedHook.ClientHook, isResolved, false, func() {
+			shutdown := func() {}
 			savedHook.state.With(func(s *clientHookState) {
-				s.calls--
-				if s.refs == 0 && s.calls == 0 {
-					close(savedHook.done)
+				s.refs--
+				if s.refs == 0 {
+					shutdown = savedHook.Shutdown
 				}
 			})
+			shutdown()
 		}
 	})
 }
@@ -653,12 +649,8 @@ func (c Client) Release() {
 		cl.Unlock()
 		return
 	}
-	if hl.Value().calls == 0 {
-		close(h.done)
-	}
 	hl.Unlock()
 	cl.Unlock()
-	<-h.done
 	h.Shutdown()
 	c.GetFlowLimiter().Release()
 }
@@ -735,7 +727,6 @@ func (cp *clientPromise) Fulfill(c Client) {
 // references to be dropped, and then shuts down the hook. The caller
 // must have previously invoked cp.fulfill().
 func (cp *clientPromise) shutdown() {
-	<-cp.h.done
 	cp.h.Shutdown()
 }
 
@@ -769,10 +760,6 @@ func (cp *clientPromise) fulfill(c Client) {
 		return
 	}
 
-	// Client still had references, so we're responsible for shutting it down.
-	if l.Value().calls == 0 {
-		close(cp.h.done)
-	}
 	rh, l = resolveHook(cp.h, l) // swaps mutex on cp.h for mutex on rh
 	if rh != nil {
 		l.Value().refs += refs
@@ -854,8 +841,8 @@ type ClientHook interface {
 
 	// Shutdown releases any resources associated with this capability.
 	// The behavior of calling any methods on the receiver after calling
-	// Shutdown is undefined.  It is expected for the ClientHook to reject
-	// any outstanding call futures.
+	// Shutdown is undefined. Shutdown must not interrupt any already
+	// outstanding calls.
 	Shutdown()
 
 	// String formats the hook as a string (same as fmt.Stringer)
@@ -1008,7 +995,6 @@ func ErrorClient(e error) Client {
 	// Avoid NewClient because it can set a finalizer.
 	h := &clientHook{
 		ClientHook: errorClient{e},
-		done:       make(chan struct{}),
 		metadata:   *NewMetadata(),
 		state: mutex.New(clientHookState{
 			resolved: closedSignal,

rpc/level0_test.go

@@ -807,8 +807,35 @@ func TestSendBootstrapPipelineCall(t *testing.T) {
 		}
 	}
 
-	// 6. Release the client, read the finish.
-	client.Release()
+	// 6. Send back a return for the bootstrap message:
+	bootstrapExportID := uint32(99)
+	{
+		outMsg, err := p2.NewMessage()
+		require.NoError(t, err)
+		iface := capnp.NewInterface(outMsg.Message().Segment(), 0)
+		require.NoError(t, pogs.Insert(rpccp.Message_TypeID, capnp.Struct(outMsg.Message()),
+			&rpcMessage{
+				Which: rpccp.Message_Which_return,
+				Return: &rpcReturn{
+					AnswerID: bootstrapQID,
+					Which:    rpccp.Return_Which_results,
+					Results: &rpcPayload{
+						Content: iface.ToPtr(),
+						CapTable: []rpcCapDescriptor{
+							{
+								Which:        rpccp.CapDescriptor_Which_senderHosted,
+								SenderHosted: bootstrapExportID,
+							},
+						},
+					},
+				},
+			},
+		))
+		require.NoError(t, outMsg.Send())
+		outMsg.Release()
+	}
+
+	// 7. Read the finish:
 	{
 		rmsg, release, err := recvMessage(ctx, p2)
 		if err != nil {
@@ -821,9 +848,22 @@ func TestSendBootstrapPipelineCall(t *testing.T) {
 		if rmsg.Finish.QuestionID != bootstrapQID {
 			t.Errorf("Received finish for question %d; want %d", rmsg.Finish.QuestionID, bootstrapQID)
 		}
-		if !rmsg.Finish.ReleaseResultCaps {
-			t.Error("Received finish that does not release bootstrap")
-		}
+		require.False(
+			t,
+			rmsg.Finish.ReleaseResultCaps,
+			"Received finish that releases bootstrap (should receive separate releasemessage)",
+		)
+	}
+
+	// 8. Release the client, read the release message.
+	client.Release()
+	{
+		rmsg, release, err := recvMessage(ctx, p2)
+		require.NoError(t, err)
+		defer release()
+		require.Equal(t, rpccp.Message_Which_release, rmsg.Which)
+		require.Equal(t, bootstrapExportID, rmsg.Release.ID)
+		require.Equal(t, uint32(1), rmsg.Release.ReferenceCount)
 	}
 }
 

rpc/question.go

@@ -15,8 +15,6 @@ type question struct {
 	c  *Conn
 	id questionID
 
-	bootstrapPromise capnp.Resolver[capnp.Client]
-
 	p       *capnp.Promise
 	release capnp.ReleaseFunc // written before resolving p
 
@@ -127,12 +125,7 @@ func (q *question) handleCancel(ctx context.Context) {
 				q.c.er.ReportError(rpcerr.Annotate(err, "send finish"))
 			}
 			close(q.finishMsgSend)
-
 			q.p.Reject(rejectErr)
-			if q.bootstrapPromise != nil {
-				q.bootstrapPromise.Reject(rejectErr)
-				q.p.ReleaseClients()
-			}
 		})
 	})
 }
@@ -278,14 +271,8 @@ func (q *question) mark(xform []capnp.PipelineOp) {
 }
 
 func (q *question) Reject(err error) {
-	if q != nil {
-		if q.bootstrapPromise != nil {
-			q.bootstrapPromise.Fulfill(capnp.ErrorClient(err))
-		}
-
-		if q.p != nil {
-			q.p.Reject(err)
-		}
+	if q != nil && q.p != nil {
+		q.p.Reject(err)
 	}
 }
 

rpc/rpc.go

@@ -308,12 +308,12 @@ func (c *Conn) Bootstrap(ctx context.Context) (bc capnp.Client) {
 		}
 		defer c.tasks.Done()
 
-		bootCtx, cancel := context.WithCancel(ctx)
 		q := c.newQuestion(capnp.Method{})
-		bc, q.bootstrapPromise = capnp.NewPromisedClient(bootstrapClient{
-			c:      q.p.Answer().Client().AddRef(),
-			cancel: cancel,
-		})
+		bc = q.p.Answer().Client().AddRef()
+		go func() {
+			q.p.ReleaseClients()
+			q.release()
+		}()
 
 		c.sendMessage(ctx, func(m rpccp.Message) error {
 			boot, err := m.NewBootstrap()
@@ -327,7 +327,7 @@ func (c *Conn) Bootstrap(ctx context.Context) (bc capnp.Client) {
 				syncutil.With(&c.lk, func() {
 					c.lk.questions[q.id] = nil
 				})
-				q.bootstrapPromise.Reject(exc.Annotate("rpc", "bootstrap", err))
+				q.p.Reject(exc.Annotate("rpc", "bootstrap", err))
 				syncutil.With(&c.lk, func() {
 					c.lk.questionID.remove(uint32(q.id))
 				})
@@ -337,40 +337,14 @@ func (c *Conn) Bootstrap(ctx context.Context) (bc capnp.Client) {
 			c.tasks.Add(1)
 			go func() {
 				defer c.tasks.Done()
-				q.handleCancel(bootCtx)
+				q.handleCancel(ctx)
 			}()
 		})
 
 		return
 	})
 }
 
-type bootstrapClient struct {
-	c      capnp.Client
-	cancel context.CancelFunc
-}
-
-func (bc bootstrapClient) String() string {
-	return "bootstrapClient{c: " + bc.c.String() + "}"
-}
-
-func (bc bootstrapClient) Send(ctx context.Context, s capnp.Send) (*capnp.Answer, capnp.ReleaseFunc) {
-	return bc.c.SendCall(ctx, s)
-}
-
-func (bc bootstrapClient) Recv(ctx context.Context, r capnp.Recv) capnp.PipelineCaller {
-	return bc.c.RecvCall(ctx, r)
-}
-
-func (bc bootstrapClient) Brand() capnp.Brand {
-	return bc.c.State().Brand
-}
-
-func (bc bootstrapClient) Shutdown() {
-	bc.cancel()
-	bc.c.Release()
-}
-
 // Close sends an abort to the remote vat and closes the underlying
 // transport.
 func (c *Conn) Close() error {
@@ -1151,23 +1125,17 @@ func (c *Conn) handleReturn(ctx context.Context, in transport.IncomingMessage) e
 			c.er.ReportError(rpcerr.Annotate(pr.err, "incoming return"))
 		}
 
-		if q.bootstrapPromise == nil && pr.err == nil {
-			// The result of the message contains actual data (not just a
-			// client or an error), so we save the ReleaseFunc for later:
+		if pr.err == nil {
+			// The result of the message contains actual data (not just
+			// an error), so we save the ReleaseFunc for later:
 			q.release = in.Release
 		}
 		// We're going to potentially block fulfilling some promises so fork
 		// off a goroutine to avoid blocking the receive loop.
 		go func() {
 			c := unlockedConn
 			q.p.Resolve(pr.result, pr.err)
-			if q.bootstrapPromise != nil {
-				q.bootstrapPromise.Fulfill(q.p.Answer().Client())
-				q.p.ReleaseClients()
-				// We can release now; root pointer of the result is a client, so the
-				// message won't be accessed:
-				in.Release()
-			} else if pr.err != nil {
+			if pr.err != nil {
 				// We can release now; the result is an error, so data from the message
 				// won't be accessed:
 				in.Release()