Improve type safety wrt locking.

rpc/answer.go

@@ -198,8 +198,8 @@ func (ans *answer) Return(e error) {
 
 	var err error
 
-	syncutil.With(&ans.c.lk, func() {
-		err = ans.sendReturn(rl)
+	ans.c.withLocked(func(c *lockedConn) {
+		err = ans.sendReturn(c, rl)
 	})
 	ans.c.tasks.Done() // added by handleCall
 
@@ -217,14 +217,19 @@ func (ans *answer) Return(e error) {
 // Finish with releaseResultCaps set to true, then sendReturn returns
 // the number of references to be subtracted from each export.
 //
-// The caller MUST be holding onto ans.c.lk. sendReturn MUST NOT be
-// called if sendException was previously called.
-func (ans *answer) sendReturn(rl *releaseList) error {
+// The caller MUST be holding onto ans.c.lk, and the lockedConn parameter
+// must be equal to ans.c (it exists to make it hard to forget to acquire
+// the lock, per usual).
+//
+// sendReturn MUST NOT be called if sendException was previously called.
+func (ans *answer) sendReturn(c *lockedConn, rl *releaseList) error {
+	c.assertIs(ans.c)
+
 	ans.pcall = nil
 	ans.flags |= resultsReady
 
 	var err error
-	ans.exportRefs, err = ans.c.fillPayloadCapTable(ans.results)
+	ans.exportRefs, err = c.fillPayloadCapTable(ans.results)
 	if err != nil {
 		// We're not going to send the message after all, so don't forget to release it.
 		ans.msgReleaser.Decr()

rpc/export.go

@@ -24,16 +24,16 @@ type exportIDKey struct {
 	Conn *Conn
 }
 
-func (c *Conn) findExportID(m *capnp.Metadata) (_ exportID, ok bool) {
-	maybeID, ok := m.Get(exportIDKey{c})
+func (c *lockedConn) findExportID(m *capnp.Metadata) (_ exportID, ok bool) {
+	maybeID, ok := m.Get(exportIDKey{(*Conn)(c)})
 	if ok {
 		return maybeID.(exportID), true
 	}
 	return 0, false
 }
 
-func (c *Conn) setExportID(m *capnp.Metadata, id exportID) {
-	m.Put(exportIDKey{c}, id)
+func (c *lockedConn) setExportID(m *capnp.Metadata, id exportID) {
+	m.Put(exportIDKey{(*Conn)(c)}, id)
 }
 
 func (c *Conn) clearExportID(m *capnp.Metadata) {
@@ -101,17 +101,16 @@ func (c *Conn) releaseExportRefs(rl *releaseList, refs map[exportID]uint32) erro
 }
 
 // sendCap writes a capability descriptor, returning an export ID if
-// this vat is hosting the capability.  The caller must be holding
-// onto c.mu.
-func (c *Conn) sendCap(d rpccp.CapDescriptor, client capnp.Client) (_ exportID, isExport bool, _ error) {
+// this vat is hosting the capability.
+func (c *lockedConn) sendCap(d rpccp.CapDescriptor, client capnp.Client) (_ exportID, isExport bool, _ error) {
 	if !client.IsValid() {
 		d.SetNone()
 		return 0, false, nil
 	}
 
 	state := client.State()
 	bv := state.Brand.Value
-	if ic, ok := bv.(*importClient); ok && ic.c == c {
+	if ic, ok := bv.(*importClient); ok && ic.c == (*Conn)(c) {
 		if ent := c.lk.imports[ic.id]; ent != nil && ent.generation == ic.generation {
 			d.SetReceiverHosted(uint32(ic.id))
 			return 0, false, nil
@@ -120,7 +119,7 @@ func (c *Conn) sendCap(d rpccp.CapDescriptor, client capnp.Client) (_ exportID,
 
 	if pc, ok := bv.(capnp.PipelineClient); ok {
 		q, ok := c.getAnswerQuestion(pc.Answer())
-		if ok && q.c == c {
+		if ok && q.c == (*Conn)(c) {
 			pcTrans := pc.Transform()
 			pa, err := d.NewReceiverAnswer()
 			if err != nil {
@@ -170,9 +169,7 @@ func (c *Conn) sendCap(d rpccp.CapDescriptor, client capnp.Client) (_ exportID,
 // fillPayloadCapTable adds descriptors of payload's message's
 // capabilities into payload's capability table and returns the
 // reference counts that have been added to the exports table.
-//
-// The caller must be holding onto c.mu.
-func (c *Conn) fillPayloadCapTable(payload rpccp.Payload) (map[exportID]uint32, error) {
+func (c *lockedConn) fillPayloadCapTable(payload rpccp.Payload) (map[exportID]uint32, error) {
 	if !payload.IsValid() {
 		return nil, nil
 	}

rpc/import.go

@@ -84,53 +84,50 @@ type importClient struct {
 }
 
 func (ic *importClient) Send(ctx context.Context, s capnp.Send) (*capnp.Answer, capnp.ReleaseFunc) {
-	ic.c.lk.Lock()
-	defer ic.c.lk.Unlock()
-
-	if !ic.c.startTask() {
-		return capnp.ErrorAnswer(s.Method, ExcClosed), func() {}
-	}
-	defer ic.c.tasks.Done()
-	ent := ic.c.lk.imports[ic.id]
-	if ent == nil || ic.generation != ent.generation {
-		return capnp.ErrorAnswer(s.Method, rpcerr.Disconnectedf("send on closed import")), func() {}
-	}
-	q := ic.c.newQuestion(s.Method)
-
-	// Send call message.
-	ic.c.sendMessage(ctx, func(m rpccp.Message) error {
-		return ic.c.newImportCallMessage(m, ic.id, q.id, s)
-	}, func(err error) {
-		if err != nil {
-			syncutil.With(&ic.c.lk, func() {
-				ic.c.lk.questions[q.id] = nil
-			})
-			q.p.Reject(rpcerr.Failedf("send message: %w", err))
-			syncutil.With(&ic.c.lk, func() {
-				ic.c.lk.questionID.remove(uint32(q.id))
-			})
-			return
+	return withLockedConn2(ic.c, func(c *lockedConn) (*capnp.Answer, capnp.ReleaseFunc) {
+		if !c.startTask() {
+			return capnp.ErrorAnswer(s.Method, ExcClosed), func() {}
+		}
+		defer c.tasks.Done()
+		ent := c.lk.imports[ic.id]
+		if ent == nil || ic.generation != ent.generation {
+			return capnp.ErrorAnswer(s.Method, rpcerr.Disconnectedf("send on closed import")), func() {}
+		}
+		q := c.newQuestion(s.Method)
+
+		// Send call message.
+		c.sendMessage(ctx, func(m rpccp.Message) error {
+			return c.newImportCallMessage(m, ic.id, q.id, s)
+		}, func(err error) {
+			if err != nil {
+				syncutil.With(&ic.c.lk, func() {
+					ic.c.lk.questions[q.id] = nil
+				})
+				q.p.Reject(rpcerr.Failedf("send message: %w", err))
+				syncutil.With(&ic.c.lk, func() {
+					ic.c.lk.questionID.remove(uint32(q.id))
+				})
+				return
+			}
+
+			q.c.tasks.Add(1)
+			go func() {
+				defer q.c.tasks.Done()
+				q.handleCancel(ctx)
+			}()
+		})
+
+		ans := q.p.Answer()
+		return ans, func() {
+			<-ans.Done()
+			q.p.ReleaseClients()
+			q.release()
 		}
-
-		q.c.tasks.Add(1)
-		go func() {
-			defer q.c.tasks.Done()
-			q.handleCancel(ctx)
-		}()
 	})
-
-	ans := q.p.Answer()
-	return ans, func() {
-		<-ans.Done()
-		q.p.ReleaseClients()
-		q.release()
-	}
 }
 
 // newImportCallMessage builds a Call message targeted to an import.
-//
-// The caller MUST hold c.mu.
-func (c *Conn) newImportCallMessage(msg rpccp.Message, imp importID, qid questionID, s capnp.Send) error {
+func (c *lockedConn) newImportCallMessage(msg rpccp.Message, imp importID, qid questionID, s capnp.Send) error {
 	call, err := msg.NewCall()
 	if err != nil {
 		return rpcerr.Failedf("build call message: %w", err)
@@ -214,29 +211,28 @@ func (ic *importClient) Brand() capnp.Brand {
 }
 
 func (ic *importClient) Shutdown() {
-	ic.c.lk.Lock()
-	defer ic.c.lk.Unlock()
-
-	if !ic.c.startTask() {
-		return
-	}
-	defer ic.c.tasks.Done()
+	ic.c.withLocked(func(c *lockedConn) {
+		if !c.startTask() {
+			return
+		}
+		defer c.tasks.Done()
 
-	ent := ic.c.lk.imports[ic.id]
-	if ic.generation != ent.generation {
-		// A new reference was added concurrently with the Shutdown.  See
-		// impent.generation documentation for an explanation.
-		return
-	}
-	delete(ic.c.lk.imports, ic.id)
-	ic.c.sendMessage(ic.c.bgctx, func(msg rpccp.Message) error {
-		rel, err := msg.NewRelease()
-		if err == nil {
-			rel.SetId(uint32(ic.id))
-			rel.SetReferenceCount(uint32(ent.wireRefs))
+		ent := c.lk.imports[ic.id]
+		if ic.generation != ent.generation {
+			// A new reference was added concurrently with the Shutdown.  See
+			// impent.generation documentation for an explanation.
+			return
 		}
-		return err
-	}, func(err error) {
-		ic.c.er.ReportError(rpcerr.Annotate(err, "send release"))
+		delete(ic.c.lk.imports, ic.id)
+		c.sendMessage(c.bgctx, func(msg rpccp.Message) error {
+			rel, err := msg.NewRelease()
+			if err == nil {
+				rel.SetId(uint32(ic.id))
+				rel.SetReferenceCount(uint32(ent.wireRefs))
+			}
+			return err
+		}, func(err error) {
+			ic.c.er.ReportError(rpcerr.Annotate(err, "send release"))
+		})
 	})
 }