internal: cache package

internal/cache/doc.go

@@ -0,0 +1,2 @@
+// Package cache provides caching implementations for Go values.
+package cache

internal/cache/live.go

@@ -0,0 +1,87 @@
+//go:build go1.24
+
+package cache
+
+import (
+	"context"
+	"runtime"
+	"sync"
+	"unique"
+	"weak"
+
+	"github.com/quay/claircore/internal/singleflight"
+)
+
+// Live is a cache that keeps a cached copy as long as the go runtime determines
+// the value is live.
+//
+// See also: [weak.Pointer].
+type Live[K comparable, V any] struct {
+	create func(context.Context, K) (*V, error)
+	m      sync.Map
+	sf     singleflight.Group[K, *V]
+}
+
+// NewLive creates a cache that relies on the runtime's liveness judgment.
+//
+// If the "create" function needs a complex type to be able to construct a
+// value, consider using [unique.Handle] to be able to satisfy the "comparable"
+// constraint.
+func NewLive[K comparable, V any](create func(context.Context, K) (*V, error)) *Live[K, V] {
+	return &Live[K, V]{create: create}
+}
+
+var _ unique.Handle[string] // for docs
+
+// Get returns a pointer to the value associated with the key, calling the
+// "create" function that was passed to [NewLive] as needed to construct values.
+func (c *Live[K, V]) Get(ctx context.Context, key K) (*V, error) {
+	for {
+		// Try to load an existing value out of the cache.
+		value, ok := c.m.Load(key)
+		if !ok {
+			// No value found. Create a new value.
+			fn := func() (*V, error) {
+				// Eagerly check the Context so that every create function
+				// doesn't need the preamble.
+				//
+				// Do this because this goroutine may have gone around the loop
+				// multiple times and found entries in the map that had
+				// invalidated weak pointers, so the context may have expired.
+				if ctx.Err() != nil {
+					return nil, context.Cause(ctx)
+				}
+				v, err := c.create(ctx, key)
+				if err != nil {
+					return nil, err
+				}
+
+				wp := weak.Make(v)
+				c.m.Store(key, wp)
+				runtime.AddCleanup(v, func(key K) {
+					// Only delete if the weak pointer is equal. If it's not,
+					// someone else already deleted the entry and installed a
+					// new pointer.
+					c.m.CompareAndDelete(key, wp)
+				}, key)
+				return v, nil
+			}
+
+			ch := c.sf.DoChan(key, fn)
+			select {
+			case res := <-ch:
+				return res.Val, res.Err
+			case <-ctx.Done():
+				c.sf.Forget(key)
+				return nil, context.Cause(ctx)
+			}
+		}
+
+		// See if our cache entry is valid.
+		if v := value.(weak.Pointer[V]).Value(); v != nil {
+			return v, nil
+		}
+		// Discovered a nil entry awaiting cleanup. Eagerly delete it.
+		c.m.CompareAndDelete(key, value)
+	}
+}

internal/singleflight/singleflight.go

