Make modules/context.Context a context.Context (#16031)

* Make modules/context.Context a context.Context

Signed-off-by: Andrew Thornton <art27@cantab.net>

* Simplify context calls

Signed-off-by: Andrew Thornton <art27@cantab.net>

* Set the base context for requests to the HammerContext

Signed-off-by: Andrew Thornton <art27@cantab.net>

Co-authored-by: Lunny Xiao <xiaolunwen@gmail.com>

modules/context/context.go

@@ -509,7 +509,7 @@ func (ctx *Context) ParamsInt64(p string) int64 {
 
 // SetParams set params into routes
 func (ctx *Context) SetParams(k, v string) {
-	chiCtx := chi.RouteContext(ctx.Req.Context())
+	chiCtx := chi.RouteContext(ctx)
 	chiCtx.URLParams.Add(strings.TrimPrefix(k, ":"), url.PathEscape(v))
 }
 
@@ -528,6 +528,26 @@ func (ctx *Context) Status(status int) {
 	ctx.Resp.WriteHeader(status)
 }
 
+// Deadline is part of the interface for context.Context and we pass this to the request context
+func (ctx *Context) Deadline() (deadline time.Time, ok bool) {
+	return ctx.Req.Context().Deadline()
+}
+
+// Done is part of the interface for context.Context and we pass this to the request context
+func (ctx *Context) Done() <-chan struct{} {
+	return ctx.Req.Context().Done()
+}
+
+// Err is part of the interface for context.Context and we pass this to the request context
+func (ctx *Context) Err() error {
+	return ctx.Req.Context().Err()
+}
+
+// Value is part of the interface for context.Context and we pass this to the request context
+func (ctx *Context) Value(key interface{}) interface{} {
+	return ctx.Req.Context().Value(key)
+}
+
 // Handler represents a custom handler
 type Handler func(*Context)
 

modules/graceful/server_http.go

@@ -5,7 +5,9 @@
 package graceful
 
 import (
+	"context"
 	"crypto/tls"
+	"net"
 	"net/http"
 )
 
@@ -16,6 +18,7 @@ func newHTTPServer(network, address, name string, handler http.Handler) (*Server
 		WriteTimeout:   DefaultWriteTimeOut,
 		MaxHeaderBytes: DefaultMaxHeaderBytes,
 		Handler:        handler,
+		BaseContext:    func(net.Listener) context.Context { return GetManager().HammerContext() },
 	}
 	server.OnShutdown = func() {
 		httpServer.SetKeepAlivesEnabled(false)