cache.go

     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
View as plain text