go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/gae/filter/dscache/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 dscache
    16  
    17  import (
    18  	"context"
    19  
    20  	ds "go.chromium.org/luci/gae/service/datastore"
    21  )
    22  
    23  var dsTxnCacheKey = "holds a *dsCache"
    24  var dsShardFunctionsKey = "holds []ShardFunction"
    25  var defaultImpl Cache = memcacheImpl{}
    26  
    27  // ShardFunction is a user-controllable function which calculates the number of
    28  // shards to use for a certain datastore key. The provided key will always be
    29  // valid and complete. It should return ok=true if it recognized the Key, and
    30  // false otherwise.
    31  //
    32  // The # of shards returned may be between 1 and 256. Values above this range
    33  // will be clamped into that range. A return value of 0 means that NO cache
    34  // operations should be done for this key, regardless of the dscache.enable
    35  // setting.
    36  type ShardFunction func(*ds.Key) (shards int, ok bool)
    37  
    38  // FilterRDS installs a caching RawDatastore filter in the context.
    39  //
    40  // It uses the given `impl` to actually do caching. If nil, uses GAE Memcache.
    41  func FilterRDS(ctx context.Context, impl Cache) context.Context {
    42  	if impl == nil {
    43  		impl = defaultImpl
    44  	}
    45  	return ds.AddRawFilters(ctx, func(ctx context.Context, rds ds.RawInterface) ds.RawInterface {
    46  		shardFns, _ := ctx.Value(&dsShardFunctionsKey).([]ShardFunction)
    47  
    48  		sc := &supportContext{
    49  			ds.GetKeyContext(ctx),
    50  			ctx,
    51  			impl,
    52  			shardFns,
    53  		}
    54  
    55  		v := ctx.Value(&dsTxnCacheKey)
    56  		if v == nil {
    57  			return &dsCache{rds, sc}
    58  		}
    59  		return &dsTxnCache{rds, v.(*dsTxnState), sc}
    60  	})
    61  }
    62  
    63  // AddShardFunctions appends the provided shardFn functions to the internal list
    64  // of shard functions. They are evaluated left to right, bottom to top.
    65  //
    66  // nil functions will cause a panic.
    67  //
    68  // So:
    69  //
    70  //	ctx = AddShardFunctions(ctx, A, B, C)
    71  //	ctx = AddShardFunctions(ctx, D, E, F)
    72  //
    73  // Would evaluate `D, E, F, A, B, C`
    74  func AddShardFunctions(ctx context.Context, shardFns ...ShardFunction) context.Context {
    75  	cur, _ := ctx.Value(&dsShardFunctionsKey).([]ShardFunction)
    76  	new := make([]ShardFunction, 0, len(cur)+len(shardFns))
    77  	for _, fn := range shardFns {
    78  		if fn == nil {
    79  			panic("nil function provided to AddShardFunctions")
    80  		}
    81  	}
    82  	return context.WithValue(ctx, &dsShardFunctionsKey, append(append(new, shardFns...), cur...))
    83  }