Merge remote-tracking branch 'origin/main' into handle-resolve

Author: zenhack

capability.go

@@ -117,11 +117,6 @@ type ClientKind = struct {
 }
 
 type client struct {
-	creatorFunc  int
-	creatorFile  string
-	creatorStack string
-	creatorLine  int
-
 	mu       sync.Mutex // protects the struct
 	limiter  flowcontrol.FlowLimiter
 	h        *clientHook // nil if resolved to nil or released
@@ -156,19 +151,6 @@ type clientHook struct {
 	resolvedHook *clientHook // valid only if resolved is closed
 }
 
-func (c Client) setupLeakReporting(creatorFunc int) {
-	if clientLeakFunc == nil {
-		return
-	}
-	c.creatorFunc = creatorFunc
-	_, c.creatorFile, c.creatorLine, _ = runtime.Caller(2)
-	buf := bufferpool.Default.Get(1e6)
-	n := runtime.Stack(buf, false)
-	c.creatorStack = string(buf[:n])
-	bufferpool.Default.Put(buf)
-	c.setFinalizer()
-}
-
 // NewClient creates the first reference to a capability.
 // If hook is nil, then NewClient returns nil.
 //
@@ -187,7 +169,7 @@ func NewClient(hook ClientHook) Client {
 	}
 	h.resolvedHook = h
 	c := Client{client: &client{h: h}}
-	c.setupLeakReporting(1)
+	setupLeakReporting(c)
 	return c
 }
 
@@ -216,7 +198,7 @@ func newPromisedClient(hook ClientHook) (Client, *clientPromise) {
 		metadata:   *NewMetadata(),
 	}
 	c := Client{client: &client{h: h}}
-	c.setupLeakReporting(2)
+	setupLeakReporting(c)
 	return c, &clientPromise{h: h}
 }
 
@@ -534,7 +516,7 @@ func (c Client) AddRef() Client {
 	c.h.refs++
 	c.h.mu.Unlock()
 	d := Client{client: &client{h: c.h}}
-	d.setupLeakReporting(3)
+	setupLeakReporting(d)
 	return d
 }
 
@@ -682,7 +664,7 @@ func (ch *clientHook) isResolved() bool {
 	}
 }
 
-var clientLeakFunc func(string)
+var setupLeakReporting func(Client) = func(Client) {}
 
 // SetClientLeakFunc sets a callback for reporting Clients that went
 // out of scope without being released.  The callback is not guaranteed
@@ -691,45 +673,19 @@ var clientLeakFunc func(string)
 //
 // SetClientLeakFunc must not be called after any calls to NewClient or
 // NewPromisedClient.
