go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/gae/filter/dscache/support.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  	"crypto/sha1"
    20  	"encoding/base64"
    21  	"fmt"
    22  
    23  	"go.chromium.org/luci/common/data/rand/mathrand"
    24  	"go.chromium.org/luci/common/logging"
    25  
    26  	ds "go.chromium.org/luci/gae/service/datastore"
    27  )
    28  
    29  type supportContext struct {
    30  	ds.KeyContext
    31  
    32  	c            context.Context
    33  	impl         Cache
    34  	shardsForKey []ShardFunction
    35  }
    36  
    37  func (s *supportContext) numShards(k *ds.Key) int {
    38  	ret := DefaultShards
    39  	for _, fn := range s.shardsForKey {
    40  		if amt, ok := fn(k); ok {
    41  			ret = amt
    42  		}
    43  	}
    44  	if ret < 1 {
    45  		return 0 // disable caching entirely
    46  	}
    47  	if ret > MaxShards {
    48  		ret = MaxShards
    49  	}
    50  	return ret
    51  }
    52  
    53  func (s *supportContext) mkRandKeys(keys []*ds.Key, metas ds.MultiMetaGetter, rnd mathrand.Rand) []string {
    54  	ret := []string(nil)
    55  	for i, key := range keys {
    56  		mg := metas.GetSingle(i)
    57  		if !ds.GetMetaDefault(mg, CacheEnableMeta, true).(bool) {
    58  			continue
    59  		}
    60  		shards := s.numShards(key)
    61  		if shards == 0 {
    62  			continue
    63  		}
    64  		shard := 0
    65  		if shards > 1 {
    66  			shard = rnd.Intn(shards)
    67  		}
    68  		if ret == nil {
    69  			ret = make([]string, len(keys))
    70  		}
    71  		ret[i] = makeMemcacheKey(shard, key)
    72  	}
    73  	return ret
    74  }
    75  
    76  func (s *supportContext) mkAllKeys(keys []*ds.Key) []string {
    77  	size := 0
    78  	nums := make([]int, len(keys))
    79  	for i, key := range keys {
    80  		if !key.IsIncomplete() {
    81  			shards := s.numShards(key)
    82  			nums[i] = shards
    83  			size += shards
    84  		}
    85  	}
    86  	if size == 0 {
    87  		return nil
    88  	}
    89  	ret := make([]string, 0, size)
    90  	for i, key := range keys {
    91  		if !key.IsIncomplete() {
    92  			keySuffix := hashKey(key)
    93  			for shard := 0; shard < nums[i]; shard++ {
    94  				ret = append(ret, fmt.Sprintf(KeyFormat, shard, keySuffix))
    95  			}
    96  		}
    97  	}
    98  	return ret
    99  }
   100  
   101  func (s *supportContext) mutation(keys []*ds.Key, f func() error) error {
   102  	itemKeys := s.mkAllKeys(keys)
   103  	if len(itemKeys) == 0 {
   104  		return f()
   105  	}
   106  	if err := s.impl.PutLocks(s.c, itemKeys, MutationLockTimeout); err != nil {
   107  		// this is a hard failure. No mutation can occur if we're unable to set
   108  		// locks out. See "DANGER ZONE" in the docs.
   109  		logging.WithError(err).Errorf(s.c, "dscache: HARD FAILURE: supportContext.mutation(): PutLocks")
   110  		return err
   111  	}
   112  	err := f()
   113  	// Note: the mutation can *eventually* succeed even if `err` is non-nil
   114  	// here. So on errors we pessimistically keep the locks until they expire.
   115  	if err == nil {
   116  		if err := s.impl.DropLocks(s.c, itemKeys); err != nil {
   117  			logging.WithError(err).Debugf(s.c, "dscache: DropLocks")
   118  		}
   119  	}
   120  	return err
   121  }
   122  
   123  // generateNonce creates a pseudo-random sequence of bytes for use as a nonce
   124  // using the non-cryptographic PRNG in "math/rand".
   125  //
   126  // The random values here are controlled entirely by the application, will never
   127  // be shown to, or provided by, the user, so this should be fine.
   128  func generateNonce(rnd mathrand.Rand) []byte {
   129  	nonce := make([]byte, NonceBytes)
   130  	_, _ = rnd.Read(nonce) // This Read will always return len(nonce), nil.
   131  	return nonce
   132  }
   133  
   134  // makeMemcacheKey generates a memcache key for the given datastore Key.
   135  func makeMemcacheKey(shard int, k *ds.Key) string {
   136  	return fmt.Sprintf(KeyFormat, shard, hashKey(k))
   137  }
   138  
   139  // hashKey generates just the hashed portion of the memcache key.
   140  func hashKey(k *ds.Key) string {
   141  	dgst := sha1.Sum(ds.Serialize.ToBytes(k))
   142  	return base64.RawStdEncoding.EncodeToString(dgst[:])
   143  }