@@ -0,0 +1,209 @@
+// Copyright 2013 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package singleflight provides a duplicate function call suppression
+// mechanism.
+//
+// This fork adds generics.
+//
+// Based on https://github.com/golang/sync/commit/22ba2078e183beec12908ea94f1d899c53dbf02c
+package singleflight // import "github.com/quay/claircore/internal/singleflight"
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"runtime"
+	"runtime/debug"
+	"sync"
+)
+
+// errGoexit indicates the runtime.Goexit was called in
+// the user given function.
+var errGoexit = errors.New("runtime.Goexit was called")
+
+// A panicError is an arbitrary value recovered from a panic
+// with the stack trace during the execution of given function.
+type panicError struct {
+	value any
+	stack []byte
+}
+
+// Error implements error interface.
+func (p *panicError) Error() string {
+	return fmt.Sprintf("%v\n\n%s", p.value, p.stack)
+}
+
+func newPanicError(v any) error {
+	stack := debug.Stack()
+
+	// The first line of the stack trace is of the form "goroutine N [status]:"
+	// but by the time the panic reaches Do the goroutine may no longer exist
+	// and its status will have changed. Trim out the misleading line.
+	if line := bytes.IndexByte(stack[:], '\n'); line >= 0 {
+		stack = stack[line+1:]
+	}
+	return &panicError{value: v, stack: stack}
+}
+
+// call is an in-flight or completed singleflight.Do call
+type call[V any] struct {
+	wg sync.WaitGroup
+
+	// These fields are written once before the WaitGroup is done
+	// and are only read after the WaitGroup is done.
+	val V
+	err error
+
+	// These fields are read and written with the singleflight
+	// mutex held before the WaitGroup is done, and are read but
+	// not written after the WaitGroup is done.
+	dups  int
+	chans []chan<- Result[V]
+}
+
+// Group represents a class of work and forms a namespace in
+// which units of work can be executed with duplicate suppression.
+type Group[K comparable, V any] struct {
+	mu sync.Mutex     // protects m
+	m  map[K]*call[V] // lazily initialized
+}
+
+// Result holds the results of Do, so they can be passed
+// on a channel.
+type Result[V any] struct {
+	Val    V
+	Err    error
+	Shared bool
+}
+
+// Do executes and returns the results of the given function, making
+// sure that only one execution is in-flight for a given key at a
+// time. If a duplicate comes in, the duplicate caller waits for the
+// original to complete and receives the same results.
+// The return value shared indicates whether v was given to multiple callers.
+func (g *Group[K, V]) Do(key K, fn func() (V, error)) (v V, err error, shared bool) {
+	g.mu.Lock()
+	if g.m == nil {
+		g.m = make(map[K]*call[V])
+	}
+	if c, ok := g.m[key]; ok {
+		c.dups++
+		g.mu.Unlock()
+		c.wg.Wait()
+
+		if e, ok := c.err.(*panicError); ok {
+			panic(e)
+		} else if c.err == errGoexit {
+			runtime.Goexit()
+		}
+		return c.val, c.err, true
+	}
+	c := new(call[V])
+	c.wg.Add(1)
+	g.m[key] = c
+	g.mu.Unlock()
+
+	g.doCall(c, key, fn)
+	return c.val, c.err, c.dups > 0
+}
+
+// DoChan is like Do but returns a channel that will receive the
+// results when they are ready.
+//
+// The returned channel will not be closed.
+func (g *Group[K, V]) DoChan(key K, fn func() (V, error)) <-chan Result[V] {
+	ch := make(chan Result[V], 1)
+	g.mu.Lock()
+	if g.m == nil {
+		g.m = make(map[K]*call[V])
+	}
+	if c, ok := g.m[key]; ok {
+		c.dups++
+		c.chans = append(c.chans, ch)
+		g.mu.Unlock()
+		return ch
+	}
+	c := &call[V]{chans: []chan<- Result[V]{ch}}
+	c.wg.Add(1)
+	g.m[key] = c
+	g.mu.Unlock()
+
+	go g.doCall(c, key, fn)
+
+	return ch
+}
+
+// doCall handles the single call for a key.
+func (g *Group[K, V]) doCall(c *call[V], key K, fn func() (V, error)) {
+	normalReturn := false
+	recovered := false
+
+	// use double-defer to distinguish panic from runtime.Goexit,
+	// more details see https://golang.org/cl/134395
+	defer func() {
+		// the given function invoked runtime.Goexit
+		if !normalReturn && !recovered {
+			c.err = errGoexit
+		}
+
+		g.mu.Lock()
+		defer g.mu.Unlock()
+		c.wg.Done()
+		if g.m[key] == c {
+			delete(g.m, key)
+		}
+
+		if e, ok := c.err.(*panicError); ok {
+			// In order to prevent the waiting channels from being blocked forever,
+			// needs to ensure that this panic cannot be recovered.
+			if len(c.chans) > 0 {
+				go panic(e)
+				select {} // Keep this goroutine around so that it will appear in the crash dump.
+			} else {
+				panic(e)
+			}
+		} else if c.err == errGoexit {
+			// Already in the process of goexit, no need to call again
+		} else {
+			// Normal return
+			for _, ch := range c.chans {
+				ch <- Result[V]{c.val, c.err, c.dups > 0}
+			}
+		}
+	}()
+
+	func() {
+		defer func() {
+			if !normalReturn {
+				// Ideally, we would wait to take a stack trace until we've determined
+				// whether this is a panic or a runtime.Goexit.
+				//
+				// Unfortunately, the only way we can distinguish the two is to see
+				// whether the recover stopped the goroutine from terminating, and by
+				// the time we know that, the part of the stack trace relevant to the
+				// panic has been discarded.
+				if r := recover(); r != nil {
+					c.err = newPanicError(r)
+				}
+			}
+		}()
+
+		c.val, c.err = fn()
+		normalReturn = true
+	}()
+
+	if !normalReturn {
+		recovered = true
+	}
+}
+
+// Forget tells the singleflight to forget about a key.  Future calls
+// to Do for this key will call the function rather than waiting for
+// an earlier call to complete.
+func (g *Group[K, V]) Forget(key K) {
+	g.mu.Lock()
+	delete(g.m, key)
+	g.mu.Unlock()
+}