github.com/google/osv-scalibr@v0.4.1/clients/datasource/cache.go (about)

     1  // Copyright 2025 Google LLC
     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 datasource
    16  
    17  import (
    18  	"bytes"
    19  	"encoding/gob"
    20  	"maps"
    21  	"sync"
    22  	"time"
    23  )
    24  
    25  const cacheExpiry = 6 * time.Hour
    26  
    27  func gobMarshal(v any) ([]byte, error) {
    28  	var b bytes.Buffer
    29  	enc := gob.NewEncoder(&b)
    30  
    31  	err := enc.Encode(v)
    32  	if err != nil {
    33  		return nil, err
    34  	}
    35  
    36  	return b.Bytes(), nil
    37  }
    38  
    39  func gobUnmarshal(b []byte, v any) error {
    40  	dec := gob.NewDecoder(bytes.NewReader(b))
    41  	return dec.Decode(v)
    42  }
    43  
    44  type requestCacheCall[V any] struct {
    45  	wg  sync.WaitGroup
    46  	val V
    47  	err error
    48  }
    49  
    50  // RequestCache is a map to cache the results of expensive functions that are called concurrently.
    51  type RequestCache[K comparable, V any] struct {
    52  	cache map[K]V
    53  	calls map[K]*requestCacheCall[V]
    54  	mu    sync.Mutex
    55  }
    56  
    57  // NewRequestCache creates a new RequestCache.
    58  func NewRequestCache[K comparable, V any]() *RequestCache[K, V] {
    59  	return &RequestCache[K, V]{
    60  		cache: make(map[K]V),
    61  		calls: make(map[K]*requestCacheCall[V]),
    62  	}
    63  }
    64  
    65  // Get gets the value from the cache map if it's cached, otherwise it will call fn to get the value and cache it.
    66  // fn will only ever be called once for a key, even if there are multiple simultaneous calls to Get before the first call is finished.
    67  func (rq *RequestCache[K, V]) Get(key K, fn func() (V, error)) (V, error) {
    68  	// Try get it from regular cache.
    69  	rq.mu.Lock()
    70  	if v, ok := rq.cache[key]; ok {
    71  		rq.mu.Unlock()
    72  		return v, nil
    73  	}
    74  
    75  	// See if there is already a pending request for this key.
    76  	if c, ok := rq.calls[key]; ok {
    77  		rq.mu.Unlock()
    78  		c.wg.Wait()
    79  
    80  		return c.val, c.err
    81  	}
    82  
    83  	// Cache miss - create the call.
    84  	c := new(requestCacheCall[V])
    85  	c.wg.Add(1)
    86  	rq.calls[key] = c
    87  	rq.mu.Unlock()
    88  
    89  	c.val, c.err = fn()
    90  	rq.mu.Lock()
    91  	defer rq.mu.Unlock()
    92  
    93  	// Allow other waiting goroutines to return
    94  	c.wg.Done()
    95  
    96  	// Store value in regular cache.
    97  	if c.err == nil {
    98  		rq.cache[key] = c.val
    99  	}
   100  
   101  	// Remove the completed call now that it's cached.
   102  	if rq.calls[key] == c {
   103  		delete(rq.calls, key)
   104  	}
   105  
   106  	return c.val, c.err
   107  }
   108  
   109  // GetMap gets a shallow clone of the stored cache map.
   110  func (rq *RequestCache[K, V]) GetMap() map[K]V {
   111  	rq.mu.Lock()
   112  	defer rq.mu.Unlock()
   113  
   114  	return maps.Clone(rq.cache)
   115  }
   116  
   117  // SetMap loads (a shallow clone of) the provided map into the cache map.
   118  func (rq *RequestCache[K, V]) SetMap(m map[K]V) {
   119  	rq.mu.Lock()
   120  	defer rq.mu.Unlock()
   121  	rq.cache = maps.Clone(m)
   122  }