github.com/bazelbuild/remote-apis-sdks@v0.0.0-20240425170053-8a36686a6350/go/pkg/cache/singleflightcache.go (about)

     1  // Package cache implements a cache that supports single-flight value computation.
     2  //
     3  // Single-flight value computation means that for concurrent callers of the cache, only one caller
     4  // will compute the value to avoid redundant work which could be system intensive.
     5  package cache
     6  
     7  import (
     8  	"sync"
     9  )
    10  
    11  // SingleFlight is a cache that supports single-flight value computation.
    12  type SingleFlight struct {
    13  	mu    sync.RWMutex // protects `store` field itself
    14  	store sync.Map
    15  }
    16  
    17  type entry struct {
    18  	compute sync.Once
    19  	val     interface{}
    20  	err     error
    21  }
    22  
    23  // LoadOrStore is similar to a sync.Map except that it receives a function that computes the value
    24  // to store instead of the value directly. It ensures that the function is only executed once for
    25  // concurrent callers of the LoadOrStore function.
    26  func (s *SingleFlight) LoadOrStore(key interface{}, valFn func() (val interface{}, err error)) (interface{}, error) {
    27  	s.mu.RLock()
    28  	defer s.mu.RUnlock()
    29  	eUntyped, _ := s.store.LoadOrStore(key, &entry{})
    30  	e := eUntyped.(*entry)
    31  	e.compute.Do(func() {
    32  		e.val, e.err = valFn()
    33  	})
    34  	return e.val, e.err
    35  }
    36  
    37  // Load is similar to a sync.Map.Load.
    38  //
    39  // Callers must check loaded before val and err. Their value is meaninful only
    40  // if loaded is true. The err is not the last returned value because it would
    41  // likely lead the reader to think that err must be checked before loaded.
    42  //
    43  //lint:ignore ST1008 loaded must be last, see above.
    44  func (s *SingleFlight) Load(key interface{}) (val interface{}, err error, loaded bool) {
    45  	s.mu.RLock()
    46  	defer s.mu.RUnlock()
    47  	var eUntyped interface{}
    48  	eUntyped, loaded = s.store.Load(key)
    49  	if loaded {
    50  		e := eUntyped.(*entry)
    51  		val = e.val
    52  		err = e.err
    53  	}
    54  	return
    55  }
    56  
    57  // Store forcefully updates the given cache key with val. Note that unlike LoadOrStore,
    58  // Store accepts a value instead of a valFn since it is intended to be only used in
    59  // cases where updates are lightweight and do not involve computing the cache value.
    60  func (s *SingleFlight) Store(key interface{}, val interface{}) {
    61  	s.mu.RLock()
    62  	defer s.mu.RUnlock()
    63  	e := &entry{val: val}
    64  	e.compute.Do(func() {}) // mark as computed
    65  	s.store.Store(key, e)
    66  }
    67  
    68  // Delete removes a key from the cache.
    69  func (s *SingleFlight) Delete(key interface{}) {
    70  	s.mu.RLock()
    71  	defer s.mu.RUnlock()
    72  	s.store.Delete(key)
    73  }
    74  
    75  // Reset invalidates all cache entries.
    76  func (s *SingleFlight) Reset() {
    77  	s.mu.Lock()
    78  	defer s.mu.Unlock()
    79  	s.store = sync.Map{}
    80  }