go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/gae/impl/prod/context.go (about)

     1  // Copyright 2015 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package prod
    16  
    17  import (
    18  	"context"
    19  	"net/http"
    20  
    21  	"google.golang.org/appengine"
    22  )
    23  
    24  var (
    25  	prodStateKey  = "contains the current *prodState"
    26  	probeCacheKey = "contains the current *infoProbeCache"
    27  )
    28  
    29  // getAEContext retrieves the raw "google.golang.org/appengine" compatible
    30  // Context.
    31  //
    32  // This is an independent Context chain from `c`. In an attempt to maintain user
    33  // expectations, the deadline of `c` is transferred to the returned Context,
    34  // RPCs. Cancelation is not transferred.
    35  func getAEContext(c context.Context) context.Context {
    36  	ps := getProdState(c)
    37  	return ps.context(c)
    38  }
    39  
    40  func setupAECtx(c, aeCtx context.Context, r *http.Request) context.Context {
    41  	c = withProdState(c, prodState{
    42  		ctx:      aeCtx,
    43  		noTxnCtx: aeCtx,
    44  	})
    45  	return useModule(useMail(useUser(useURLFetch(useRDS(useMC(useTQ(useGI(useLogging(c, r)))))))))
    46  }
    47  
    48  // Use adds production implementations for all the gae services to the
    49  // context. The implementations are all backed by the real appengine SDK
    50  // functionality.
    51  //
    52  // The services added are:
    53  //   - github.com/luci-go/common/logging
    54  //   - go.chromium.org/luci/gae/service/datastore
    55  //   - go.chromium.org/luci/gae/service/info
    56  //   - go.chromium.org/luci/gae/service/mail
    57  //   - go.chromium.org/luci/gae/service/memcache
    58  //   - go.chromium.org/luci/gae/service/module
    59  //   - go.chromium.org/luci/gae/service/taskqueue
    60  //   - go.chromium.org/luci/gae/service/urlfetch
    61  //   - go.chromium.org/luci/gae/service/user
    62  //
    63  // These can be retrieved with the <service>.Get functions.
    64  //
    65  // It is important to note that this DOES NOT install the AppEngine SDK into the
    66  // supplied Context. In general, using the raw AppEngine SDK to access a service
    67  // that is covered by luci/gae is dangerous, leading to a number of potential
    68  // pitfalls including inconsistent transaction management and data corruption.
    69  //
    70  // Users who wish to access the raw AppEngine SDK must derive their own
    71  // AppEngine Context at their own risk.
    72  func Use(c context.Context, r *http.Request) context.Context {
    73  	return setupAECtx(c, appengine.NewContext(r), r)
    74  }
    75  
    76  // prodState is the current production state.
    77  type prodState struct {
    78  	// ctx is the current derived GAE context.
    79  	ctx context.Context
    80  
    81  	// noTxnCtx is a Context maintained alongside ctx. When a transaction is
    82  	// entered, ctx will be updated, but noTxnCtx will not, allowing extra-
    83  	// transactional Context access.
    84  	noTxnCtx context.Context
    85  
    86  	// inTxn if true if this is in a transaction, false otherwise.
    87  	inTxn bool
    88  }
    89  
    90  func getProdState(c context.Context) prodState {
    91  	if v := c.Value(&prodStateKey).(*prodState); v != nil {
    92  		return *v
    93  	}
    94  	return prodState{}
    95  }
    96  
    97  func withProdState(c context.Context, ps prodState) context.Context {
    98  	return context.WithValue(c, &prodStateKey, &ps)
    99  }
   100  
   101  // context returns the current AppEngine-bound Context. Prior to returning,
   102  // the deadline from "c" (if any) is applied.
   103  //
   104  // Note that this does not (currently) apply any other Done state or propagate
   105  // cancellation from "c".
   106  //
   107  // Tracking at:
   108  // https://go.chromium.org/luci/gae/issues/59
   109  func (ps *prodState) context(c context.Context) context.Context {
   110  	aeCtx := ps.ctx
   111  	if aeCtx == nil {
   112  		return nil
   113  	}
   114  
   115  	if deadline, ok := c.Deadline(); ok {
   116  		aeCtx, _ = context.WithDeadline(aeCtx, deadline)
   117  	}
   118  	return aeCtx
   119  }