cache.mx raw

   1  // Copyright 2025 The Go Authors. All rights reserved.
   2  // Use of this source code is governed by a BSD-style
   3  // license that can be found in the LICENSE file.
   4  
   5  // Package fips140cache provides a weak map that associates the lifetime of
   6  // values with the lifetime of keys.
   7  //
   8  // It can be used to associate a precomputed value (such as an internal/fips140
   9  // PrivateKey value, which in FIPS 140-3 mode may have required an expensive
  10  // pairwise consistency test) with a type that doesn't have private fields (such
  11  // as an ed25519.PrivateKey), or that can't be safely modified because it may be
  12  // concurrently copied (such as an ecdsa.PrivateKey).
  13  package fips140cache
  14  
  15  import (
  16  	"runtime"
  17  	"sync"
  18  	"weak"
  19  )
  20  
  21  type Cache[K, V any] struct {
  22  	m sync.Map
  23  }
  24  
  25  // Get returns the result of new, for an associated key k.
  26  //
  27  // If Get was called with k before and didn't return an error, Get may return
  28  // the same value it returned from the previous call if check returns true on
  29  // it. If check returns false, Get will call new again and return the result.
  30  //
  31  // The cache is evicted some time after k becomes unreachable.
  32  func (c *Cache[K, V]) Get(k *K, new func() (*V, error), check func(*V) bool) (*V, error) {
  33  	p := weak.Make(k)
  34  	if cached, ok := c.m.Load(p); ok {
  35  		v := cached.(*V)
  36  		if check(v) {
  37  			return v, nil
  38  		}
  39  	}
  40  	v, err := new()
  41  	if err != nil {
  42  		return nil, err
  43  	}
  44  	if _, present := c.m.Swap(p, v); !present {
  45  		runtime.AddCleanup(k, c.evict, p)
  46  	}
  47  	return v, nil
  48  }
  49  
  50  func (c *Cache[K, V]) evict(p weak.Pointer[K]) {
  51  	c.m.Delete(p)
  52  }
  53