Repository: creachadair/jrpc2
Pull request #29 (2020-09-05 19:16:30)

Add a cancellation callback to ClientOptions.
When set, the new OnCancel callback overrides the client's built-in support for
sending rpc.cancel notifications. The hook can implement the same functionality
or take other actions using the client and the cancelled response.

Related changes:
- Add tests for OnCancel hook handling.
- Fix up some comment typos.


Patch:
diff --git a/base.go b/base.go
index 8e501a9..fda0873 100644
--- a/base.go
+++ b/base.go
@@ -167,7 +167,7 @@ func (r *Response) MarshalJSON() ([]byte, error) {
 	})
 }
 
-// wait blocks until p is complete. It is safe to call this multiple times and
+// wait blocks until r is complete. It is safe to call this multiple times and
 // from concurrent goroutines.
 func (r *Response) wait() {
 	raw, ok := <-r.ch
diff --git a/client.go b/client.go
index 780b5b1..5c54cc0 100644
--- a/client.go
+++ b/client.go
@@ -22,6 +22,7 @@ type Client struct {
 	enctx encoder
 	snote func(*jmessage)
 	scall func(*jmessage) ([]byte, error)
+	chook func(*Client, *Response)
 
 	allow1 bool // tolerate v1 replies with no version marker
 	allowC bool // send rpc.cancel when a request context ends
@@ -43,6 +44,7 @@ func NewClient(ch channel.Channel, opts *ClientOptions) *Client {
 		enctx:  opts.encodeContext(),
 		snote:  opts.handleNotification(),
 		scall:  opts.handleCallback(),
+		chook:  opts.handleCancel(),
 
 		// Lock-protected fields
 		ch:      ch,
@@ -215,7 +217,7 @@ func (c *Client) send(ctx context.Context, reqs jmessages) ([]*Response, error)
 	}
 
 	// Now that we have sent them, record the requests for which we are awaiting
-	// replies. We do this after transsmission so that an error in sending does
+	// replies. We do this after transmission so that an error in sending does
 	// not leave us with zombies that will never be fulfilled.
 	for i, p := range pends {
 		c.pending[p.id] = p
@@ -259,7 +261,13 @@ func (c *Client) waitComplete(pctx context.Context, id string, p *Response) {
 
 	// Inform the server, best effort only. N.B. Use a background context here,
 	// as the original context has ended by the time we get here.
-	if c.allowC {
+	if c.chook != nil {
+		cleanup = func() {
+			p.wait() // ensure the response has settled
+			c.log("Calling OnCancel for id %q", id)
+			c.chook(c, p)
+		}
+	} else if c.allowC {
 		cleanup = func() {
 			c.log("Sending rpc.cancel for id %q to the server", id)
 			c.Notify(context.Background(), rpcCancel, []json.RawMessage{json.RawMessage(id)})
diff --git a/opts.go b/opts.go
index 7eb626d..e5b6a98 100644
--- a/opts.go
+++ b/opts.go
@@ -157,6 +157,15 @@ type ClientOptions struct {
 	// invocation of this callback will be active at a time.
 	// Server callbacks are a non-standard extension of JSON-RPC.
 	OnCallback func(context.Context, *Request) (interface{}, error)
+
+	// If set, this function is called when the context for a request terminates.
+	// The function receives the client and the response that was cancelled.
+	// The hook can obtain the ID and error value from rsp.
+	//
+	// Setting this option disables the default rpc.cancel handling (as DisableCancel).
+	// Note that the hook does not receive the client context, which has already
+	// ended by the time the hook is called.
+	OnCancel func(cli *Client, rsp *Response)
 }
 
 func (c *ClientOptions) logger() logger {
@@ -189,6 +198,13 @@ func (c *ClientOptions) handleNotification() func(*jmessage) {
 	return func(req *jmessage) { h(&Request{method: req.M, params: req.P}) }
 }
 
+func (c *ClientOptions) handleCancel() func(*Client, *Response) {
+	if c == nil {
+		return nil
+	}
+	return c.OnCancel
+}
+
 func (c *ClientOptions) handleCallback() func(*jmessage) ([]byte, error) {
 	if c == nil || c.OnCallback == nil {
 		return nil


Test patch:
diff --git a/jrpc2_test.go b/jrpc2_test.go
index c857181..c258153 100644
--- a/jrpc2_test.go
+++ b/jrpc2_test.go
@@ -772,6 +772,59 @@ func TestDeadServerPush(t *testing.T) {
 	}
 }
 
+// Verify that an OnCancel hook is called when expected.
+func TestOnCancel(t *testing.T) {
+	// Set up a plumbing context so the test can unblock the server.
+	sctx, cancelServer := context.WithCancel(context.Background())
+	defer cancelServer()
+	loc := server.NewLocal(handler.Map{
+		// Block until explicitly cancelled, via sctx.
+		"Stall": handler.New(func(_ context.Context) error {
+			select {
+			case <-sctx.Done():
+				t.Logf("Server unblocked; returning err=%v", sctx.Err())
+				return sctx.Err()
+			case <-time.After(10 * time.Second): // shouldn't happen
+				t.Error("Timeout waiting for server cancellation")
+			}
+			return nil
+		}),
+
+		// Verify that setting the cancellation hook prevents the client from
+		// sending the default rpc.cancel notification.
+		"rpc.cancel": handler.New(func(ctx context.Context, ids json.RawMessage) error {
+			t.Errorf("Server-side rpc.cancel unexpectedly called: %s", string(ids))
+			return nil
+		}),
+	}, &server.LocalOptions{
+		// Disable handling of built-in methods on the server.
+		Server: &jrpc2.ServerOptions{DisableBuiltin: true},
+		Client: &jrpc2.ClientOptions{
+			OnCancel: func(cli *jrpc2.Client, rsp *jrpc2.Response) {
+				t.Logf("OnCancel hook called with id=%q, err=%v", rsp.ID(), rsp.Error())
+				cancelServer()
+			},
+		},
+	})
+
+	// Call a method on the server that will stall until cancelServer is called.
+	// On the client side, set a deadline to expire the caller's context.
+	// The cancellation hook will unblock the server.
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer cancel()
+
+	got, err := loc.Client.Call(ctx, "Stall", nil)
+	if err == nil {
+		t.Errorf("Stall: got %+v, wanted error", got)
+	} else if err != context.DeadlineExceeded {
+		t.Errorf("Stall: got error %v, want %v", err, context.Canceled)
+	}
+	loc.Client.Close()
+	if err := loc.Server.Wait(); err != nil {
+		t.Errorf("Server exit status: %v", err)
+	}
+}
+
 // Verify that the context encoding/decoding hooks work.
 func TestContextPlumbing(t *testing.T) {
 	want := time.Now().Add(10 * time.Second)