Add request_context flag to utilize (*http.Request).Context() in handlers (#265)

* Add usage of http.Request's Context
* add request_context flag

Contribution by @kokaz, @tmc and @yugui

Author: tmc

.travis.yml

@@ -23,10 +23,13 @@ before_script:
 - sh -c 'cd examples/browser && npm install'
 script:
 - make realclean && make examples SWAGGER_CODEGEN="java -jar $HOME/local/swagger-codegen-cli.jar"
-- if ! go version | grep devel; then test -z "$(git status --porcelain)" || (git status; git diff; exit 1); fi
+- if (go version | grep -qv devel) && [ -z "${GATEWAY_PLUGIN_FLAGS}" ]; then test -z "$(git status --porcelain)" || (git status; git diff; exit 1); fi
 - env GLOG_logtostderr=1 go test -race -v github.com/grpc-ecosystem/grpc-gateway/...
 - make lint
 - sh -c 'cd examples/browser && gulp'
 env:
   global:
   - "PATH=$PATH:$HOME/local/bin"
+  matrix:
+  - GATEWAY_PLUGIN_FLAGS=
+  - GATEWAY_PLUGIN_FLAGS=request_context=true

Makefile

@@ -35,6 +35,7 @@ GATEWAY_PLUGIN_SRC= utilities/doc.go \
 		    protoc-gen-grpc-gateway/httprule/parse.go \
 		    protoc-gen-grpc-gateway/httprule/types.go \
 		    protoc-gen-grpc-gateway/main.go
+GATEWAY_PLUGIN_FLAGS?=
 
 GOOGLEAPIS_DIR=third_party/googleapis
 OPTIONS_PROTO=$(GOOGLEAPIS_DIR)/google/api/annotations.proto $(GOOGLEAPIS_DIR)/google/api/http.proto
@@ -45,6 +46,10 @@ RUNTIME_PROTO=runtime/internal/stream_chunk.proto
 RUNTIME_GO=$(RUNTIME_PROTO:.proto=.pb.go)
 
 PKGMAP=Mgoogle/protobuf/descriptor.proto=$(GO_PLUGIN_PKG)/descriptor,Mgoogle/api/annotations.proto=$(PKG)/$(GOOGLEAPIS_DIR)/google/api,Mexamples/sub/message.proto=$(PKG)/examples/sub
+ADDITIONAL_FLAGS=
+ifneq "$(GATEWAY_PLUGIN_FLAGS)" ""
+	ADDITIONAL_FLAGS=,$(GATEWAY_PLUGIN_FLAGS)
+endif
 SWAGGER_EXAMPLES=examples/examplepb/echo_service.proto \
 	 examples/examplepb/a_bit_of_everything.proto
 EXAMPLES=examples/examplepb/echo_service.proto \
@@ -102,7 +107,7 @@ $(EXAMPLE_DEPSRCS): $(GO_PLUGIN) $(EXAMPLE_DEPS)
 	protoc -I $(PROTOC_INC_PATH) -I. --plugin=$(GO_PLUGIN) --go_out=$(PKGMAP),plugins=grpc:$(OUTPUT_DIR) $(@:.pb.go=.proto)
 	cp $(OUTPUT_DIR)/$(PKG)/$@ $@ || cp $(OUTPUT_DIR)/$@ $@
 $(EXAMPLE_GWSRCS): $(GATEWAY_PLUGIN) $(EXAMPLES)
-	protoc -I $(PROTOC_INC_PATH) -I. -I$(GOOGLEAPIS_DIR) --plugin=$(GATEWAY_PLUGIN) --grpc-gateway_out=logtostderr=true,$(PKGMAP):. $(EXAMPLES)
+	protoc -I $(PROTOC_INC_PATH) -I. -I$(GOOGLEAPIS_DIR) --plugin=$(GATEWAY_PLUGIN) --grpc-gateway_out=logtostderr=true,$(PKGMAP)$(ADDITIONAL_FLAGS):. $(EXAMPLES)
 $(EXAMPLE_SWAGGERSRCS): $(SWAGGER_PLUGIN) $(SWAGGER_EXAMPLES)
 	protoc -I $(PROTOC_INC_PATH) -I. -I$(GOOGLEAPIS_DIR) --plugin=$(SWAGGER_PLUGIN) --swagger_out=logtostderr=true,$(PKGMAP):. $(SWAGGER_EXAMPLES)
 

README.md

@@ -186,6 +186,9 @@ Make sure that your `$GOPATH/bin` is in your `$PATH`.
 `protoc-gen-grpc-gateway` supports custom mapping from Protobuf `import` to Golang import path.
 They are compatible to [the parameters with same names in `protoc-gen-go`](https://github.com/golang/protobuf#parameters).
 
+In addition we also support the `request_context` parameter in order to use the `http.Request`'s Context (only for Go 1.7 and above).
+This parameter can be useful to pass request scoped context between the gateway and the gRPC service.
+
 `protoc-gen-grpc-gateway` also supports some more command line flags to control logging. You can give these flags together with parameters above. Run `protoc-gen-grpc-gateway --help` for more details about the flags.
 
 ## More Examples

protoc-gen-grpc-gateway/gengateway/generator.go

@@ -20,12 +20,13 @@ var (
 )
 
 type generator struct {
-	reg         *descriptor.Registry
-	baseImports []descriptor.GoPackage
+	reg               *descriptor.Registry
+	baseImports       []descriptor.GoPackage
+	useRequestContext bool
 }
 
 // New returns a new generator which generates grpc gateway files.
-func New(reg *descriptor.Registry) gen.Generator {
+func New(reg *descriptor.Registry, useRequestContext bool) gen.Generator {
 	var imports []descriptor.GoPackage
 	for _, pkgpath := range []string{
 		"io",
@@ -54,7 +55,7 @@ func New(reg *descriptor.Registry) gen.Generator {
 		}
 		imports = append(imports, pkg)
 	}
-	return &generator{reg: reg, baseImports: imports}
+	return &generator{reg: reg, baseImports: imports, useRequestContext: useRequestContext}
 }
 
 func (g *generator) Generate(targets []*descriptor.File) ([]*plugin.CodeGeneratorResponse_File, error) {
@@ -107,5 +108,5 @@ func (g *generator) generate(file *descriptor.File) (string, error) {
 			imports = append(imports, pkg)
 		}
 	}
-	return applyTemplate(param{File: file, Imports: imports})
+	return applyTemplate(param{File: file, Imports: imports, UseRequestContext: g.useRequestContext})
 }

protoc-gen-grpc-gateway/gengateway/template.go

@@ -13,7 +13,8 @@ import (
 
 type param struct {
 	*descriptor.File
-	Imports []descriptor.GoPackage
+	Imports           []descriptor.GoPackage
+	UseRequestContext bool
 }
 
 type binding struct {
@@ -66,6 +67,11 @@ func (f queryParamFilter) String() string {
 	return fmt.Sprintf("&utilities.DoubleArray{Encoding: map[string]int{%s}, Base: %#v, Check: %#v}", e, f.Base, f.Check)
 }
 
+type trailerParams struct {
+	Services          []*descriptor.Service
+	UseRequestContext bool
+}
+
 func applyTemplate(p param) (string, error) {
 	w := bytes.NewBuffer(nil)
 	if err := headerTemplate.Execute(w, p); err != nil {
@@ -90,7 +96,12 @@ func applyTemplate(p param) (string, error) {
 	if len(targetServices) == 0 {
 		return "", errNoTargetService
 	}
-	if err := trailerTemplate.Execute(w, targetServices); err != nil {
+
+	tp := trailerParams{
+		Services:          targetServices,
+		UseRequestContext: p.UseRequestContext,
+	}
+	if err := trailerTemplate.Execute(w, tp); err != nil {
 		return "", err
 	}
 	return w.String(), nil
@@ -295,7 +306,8 @@ var (
 `))
 
 	trailerTemplate = template.Must(template.New("trailer").Parse(`
-{{range $svc := .}}
+{{$UseRequestContext := .UseRequestContext}}
+{{range $svc := .Services}}
 // Register{{$svc.GetName}}HandlerFromEndpoint is same as Register{{$svc.GetName}}Handler but
 // automatically dials to "endpoint" and closes the connection when "ctx" gets done.
 func Register{{$svc.GetName}}HandlerFromEndpoint(ctx context.Context, mux *runtime.ServeMux, endpoint string, opts []grpc.DialOption) (err error) {
@@ -328,7 +340,11 @@ func Register{{$svc.GetName}}Handler(ctx context.Context, mux *runtime.ServeMux,
 	{{range $m := $svc.Methods}}
 	{{range $b := $m.Bindings}}
 	mux.Handle({{$b.HTTPMethod | printf "%q"}}, pattern_{{$svc.GetName}}_{{$m.GetName}}_{{$b.Index}}, func(w http.ResponseWriter, req *http.Request, pathParams map[string]string) {
+	{{- if $UseRequestContext }}
+		ctx, cancel := context.WithCancel(req.Context())
+	{{- else -}}
 		ctx, cancel := context.WithCancel(ctx)
+	{{- end }}
 		defer cancel()
 		if cn, ok := w.(http.CloseNotifier); ok {
 			go func(done <-chan struct{}, closed <-chan bool) {

protoc-gen-grpc-gateway/main.go

@@ -23,7 +23,8 @@ import (
 )
 
 var (
-	importPrefix = flag.String("import_prefix", "", "prefix to be added to go package paths for imported proto files")
+	importPrefix      = flag.String("import_prefix", "", "prefix to be added to go package paths for imported proto files")
+	useRequestContext = flag.Bool("request_context", false, "determine whether to use http.Request's context or not")
 )
 
 func parseReq(r io.Reader) (*plugin.CodeGeneratorRequest, error) {
@@ -73,7 +74,7 @@ func main() {
 		}
 	}
 
-	g := gengateway.New(reg)
+	g := gengateway.New(reg, *useRequestContext)
 
 	reg.SetPrefix(*importPrefix)
 	if err := reg.Load(req); err != nil {