internal: for Flex, use req.Context instead of a context map (#107)

api_pre17.go is a copy of the previous api.go, kept for backwards
compatibility.

Some other clean-up, too, like using a sync.Once to initialize the background context.

Author: broady

aetest/instance_vm.go

@@ -71,7 +71,7 @@ func (i *instance) NewRequest(method, urlStr string, body io.Reader) (*http.Requ
 	}
 
 	// Associate this request.
-	release := internal.RegisterTestRequest(req, i.apiURL, func(ctx context.Context) context.Context {
+	req, release := internal.RegisterTestRequest(req, i.apiURL, func(ctx context.Context) context.Context {
 		ctx = internal.WithAppIDOverride(ctx, "dev~"+i.appID)
 		return ctx
 	})

appengine.go

@@ -63,7 +63,7 @@ func IsDevAppServer() bool {
 // NewContext returns a context for an in-flight HTTP request.
 // This function is cheap.
 func NewContext(req *http.Request) context.Context {
-	return WithContext(context.Background(), req)
+	return internal.ReqContext(req)
 }
 
 // WithContext returns a copy of the parent context

internal/api.go

@@ -3,6 +3,7 @@
 // license that can be found in the LICENSE file.
 
 // +build !appengine
+// +build go1.7
 
 package internal
 
@@ -62,8 +63,10 @@ var (
 		},
 	}
 
-	defaultTicketOnce sync.Once
-	defaultTicket     string
+	defaultTicketOnce     sync.Once
+	defaultTicket         string
+	backgroundContextOnce sync.Once
+	backgroundContext     netcontext.Context
 )
 
 func apiURL() *url.URL {
@@ -87,16 +90,10 @@ func handleHTTP(w http.ResponseWriter, r *http.Request) {
 		outHeader: w.Header(),
 		apiURL:    apiURL(),
 	}
-	stopFlushing := make(chan int)
+	r = r.WithContext(withContext(r.Context(), c))
+	c.req = r
 
-	ctxs.Lock()
-	ctxs.m[r] = c
-	ctxs.Unlock()
-	defer func() {
-		ctxs.Lock()
-		delete(ctxs.m, r)
-		ctxs.Unlock()
-	}()
+	stopFlushing := make(chan int)
 
 	// Patch up RemoteAddr so it looks reasonable.
 	if addr := r.Header.Get(userIPHeader); addr != "" {
@@ -195,18 +192,6 @@ func renderPanic(x interface{}) string {
 	return string(buf)
 }
 
-var ctxs = struct {
-	sync.Mutex
-	m  map[*http.Request]*context
-	bg *context // background context, lazily initialized
-	// dec is used by tests to decorate the netcontext.Context returned
-	// for a given request. This allows tests to add overrides (such as
-	// WithAppIDOverride) to the context. The map is nil outside tests.
-	dec map[*http.Request]func(netcontext.Context) netcontext.Context
-}{
-	m: make(map[*http.Request]*context),
-}
-
 // context represents the context of an in-flight HTTP request.
 // It implements the appengine.Context and http.ResponseWriter interfaces.
 type context struct {
@@ -227,6 +212,32 @@ type context struct {
 
 var contextKey = "holds a *context"
 
+// jointContext joins two contexts in a superficial way.
+// It takes values and timeouts from a base context, and only values from another context.
+type jointContext struct {
+	base       netcontext.Context
+	valuesOnly netcontext.Context
+}
+
+func (c jointContext) Deadline() (time.Time, bool) {
+	return c.base.Deadline()
+}
+
+func (c jointContext) Done() <-chan struct{} {
+	return c.base.Done()
+}
+
+func (c jointContext) Err() error {
+	return c.base.Err()
+}
+
+func (c jointContext) Value(key interface{}) interface{} {
+	if val := c.base.Value(key); val != nil {
+		return val
+	}
+	return c.valuesOnly.Value(key)
+}
+
 // fromContext returns the App Engine context or nil if ctx is not
 // derived from an App Engine context.
 func fromContext(ctx netcontext.Context) *context {
@@ -253,23 +264,15 @@ func IncomingHeaders(ctx netcontext.Context) http.Header {
 	return nil
 }
 
-func WithContext(parent netcontext.Context, req *http.Request) netcontext.Context {
-	ctxs.Lock()
-	c := ctxs.m[req]
-	d := ctxs.dec[req]
-	ctxs.Unlock()
-
-	if d != nil {
-		parent = d(parent)
-	}
+func ReqContext(req *http.Request) netcontext.Context {
+	return req.Context()
+}
 
-	if c == nil {
-		// Someone passed in an http.Request that is not in-flight.
-		// We panic here rather than panicking at a later point
-		// so that stack traces will be more sensible.
-		log.Panic("appengine: NewContext passed an unknown http.Request")
+func WithContext(parent netcontext.Context, req *http.Request) netcontext.Context {
+	return jointContext{
+		base:       parent,
+		valuesOnly: req.Context(),
 	}
-	return withContext(parent, c)
 }
 
 // DefaultTicket returns a ticket used for background context or dev_appserver.
@@ -291,60 +294,40 @@ func DefaultTicket() string {
 }
 
 func BackgroundContext() netcontext.Context {
-	ctxs.Lock()
-	defer ctxs.Unlock()
-
-	if ctxs.bg != nil {
-		return toContext(ctxs.bg)
-	}
-
-	// Compute background security ticket.
-	ticket := DefaultTicket()
-
-	ctxs.bg = &context{
-		req: &http.Request{
-			Header: http.Header{
-				ticketHeader: []string{ticket},
+	backgroundContextOnce.Do(func() {
+		// Compute background security ticket.
+		ticket := DefaultTicket()
+
+		c := &context{
+			req: &http.Request{
+				Header: http.Header{
+					ticketHeader: []string{ticket},
+				},
 			},
-		},
-		apiURL: apiURL(),
-	}
+			apiURL: apiURL(),
+		}
+		backgroundContext = toContext(c)
 
-	// TODO(dsymonds): Wire up the shutdown handler to do a final flush.
-	go ctxs.bg.logFlusher(make(chan int))
+		// TODO(dsymonds): Wire up the shutdown handler to do a final flush.
+		go c.logFlusher(make(chan int))
+	})
 
-	return toContext(ctxs.bg)
+	return backgroundContext
 }
 
 // RegisterTestRequest registers the HTTP request req for testing, such that
 // any API calls are sent to the provided URL. It returns a closure to delete
 // the registration.
 // It should only be used by aetest package.
-func RegisterTestRequest(req *http.Request, apiURL *url.URL, decorate func(netcontext.Context) netcontext.Context) func() {
+func RegisterTestRequest(req *http.Request, apiURL *url.URL, decorate func(netcontext.Context) netcontext.Context) (*http.Request, func()) {
 	c := &context{
 		req:    req,
 		apiURL: apiURL,
 	}
-	ctxs.Lock()
-	defer ctxs.Unlock()
-	if _, ok := ctxs.m[req]; ok {
-		log.Panic("req already associated with context")
-	}
-	if _, ok := ctxs.dec[req]; ok {
-		log.Panic("req already associated with context")
-	}
-	if ctxs.dec == nil {
-		ctxs.dec = make(map[*http.Request]func(netcontext.Context) netcontext.Context)
-	}
-	ctxs.m[req] = c
-	ctxs.dec[req] = decorate
-
-	return func() {
-		ctxs.Lock()
-		delete(ctxs.m, req)
-		delete(ctxs.dec, req)
-		ctxs.Unlock()
-	}
+	ctx := withContext(decorate(req.Context()), c)
+	req = req.WithContext(ctx)
+	c.req = req
+	return req, func() {}
 }
 
 var errTimeout = &CallError{

internal/api_classic.go

@@ -59,6 +59,10 @@ func IncomingHeaders(ctx netcontext.Context) http.Header {
 	return nil
 }
 
+func ReqContext(req *http.Request) netcontext.Context {
+	return WithContext(netcontext.Background(), req)
+}
+
 func WithContext(parent netcontext.Context, req *http.Request) netcontext.Context {
 	c := appengine.NewContext(req)
 	return withContext(parent, c)