Repository: creachadair/jrpc2
Pull request #103 (2023-07-12 18:06:57)

client: improve handling of stopped clients
A client can be stopped by calling its Close method, or by a failure of its
connection to the server. The user of the client does not have a way to handle
this gracefully.

Add an IsStopped method to the *jrpc2.Client, allowing the caller to check
whether the client is still alive explicitly.

Add an optional OnStop callback that will be invoked when the client closes,
allowing the caller to handle cleanup tasks automatically.

Updates #102


pbedat: This is perfect! Thanks!

Patch:
diff --git a/client.go b/client.go
index 0d0b7f9..c3cbc20 100644
--- a/client.go
+++ b/client.go
@@ -22,6 +22,7 @@ type Client struct {
 	snote func(*jmessage)
 	scall func(context.Context, *jmessage) []byte
 	chook func(*Client, *Response)
+	shook func(*Client, error)
 
 	cbctx    context.Context    // terminates when the client is closed
 	cbcancel context.CancelFunc // cancels cbctx
@@ -42,6 +43,7 @@ func NewClient(ch channel.Channel, opts *ClientOptions) *Client {
 		snote: opts.handleNotification(),
 		scall: opts.handleCallback(),
 		chook: opts.handleCancel(),
+		shook: opts.handleStop(),
 
 		cbctx:    cbctx,
 		cbcancel: cbcancel,
@@ -82,7 +84,7 @@ func (c *Client) accept(ch receiver) error {
 			c.log("Decoding error: %v", err)
 		}
 		c.mu.Lock()
-		c.stopLocked(err)
+		defer c.stopLocked(err)()
 		c.mu.Unlock()
 		return err
 	}
@@ -377,7 +379,7 @@ func (c *Client) Notify(ctx context.Context, method string, params any) error {
 // Close shuts down the client, terminating any pending in-flight requests.
 func (c *Client) Close() error {
 	c.mu.Lock()
-	c.stopLocked(errClientStopped)
+	defer c.stopLocked(errClientStopped)()
 	c.mu.Unlock()
 	c.done.Wait()
 
@@ -388,6 +390,14 @@ func (c *Client) Close() error {
 	return c.err
 }
 
+// IsStopped reports whether the client has been stopped, either by a call to
+// its Close method or by a failure of the channel to the server.
+func (c *Client) IsStopped() bool {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.err != nil
+}
+
 func isUninteresting(err error) bool {
 	return err == io.EOF || channel.IsErrClosing(err) || err == errClientStopped
 }
@@ -395,9 +405,10 @@ func isUninteresting(err error) bool {
 // stopLocked closes down the reader for c and records err as its final state.
 // The caller must hold c.mu. If multiple callers invoke stop, only the first
 // will successfully record its error status.
-func (c *Client) stopLocked(err error) {
+// The caller should defer a call to the returned function.
+func (c *Client) stopLocked(err error) func() {
 	if c.ch == nil {
-		return // nothing is running
+		return func() {} // nothing is running
 	}
 	c.ch.Close()
 
@@ -411,6 +422,7 @@ func (c *Client) stopLocked(err error) {
 
 	c.err = err
 	c.ch = nil
+	return func() { c.shook(c, err) }
 }
 
 // marshalParams validates and marshals params to JSON for a request.  The
diff --git a/opts.go b/opts.go
index 8c4d551..b1da3b3 100644
--- a/opts.go
+++ b/opts.go
@@ -120,6 +120,11 @@ type ClientOptions struct {
 	// Note that the hook does not receive the request context, which has
 	// already ended by the time the hook is called.
 	OnCancel func(cli *Client, rsp *Response)
+
+	// If set, this function is called when the client is stopped, either by
+	// calling its Close method or by disconnection of its channel.  The
+	// arguments are the client itself and the error that caused it to stop.
+	OnStop func(cli *Client, err error)
 }
 
 func (c *ClientOptions) logFunc() func(string, ...any) {
@@ -144,6 +149,13 @@ func (c *ClientOptions) handleCancel() func(*Client, *Response) {
 	return c.OnCancel
 }
 
+func (c *ClientOptions) handleStop() func(*Client, error) {
+	if c == nil || c.OnStop == nil {
+		return func(*Client, error) {}
+	}
+	return c.OnStop
+}
+
 func (c *ClientOptions) handleCallback() func(context.Context, *jmessage) []byte {
 	if c == nil || c.OnCallback == nil {
 		return nil


Test patch:
diff --git a/jrpc2_test.go b/jrpc2_test.go
index 7431b85..539319a 100644
--- a/jrpc2_test.go
+++ b/jrpc2_test.go
@@ -1348,3 +1348,32 @@ func TestServer_newContext(t *testing.T) {
 		t.Errorf("Call failed: %v", err)
 	}
 }
+
+func TestClient_IsStopped(t *testing.T) {
+	clich := make(chan *jrpc2.Client, 1)
+	loc := server.NewLocal(handler.Map{}, &server.LocalOptions{
+		Client: &jrpc2.ClientOptions{
+			OnStop: func(cli *jrpc2.Client, _ error) { clich <- cli },
+		},
+	})
+
+	// A client should not be initially stopped.
+	if loc.Client.IsStopped() {
+		t.Error("IsStopped is true, want false")
+	}
+
+	// Closing the client should trigger its callback and leave it stopped.
+	loc.Client.Close()
+	loc.Server.Wait() // ensure callbacks have settled
+	close(clich)      // unblock
+
+	// A closed client should appear stopped.
+	if !loc.Client.IsStopped() {
+		t.Error("IsStopped is false, want true")
+	}
+
+	// The stop hook should have populated the same client value.
+	if cli := <-clich; cli != loc.Client {
+		t.Errorf("OnStop did not work: got %v, want %v", cli, loc.Client)
+	}
+}