-func SetClientLeakFunc(f func(msg string)) {
-	clientLeakFunc = f
-}
-
-func (c Client) setFinalizer() {
-	runtime.SetFinalizer(c.client, finalizeClient)
-}
-
-func finalizeClient(c *client) {
-	// Since there are no other references to c, then we don't have to
-	// acquire the mutex to read.
-	if c.released {
-		return
-	}
-
-	var fname string
-	switch c.creatorFunc {
-	case 1:
-		fname = "NewClient"
-	case 2:
-		fname = "NewPromisedClient"
-	case 3:
-		fname = "AddRef"
-	default:
-		fname = "<???>"
-	}
-	var msg string
-	if c.creatorFile == "" {
-		msg = "leaked client created by " + fname
-	} else {
-		msg = "leaked client created by " + fname + " on " +
-			c.creatorFile + ":" + str.Itod(c.creatorLine)
-	}
-	if c.creatorStack != "" {
-		msg += "\nCreation stack trace:\n" + c.creatorStack + "\n"
+func SetClientLeakFunc(clientLeakFunc func(msg string)) {
+	setupLeakReporting = func(c Client) {
+		buf := bufferpool.Default.Get(1e6)
+		n := runtime.Stack(buf, false)
+		stack := string(buf[:n])
+		bufferpool.Default.Put(buf)
+		runtime.SetFinalizer(c.client, func(c *client) {
+			if c.released {
+				return
+			}
+			clientLeakFunc("leaked client created at:\n\n" + stack)
+		})
 	}
-
-	// finalizeClient will only be called if clientLeakFunc != nil.
-	go clientLeakFunc(msg)
 }
 
 // A ClientPromise resolves the identity of a client created by NewPromisedClient.
@@ -830,11 +786,7 @@ func (wc *WeakClient) AddRef() (c Client, ok bool) {
 	wc.h.refs++
 	wc.h.mu.Unlock()
 	c = Client{client: &client{h: wc.h}}
-	if clientLeakFunc != nil {
-		c.creatorFunc = 3
-		_, c.creatorFile, c.creatorLine, _ = runtime.Caller(1)
-		c.setFinalizer()
-	}
+	setupLeakReporting(c)
 	return c, true
 }
 

message.go

@@ -5,6 +5,7 @@ import (
 	"encoding/binary"
 	"errors"
 	"io"
+	"net"
 	"sync"
 	"sync/atomic"
 
@@ -653,6 +654,11 @@ func (e *Encoder) Encode(m *Message) error {
 	return nil
 }
 
+func (e *Encoder) write(bufs [][]byte) error {
+	_, err := (*net.Buffers)(&bufs).WriteTo(e.w)
+	return err
+}
+
 // Marshal concatenates the segments in the message into a single byte
 // slice including framing.
 func (m *Message) Marshal() ([]byte, error) {

message_go18.go

@@ -918,11 +918,7 @@ func TestRecvBootstrapCall(t *testing.T) {
 	})
 	defer func() {
 		finishTest(t, conn, p2)
-		select {
-		case <-srvShutdown:
-		default:
-			t.Error("Bootstrap client still alive after Close returned")
-		}
+		<-srvShutdown // Hangs if bootstrap client is never shut down.
 	}()
 
 	ctx := context.Background()
@@ -1233,11 +1229,7 @@ func TestRecvBootstrapPipelineCall(t *testing.T) {
 	})
 	defer func() {
 		finishTest(t, conn, p2)
-		select {
-		case <-srvShutdown:
-		default:
-			t.Error("Bootstrap client still alive after Close returned")
-		}
+		<-srvShutdown // Will hang if closing does not shut down the client.
 	}()
 	ctx := context.Background()
 
@@ -1699,13 +1691,9 @@ func TestRecvCancel(t *testing.T) {
 	}
 
 	// 8. Verify that returned capability was shut down.
-	// There's no guarantee when the release/shutdown will happen, other
-	// than it will be released before Close returns.
-	select {
-	case <-retcapShutdown:
-	default:
-		t.Error("returned capability was not shut down")
-	}
+	// There's no guarantee exactly when the release/shutdown will happen,
+	// but Close should trigger it. Otherwise, this will hang:
+	<-retcapShutdown
 }
 
 // TestSendCancel makes a call, cancels the Context, then checks to

message_other.go

@@ -153,15 +153,15 @@ func TestResolveUnimplementedDrop(t *testing.T) {
 		assert.Equal(t, rpccp.CapDescriptor_Which_senderPromise, desc.Which)
 		emptyExportID = desc.SenderPromise
 	}
-	// 7. Fulfill promise
+	// 6. Fulfill promise
 	{
 		pp := testcapnp.Empty_ServerToClient(emptyShutdowner{
 			onShutdown: onShutdown,
 		})
 		r.Fulfill(pp)
 		pp.Release()
 	}
-	// 8. Receive resolve, send unimplemented
+	// 7. Receive resolve, send unimplemented
 	{
 		rmsg, release, err := recvMessage(ctx, p2)
 		assert.NoError(t, err)
@@ -176,12 +176,12 @@ func TestResolveUnimplementedDrop(t *testing.T) {
 			Unimplemented: rmsg,
 		}))
 	}
-	// 9. Drop the promise on our side. Otherwise it will stay alive because of
+	// 8. Drop the promise on our side. Otherwise it will stay alive because of
 	// the bootstrap interface:
 	{
 		p.Release()
 	}
