[ADDED] ContextHandler helper in package micro (#1215)

micro/example_test.go

@@ -14,6 +14,7 @@
 package micro_test
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"reflect"
@@ -313,6 +314,37 @@ func ExampleService_Reset() {
 	}
 }
 
+func ExampleContextHandler() {
+	nc, err := nats.Connect("127.0.0.1:4222")
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer nc.Close()
+
+	handler := func(ctx context.Context, req micro.Request) {
+		select {
+		case <-ctx.Done():
+			req.Error("400", "context canceled", nil)
+		default:
+			req.Respond([]byte("ok"))
+		}
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	config := micro.Config{
+		Name:    "EchoService",
+		Version: "0.1.0",
+		Endpoint: &micro.EndpointConfig{
+			Subject: "echo",
+			Handler: micro.ContextHandler(ctx, handler),
+		},
+	}
+
+	srv, _ := micro.AddService(nc, config)
+	defer srv.Stop()
+}
+
 func ExampleControlSubject() {
 
 	// subject used to get PING from all services

micro/request.go

@@ -14,6 +14,7 @@
 package micro
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -82,6 +83,14 @@ func (fn HandlerFunc) Handle(req Request) {
 	fn(req)
 }
 
+// ContextHandler is a helper function used to utilize [context.Context]
+// in request handlers.
+func ContextHandler(ctx context.Context, handler func(context.Context, Request)) Handler {
+	return HandlerFunc(func(req Request) {
+		handler(ctx, req)
+	})
+}
+
 // Respond sends the response for the request.
 // Additional headers can be passed using [WithHeaders] option.
 func (r *request) Respond(response []byte, opts ...RespondOpt) error {

micro/test/service_test.go

@@ -15,6 +15,7 @@ package micro_test
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -984,6 +985,69 @@ func TestMonitoringHandlers(t *testing.T) {
 	}
 }
 
+func TestContextHandler(t *testing.T) {
+	s := RunServerOnPort(-1)
+	defer s.Shutdown()
+
+	nc, err := nats.Connect(s.ClientURL())
+	if err != nil {
+		t.Fatalf("Expected to connect to server, got %v", err)
+	}
+	defer nc.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	type key string
+	ctx = context.WithValue(ctx, key("key"), []byte("val"))
+
+	handler := func(ctx context.Context, req micro.Request) {
+		select {
+		case <-ctx.Done():
+			req.Error("400", "context canceled", nil)
+		default:
+			v := ctx.Value(key("key"))
+			req.Respond(v.([]byte))
+		}
+	}
+	config := micro.Config{
+		Name:    "test_service",
+		Version: "0.1.0",
+		APIURL:  "http://someapi.com/v1",
+		Endpoint: &micro.EndpointConfig{
+			Subject: "test.func",
+			Handler: micro.ContextHandler(ctx, handler),
+			Schema: &micro.Schema{
+				Request:  "request_schema",
+				Response: "response_schema",
+			},
+		},
+	}
+
+	srv, err := micro.AddService(nc, config)
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	defer srv.Stop()
+
+	resp, err := nc.Request("test.func", nil, 1*time.Second)
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err)
+	}
+	if string(resp.Data) != "val" {
+		t.Fatalf("Invalid response; want: %q; got: %q", "val", string(resp.Data))
+	}
+	cancel()
+	resp, err = nc.Request("test.func", nil, 1*time.Second)
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err)
+	}
+	if resp.Header.Get(micro.ErrorCodeHeader) != "400" {
+		t.Fatalf("Expected error response after canceling context; got: %q", string(resp.Data))
+	}
+
+}
+
 func TestServiceNilSchema(t *testing.T) {
 	s := RunServerOnPort(-1)
 	defer s.Shutdown()