From 7c5521601bed01cef4d815edd59a322783193c6e Mon Sep 17 00:00:00 2001
From: "M. J. Fromberger" <michael.j.fromberger@gmail.com>
Date: Thu, 23 Sep 2021 11:57:03 -0700
Subject: [PATCH] server: Handle context termination for pending callbacks.
 (#55)

Prior to this change, calls pushed from the server to the client would ignore
the context passed to the Callback method. This meant that if the client did
not reply (for example, the client has no support for push calls), the callback
would block forever.

This change makes pending callbacks on the server respect the context that was
passed to the Callback method: When it ends, if the call has not yet received a
response, it is delivered an error from the context and resolved.

Update the documentation on Callback to recommend setting a timeout or deadline
on the context, in cases where the client might not reply.
---
 jrpc2_test.go | 29 +++++++++++++++++++++++++++++
 server.go     | 28 ++++++++++++++++++++++++++--
 2 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/jrpc2_test.go b/jrpc2_test.go
index 5ec87bb..b6cc507 100644
--- a/jrpc2_test.go
+++ b/jrpc2_test.go
@@ -730,6 +730,35 @@ func TestPushNotify(t *testing.T) {
 	}
 }
 
+// Verify that server-side callbacks can time out.
+func TestCallbackTimeout(t *testing.T) {
+	loc := server.NewLocal(handler.Map{
+		"Test": handler.New(func(ctx context.Context) error {
+			tctx, cancel := context.WithTimeout(ctx, 5*time.Millisecond)
+			defer cancel()
+			rsp, err := jrpc2.ServerFromContext(ctx).Callback(tctx, "hey", nil)
+			if err == context.DeadlineExceeded {
+				t.Logf("Callback correctly failed: %v", err)
+				return nil
+			} else if err != nil {
+				return fmt.Errorf("unexpected error: %v", err)
+			}
+			return fmt.Errorf("got rsp=%+v, want error", rsp)
+		}),
+	}, &server.LocalOptions{
+		Server: &jrpc2.ServerOptions{AllowPush: true},
+
+		// N.B. Client does not have a callback handler, so calls will be ignored
+		// and no response will be generated.
+	})
+	defer loc.Close()
+	ctx := context.Background()
+
+	if _, err := loc.Client.Call(ctx, "Test", nil); err != nil {
+		t.Errorf("Call failed: %v", err)
+	}
+}
+
 // Verify that server-side callbacks work.
 func TestPushCall(t *testing.T) {
 	loc := server.NewLocal(handler.Map{
diff --git a/server.go b/server.go
index 46b1ae7..571905d 100644
--- a/server.go
+++ b/server.go
@@ -395,8 +395,12 @@ func (s *Server) Notify(ctx context.Context, method string, params interface{})
 // client have concrete type *jrpc2.Error.
 //
 // This is a non-standard extension of JSON-RPC, and may not be supported by
-// all clients. Unless s was constructed with the AllowPush option set true,
-// this method will always report an error (ErrPushUnsupported) without sending
+// all clients. If you are not sure whether the client supports push calls, you
+// should set a deadline on ctx, otherwise the callback may block forever for a
+// client response that will never arrive.
+//
+// Unless s was constructed with the AllowPush option set true, this method
+// will always report an error (ErrPushUnsupported) without sending
 // anything. If Callback is called after the client connection is closed, it
 // returns ErrConnClosed.
 func (s *Server) Callback(ctx context.Context, method string, params interface{}) (*Response, error) {
@@ -414,6 +418,25 @@ func (s *Server) Callback(ctx context.Context, method string, params interface{}
 	return rsp, nil
 }
 
+// waitCallback blocks until pctx ends, and then if p is still waiting for a
+// response, deliver an error to the caller.
+func (s *Server) waitCallback(pctx context.Context, id string, p *Response) {
+	<-pctx.Done()
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if _, ok := s.call[id]; !ok {
+		return
+	}
+	delete(s.call, id)
+	err := pctx.Err()
+	s.log("Context ended for callback id %q, err=%v", id, err)
+
+	p.ch <- &jmessage{
+		ID: json.RawMessage(id),
+		E:  &Error{Code: code.FromError(err), Message: err.Error()},
+	}
+}
+
 func (s *Server) pushReq(ctx context.Context, wantID bool, method string, params interface{}) (rsp *Response, _ error) {
 	var bits []byte
 	if params != nil {
@@ -443,6 +466,7 @@ func (s *Server) pushReq(ctx context.Context, wantID bool, method string, params
 			cancel: func() {},
 		}
 		s.call[id] = rsp
+		go s.waitCallback(ctx, id, rsp)
 	}
 
 	s.log("Posting server %s %q %s", kind, method, string(bits))