Rework flight cache into a strongly typed version (#1989)

This reworks the flight cache we already have into a strongly typed
version based on a map and `sync.OnceValues`. That allows us to get rid
of type assertions and unnecessary error handling in its implementation.
Also add some tests for the behavior of flight cache.

Author: markusthoemmes

pkg/apk/apk/cache.go

@@ -32,66 +32,62 @@ import (
 	"chainguard.dev/apko/pkg/paths"
 )
 
-type flightCache[T any] struct {
-	flight *singleflight.Group
-	cache  *sync.Map
+type flightCache[K comparable, V any] struct {
+	mux   sync.RWMutex
+	cache map[K]func() (V, error)
 }
 
-// TODO: Consider [K, V] if we need a non-string key type.
-func newFlightCache[T any]() *flightCache[T] {
-	return &flightCache[T]{
-		flight: &singleflight.Group{},
-		cache:  &sync.Map{},
+func newFlightCache[K comparable, V any]() *flightCache[K, V] {
+	return &flightCache[K, V]{
+		cache: make(map[K]func() (V, error)),
 	}
 }
 
 // Do returns coalesces multiple calls, like singleflight, but also caches
 // the result if the call is successful. Failures are not cached to avoid
 // permanently failing for transient errors.
-func (f *flightCache[T]) Do(key string, fn func() (T, error)) (T, error) {
-	v, ok := f.cache.Load(key)
-	if ok {
-		if t, ok := v.(T); ok {
-			return t, nil
-		} else {
-			// This can't happen but just in case things change.
-			return t, fmt.Errorf("unexpected type %T", v)
-		}
+func (f *flightCache[K, V]) Do(key K, fn func() (V, error)) (V, error) {
+	f.mux.RLock()
+	if v, ok := f.cache[key]; ok {
+		f.mux.RUnlock()
+		return v()
 	}
+	f.mux.RUnlock()
 
-	v, err, _ := f.flight.Do(key, func() (any, error) {
-		if v, ok := f.cache.Load(key); ok {
-			return v, nil
-		}
+	f.mux.Lock()
 
-		// Don't cache errors, but maybe we should.
-		v, err := fn()
-		if err != nil {
-			return nil, err
-		}
+	// Doubly-checked-locking in case of race conditions.
+	if v, ok := f.cache[key]; ok {
+		f.mux.Unlock()
+		return v()
+	}
 
-		f.cache.Store(key, v)
+	v := sync.OnceValues(fn)
+	f.cache[key] = v
 
-		return v, nil
-	})
+	// Unlock before calling the function to avoid holding the lock for a potentially long time.
+	f.mux.Unlock()
 
-	t, ok := v.(T)
+	val, err := v()
 	if err != nil {
-		return t, err
-	}
-	if !ok {
-		// This can't happen but just in case things change.
-		return t, fmt.Errorf("unexpected type %T", v)
+		f.Forget(key)
 	}
-	return t, nil
+	return val, err
+}
+
+// Forget removes the given key from the cache.
+func (f *flightCache[K, V]) Forget(key K) {
+	f.mux.Lock()
+	defer f.mux.Unlock()
+	delete(f.cache, key)
 }
 
 type Cache struct {
 	etagCache  *sync.Map
 	headFlight *singleflight.Group
 	getFlight  *singleflight.Group
 
-	discoverKeys *flightCache[[]Key]
+	discoverKeys *flightCache[string, []Key]
 }
 
 // NewCache returns a new Cache, which allows us to persist the results of HEAD requests
@@ -109,7 +105,7 @@ func NewCache(etag bool) *Cache {
 	c := &Cache{
 		headFlight:   &singleflight.Group{},
 		getFlight:    &singleflight.Group{},
-		discoverKeys: newFlightCache[[]Key](),
+		discoverKeys: newFlightCache[string, []Key](),
 	}
 
 	if etag {

pkg/apk/apk/cache_test.go

@@ -0,0 +1,102 @@
+// Copyright 2025 Chainguard, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package apk
+
+import (
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/sync/errgroup"
+)
+
+func TestFlightCache(t *testing.T) {
+	s := newFlightCache[string, int]()
+	var called int
+	r1, err := s.Do("test", func() (int, error) {
+		called++
+		return 42, nil
+	})
+	require.NoError(t, err)
+	require.Equal(t, 42, r1)
+
+	r2, err := s.Do("test", func() (int, error) {
+		called++
+		return 1337, nil
+	})
+	require.NoError(t, err)
+	require.Equal(t, r1, r2)
+	require.Equal(t, 1, called, "Function should only be called once")
+
+	s.Forget("test")
+
+	r3, err := s.Do("test", func() (int, error) {
+		called++
+		return 1337, nil
+	})
+	require.NoError(t, err)
+	require.Equal(t, 1337, r3)
+	require.Equal(t, 2, called, "Function should be called twice, once before and once after Forget")
+
+	differentKey, err := s.Do("test2", func() (int, error) {
+		return 7, nil
+	})
+	require.NoError(t, err)
+	require.Equal(t, 7, differentKey)
+}
+
+func TestFlightCacheCachesNoErrors(t *testing.T) {
+	s := newFlightCache[string, int]()
+	var called int
+	_, err := s.Do("test", func() (int, error) {
+		called++
+		return 42, assert.AnError
+	})
+	require.ErrorIs(t, assert.AnError, err)
+
+	r2, err := s.Do("test", func() (int, error) {
+		called++
+		return 1337, nil
+	})
+	require.NoError(t, err)
+	require.Equal(t, 1337, r2)
+	require.Equal(t, 2, called, "Function should be called twice, once for the error and once for the success")
+}
+
+func TestFlightCacheCoalescesCalls(t *testing.T) {
+	s := newFlightCache[string, int]()
+
+	var called atomic.Int32
+	var mux sync.Mutex
+	mux.Lock() // Lock to ensure the call below hangs until we unlock.
+
+	var eg errgroup.Group
+	for range 10 {
+		eg.Go(func() error {
+			_, err := s.Do("test", func() (int, error) {
+				mux.Lock() // Hangs until the unlock below.
+				called.Add(1)
+				return 42, nil
+			})
+			return err
+		})
+	}
+	mux.Unlock() // Allow the calls to proceed.
+	require.NoError(t, eg.Wait())
+
+	require.EqualValues(t, 1, called.Load(), "Function should only be called once")
+}