Repository: creachadair/jrpc2 (go)

Issue:
Push notifications/requests from outside handlers
In [our language server](https://github.com/hashicorp/terraform-ls) we spin up a few routines as part of compatible `server.Service`

https://github.com/creachadair/jrpc2/blob/446228c04965bd4328cac6e1a43af81198ddf998/server/loop.go#L13-L20

These routines can be long-lived, but never longer than the lifetime of a session - hence we spin them up as part of `Assigner` and stop as part of `Finish`.

In particular we have three things running outside of handlers currently:

 - watcher (watching the filesystem for relevant changes)
 - walker (walking the filesystem on `initalize` to understand relationships between different folders user has in the hierarchy)
 - "module loader" - this is basically a task queue with a worker pool which executes various tasks on the background without blocking any handler response, reflecting that these tasks may take a while to finish and there can be many queued
    - config parsing
    - obtaining version of Terraform
    - parsing metadata relevant to open files/folders

This all works aok, but we'd also like to report progress on some of these operations, i.e. send `$/progress` notifications back to the client.

I think the best way of solving that might be introduction of a new function which can extract the `*Server` out of the context as part of a handler. This then allows for it to be persisted somewhere outside of the short-lived handler context.

```go
func ServerFromContext(ctx context.Context) (*Server, bool) {
  s, ok := ctx.Value(serverKey{}).(*Server)
  return s, ok
}
```
alternatively instead of exposing the whole `*Server` the function could extract an interface such as this one
```go
type ServerPusher interface {
  Callback(ctx context.Context, method string, params interface{}) (*Response, error)
  Notify(ctx context.Context, method string, params interface{}) error
}
```
just to discourage from using other parts of the API in that context. Obviously the interface can easily be just casted back to `*Server`, but it would still communicate the intentions better?

--- 

I would be happy to send a PR - let me know what you think.

Pull request description:
Add a function to expose a Server to its handlers.
Add a top-level `ServerFromContext` function that recovers the `*Server` value from
the context passed to a handler.

Related changes:
- Clean up some redundant flag checks in the context handlers.
- Rewrite the existing context handlers in terms of this new function.
- Add a top-level `CancelRequest` method on the `Server`.

Fixes #38.



Patch:
diff --git a/ctx.go b/ctx.go
index 17219c8..da69d09 100644
--- a/ctx.go
+++ b/ctx.go
@@ -2,18 +2,16 @@ package jrpc2
 
 import (
 	"context"
-	"encoding/json"
 	"errors"
 
 	"github.com/creachadair/jrpc2/metrics"
 )
 
-// ServerMetrics returns the server metrics collector associated with the given
-// context, or nil if ctx does not have a collector attached.  The context
-// passed to a handler by *jrpc2.Server will include this value.
-func ServerMetrics(ctx context.Context) *metrics.M {
-	return ctx.Value(serverKey{}).(*Server).metrics
-}
+// ServerMetrics returns the server metrics collector. If the server does not
+// have a metrics collector, it returns nil, which is ready for use but
+// discards all posted metrics.  This function is for use by handlers, and will
+// panic for a non-handler context.
+func ServerMetrics(ctx context.Context) *metrics.M { return ServerFromContext(ctx).metrics }
 
 // InboundRequest returns the inbound request associated with the given
 // context, or nil if ctx does not have an inbound request. The context passed
@@ -32,40 +30,38 @@ func InboundRequest(ctx context.Context) *Request {
 
 type inboundRequestKey struct{}
 
-// PushNotify posts a server notification to the client. If ctx does not
-// contain a server notifier, this reports ErrPushUnsupported. The context
-// passed to the handler by *jrpc2.Server will support notifications if the
-// server was constructed with the AllowPush option set true.
+// PushNotify posts a server notification to the client. If the server does not
+// have push enabled (via the AllowPush option), it reports ErrPushUnsupported.
+// This function is for use by handlers, and will panic for a non-handler context.
 func PushNotify(ctx context.Context, method string, params interface{}) error {
-	s := ctx.Value(serverKey{}).(*Server)
-	if !s.allowP {
-		return ErrPushUnsupported
-	}
-	return s.Notify(ctx, method, params)
+	return ServerFromContext(ctx).Notify(ctx, method, params)
 }
 
-// PushCall posts a server call to the client. If ctx does not contain a server
-// caller, this reports ErrPushUnsupported. The context passed to the handler
-// by *jrpc2.Server will support callbacks if the server was constructed with
-// the AllowPush option set true.
+// PushCall posts a server call to the client. If the server does not have push
+// enabled (via the AllowPush option), it reports ErrPushUnsupported.
+// This function is for use by handlers, and will panic for a non-handler context.
 //
 // A successful callback reports a nil error and a non-nil response. Errors
-// returned by the client have concrete type *jrpc2.Error.
+// reported by the client have concrete type *jrpc2.Error.
 func PushCall(ctx context.Context, method string, params interface{}) (*Response, error) {
-	s := ctx.Value(serverKey{}).(*Server)
-	if !s.allowP {
-		return nil, ErrPushUnsupported
-	}
-	return s.Callback(ctx, method, params)
+	return ServerFromContext(ctx).Callback(ctx, method, params)
 }
 
 // CancelRequest requests the server associated with ctx to cancel the pending
 // or in-flight request with the specified ID.  If no request exists with that
 // ID, this is a no-op without error.
-func CancelRequest(ctx context.Context, id string) {
-	s := ctx.Value(serverKey{}).(*Server)
-	s.cancelRequests(ctx, []json.RawMessage{json.RawMessage(id)})
-}
+// This function is for use by handlers, and will panic for a non-handler context.
+func CancelRequest(ctx context.Context, id string) { ServerFromContext(ctx).CancelRequest(id) }
+
+// ServerFromContext returns the server associated with the given context.
+// This will be populated on the context passed to request handlers.
+// This function is for use by handlers, and will panic for a non-handler context.
+//
+// It is safe to retain the server and invoke its methods beyond the lifetime
+// of the context from which it was extracted; however, a handler must not
+// block on the Wait or WaitStatus methods of the server, as the server will
+// deadlock waiting for the handler to return.
+func ServerFromContext(ctx context.Context) *Server { return ctx.Value(serverKey{}).(*Server) }
 
 type serverKey struct{}
 
diff --git a/special.go b/special.go
index c17875c..715ae55 100644
--- a/special.go
+++ b/special.go
@@ -22,11 +22,11 @@ func (s *Server) handleRPCCancel(ctx context.Context, req *Request) (interface{}
 	if err := req.UnmarshalParams(&ids); err != nil {
 		return nil, err
 	}
-	s.cancelRequests(ctx, ids)
+	s.cancelRequests(ids)
 	return nil, nil
 }
 
-func (s *Server) cancelRequests(ctx context.Context, ids []json.RawMessage) {
+func (s *Server) cancelRequests(ids []json.RawMessage) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	for _, raw := range ids {
@@ -37,6 +37,12 @@ func (s *Server) cancelRequests(ctx context.Context, ids []json.RawMessage) {
 	}
 }
 
+// CancelRequest instructs s to cancel the pending or in-flight request with
+// the specified ID. If no request exists with that ID, this is a no-op.
+func (s *Server) CancelRequest(id string) {
+	s.cancelRequests([]json.RawMessage{json.RawMessage(id)})
+}
+
 // methodFunc is a replication of handler.Func redeclared to avert a cycle.
 type methodFunc func(context.Context, *Request) (interface{}, error)
 


Test patch:
diff --git a/jrpc2_test.go b/jrpc2_test.go
index 190d94e..1f2d283 100644
--- a/jrpc2_test.go
+++ b/jrpc2_test.go
@@ -1129,3 +1129,22 @@ func TestStrictFields(t *testing.T) {
 		}
 	})
 }
+
+func TestServerFromContext(t *testing.T) {
+	var got *jrpc2.Server
+	loc := server.NewLocal(handler.Map{
+		"Test": handler.New(func(ctx context.Context) error {
+			got = jrpc2.ServerFromContext(ctx)
+			return nil
+		}),
+	}, nil)
+	if _, err := loc.Client.Call(context.Background(), "Test", nil); err != nil {
+		t.Fatalf("Call failed: %v", err)
+	}
+	if err := loc.Close(); err != nil {
+		t.Errorf("Close failed: %v", err)
+	}
+	if got != loc.Server {
+		t.Errorf("ServerFromContext: got %p, want %p", got, loc.Server)
+	}
+}