-	// 6. Send finish
+	// 9. Send finish
 	{
 		assert.NoError(t, sendMessage(ctx, p2, &rpcMessage{
 			Which: rpccp.Message_Which_finish,

rpc/level0_test.go

@@ -69,7 +69,7 @@ func (c *Call) Go() {
 		return
 	}
 	c.acked = true
-	go c.srv.handleCalls(c.srv.handleCallsCtx)
+	go c.srv.handleCalls()
 }
 
 // Shutdowner is the interface that wraps the Shutdown method.
@@ -84,15 +84,6 @@ type Server struct {
 	brand    any
 	shutdown Shutdowner
 
-	// Cancels handleCallsCtx
-	cancelHandleCalls context.CancelFunc
-
-	// Context used by the goroutine running handleCalls(). Note
-	// the calls themselves will have different contexts, which
-	// are not children of this context, but are supplied by
-	// start().
-	handleCallsCtx context.Context
-
 	// wg is incremented each time a method is queued, and
 	// decremented after it is handled.
 	wg sync.WaitGroup
@@ -114,19 +105,15 @@ func (s *Server) String() string {
 // guarantees message delivery order by blocking each call on the
 // return of the previous call or a call to Call.Go.
 func New(methods []Method, brand any, shutdown Shutdowner) *Server {
-	ctx, cancel := context.WithCancel(context.Background())
-
 	srv := &Server{
-		methods:           make(sortedMethods, len(methods)),
-		brand:             brand,
-		shutdown:          shutdown,
-		callQueue:         mpsc.New[*Call](),
-		cancelHandleCalls: cancel,
-		handleCallsCtx:    ctx,
+		methods:   make(sortedMethods, len(methods)),
+		brand:     brand,
+		shutdown:  shutdown,
+		callQueue: mpsc.New[*Call](),
 	}
 	copy(srv.methods, methods)
 	sort.Sort(srv.methods)
-	go srv.handleCalls(ctx)
+	go srv.handleCalls()
 	return srv
 }
 
@@ -171,38 +158,20 @@ func (srv *Server) Recv(ctx context.Context, r capnp.Recv) capnp.PipelineCaller
 	return srv.start(ctx, mm, r)
 }
 
-func (srv *Server) handleCalls(ctx context.Context) {
+func (srv *Server) handleCalls() {
+	ctx := context.Background()
 	for {
 		call, err := srv.callQueue.Recv(ctx)
 		if err != nil {
-			// Context has been canceled; drain the rest of the queue,
-			// invoking handleCall() with the cancelled context to
-			// trigger cleanup.
-			var ok bool
-			call, ok = srv.callQueue.TryRecv()
-			if !ok {
-				return
+			// Queue closed; wait for outstanding calls and shut down.
+			if srv.shutdown != nil {
+				srv.wg.Wait()
+				srv.shutdown.Shutdown()
 			}
+			return
 		}
 
-		// The context for the individual call is not necessarily
-		// related to the context managing the server's lifetime
-		// (ctx); we need to monitor both and pass the call a
-		// context that will be canceled if *either* context is
-		// cancelled.
-		callCtx, cancelCall := context.WithCancel(call.ctx)
-		go func() {
-			defer cancelCall()
-			select {
-			case <-callCtx.Done():
-			case <-ctx.Done():
-			}
-		}()
-		func() {
-			defer cancelCall()
-			srv.handleCall(callCtx, call)
-		}()
-
+		srv.handleCall(call)
 		if call.acked {
 			// Another goroutine has taken over; time
 			// to retire.
@@ -211,10 +180,10 @@ func (srv *Server) handleCalls(ctx context.Context) {
 	}
 }
 
-func (srv *Server) handleCall(ctx context.Context, c *Call) {
+func (srv *Server) handleCall(c *Call) {
 	defer srv.wg.Done()
 
-	err := c.method.Impl(ctx, c)
+	err := c.method.Impl(c.ctx, c)
 
 	c.recv.ReleaseArgs()
 	c.recv.Returner.PrepareReturn(err)
@@ -246,15 +215,11 @@ func (srv *Server) Brand() capnp.Brand {
 	return capnp.Brand{Value: serverBrand{srv.brand}}
 }
 
-// Shutdown waits for ongoing calls to finish and calls Shutdown on the
-// Shutdowner passed into NewServer.  Shutdown must not be called more
-// than once.
+// Shutdown arranges for Shutdown to be called on the Shutdowner passed
+// into NewServer after outstanding all calls have been serviced.
+// Shutdown must not be called more than once.
 func (srv *Server) Shutdown() {
-	srv.cancelHandleCalls()
-	srv.wg.Wait()
-	if srv.shutdown != nil {
-		srv.shutdown.Shutdown()
-	}
+	srv.callQueue.Close()
 }
 
 // IsServer reports whether a brand returned by capnp.Client.Brand