Improve performance of controllers

controllers/actions.github.com/resourcecache.go

@@ -0,0 +1,230 @@
+package actionsgithubcom
+
+import (
+	"reflect"
+	"slices"
+	"strings"
+	"sync"
+
+	"github.com/actions/actions-runner-controller/apis/actions.github.com/v1alpha1"
+	"github.com/actions/actions-runner-controller/hash"
+	corev1 "k8s.io/api/core/v1"
+	rbacv1 "k8s.io/api/rbac/v1"
+	"k8s.io/apimachinery/pkg/types"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+)
+
+var resourceCacheObjectTypes sync.Map
+
+type ResourceCacheObjectRef struct {
+	ObjectType      string
+	Namespace       string
+	Name            string
+	UID             types.UID
+	ResourceVersion string
+}
+
+type ResourceCacheKey struct {
+	MainUID   types.UID
+	Namespace string
+	Name      string
+}
+
+type ResourceCacheValue[T client.Object] struct {
+	MainObject      ResourceCacheObjectRef
+	ResourceVersion string
+	Dependencies    []ResourceCacheObjectRef
+	Object          T
+}
+
+type ResourceCache struct {
+	autoscalingListener    *resourceCacheState[*v1alpha1.AutoscalingListener]
+	ephemeralRunnerSet     *resourceCacheState[*v1alpha1.EphemeralRunnerSet]
+	listenerPod            *resourceCacheState[*corev1.Pod]
+	listenerServiceAccount *resourceCacheState[*corev1.ServiceAccount]
+	listenerRole           *resourceCacheState[*rbacv1.Role]
+	listenerRoleBinding    *resourceCacheState[*rbacv1.RoleBinding]
+}
+
+func NewResourceCache() ResourceCache {
+	return ResourceCache{
+		autoscalingListener:    newResourceCacheState[*v1alpha1.AutoscalingListener](),
+		ephemeralRunnerSet:     newResourceCacheState[*v1alpha1.EphemeralRunnerSet](),
+		listenerPod:            newResourceCacheState[*corev1.Pod](),
+		listenerServiceAccount: newResourceCacheState[*corev1.ServiceAccount](),
+		listenerRole:           newResourceCacheState[*rbacv1.Role](),
+		listenerRoleBinding:    newResourceCacheState[*rbacv1.RoleBinding](),
+	}
+}
+
+type resourceCacheState[T client.Object] struct {
+	mu      sync.RWMutex
+	entries map[ResourceCacheKey]ResourceCacheValue[T]
+}
+
+func newResourceCacheState[T client.Object]() *resourceCacheState[T] {
+	return &resourceCacheState[T]{
+		entries: make(map[ResourceCacheKey]ResourceCacheValue[T], 512),
+	}
+}
+
+func (s *resourceCacheState[T]) Get(
+	mainObject client.Object,
+	desiredObject T,
+	dependencies ...client.Object,
+) (T, bool) {
+	key := newResourceCacheKey(mainObject, desiredObject)
+
+	s.mu.RLock()
+	value, ok := s.entries[key]
+	s.mu.RUnlock()
+	if !ok || !value.Matches(mainObject, dependencies...) {
+		var zero T
+		return zero, false
+	}
+
+	return cloneResourceCacheObject(value.Object), true
+}
+
+func (s *resourceCacheState[T]) Upsert(
+	mainObject client.Object,
+	desiredObject T,
+	dependencies ...client.Object,
+) (ResourceCacheValue[T], bool) {
+	key := newResourceCacheKey(mainObject, desiredObject)
+	mainObjectRef := newResourceCacheObjectRef(mainObject)
+	resourceVersion := desiredObject.GetResourceVersion()
+
+	s.mu.RLock()
+	previous, ok := s.entries[key]
+	if ok && previous.MainObject == mainObjectRef && previous.ResourceVersion == resourceVersion && previous.dependenciesMatch(dependencies...) {
+		s.mu.RUnlock()
+		return previous, false
+	}
+	s.mu.RUnlock()
+
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	previous, ok = s.entries[key]
+	if ok && previous.MainObject == mainObjectRef && previous.ResourceVersion == resourceVersion && previous.dependenciesMatch(dependencies...) {
+		return previous, false
+	}
+
+	dependencyRefs := newResourceCacheObjectRefs(dependencies...)
+	value := newResourceCacheValue(mainObjectRef, resourceVersion, dependencyRefs, cloneResourceCacheObject(desiredObject))
+	s.entries[key] = value
+	return value, true
+}
+
+func (v ResourceCacheValue[T]) Matches(mainObject client.Object, dependencies ...client.Object) bool {
+	if v.MainObject != newResourceCacheObjectRef(mainObject) {
+		return false
+	}
+
+	return v.dependenciesMatch(dependencies...)
+}
+
+func newResourceCacheKey(mainObject client.Object, desiredObject client.Object) ResourceCacheKey {
+	return ResourceCacheKey{
+		MainUID:   mainObject.GetUID(),
+		Namespace: desiredObject.GetNamespace(),
+		Name:      resourceCacheObjectName(desiredObject),
+	}
+}
+
+func newResourceCacheValue[T client.Object](
+	mainObjectRef ResourceCacheObjectRef,
+	resourceVersion string,
+	dependencyRefs []ResourceCacheObjectRef,
+	object T,
+) ResourceCacheValue[T] {
+	return ResourceCacheValue[T]{
+		MainObject:      mainObjectRef,
+		ResourceVersion: resourceVersion,
+		Dependencies:    dependencyRefs,
+		Object:          object,
+	}
+}
+
+func cloneResourceCacheObject[T client.Object](object T) T {
+	return object.DeepCopyObject().(T)
+}
+
+func newResourceCacheObjectRefs(objects ...client.Object) []ResourceCacheObjectRef {
+	refs := make([]ResourceCacheObjectRef, 0, len(objects))
+	for _, object := range objects {
+		refs = append(refs, newResourceCacheObjectRef(object))
+	}
+	slices.SortFunc(refs, func(a, b ResourceCacheObjectRef) int {
+		return compareResourceCacheObjectRefs(a, b)
+	})
+	return refs
+}
+
+func (v ResourceCacheValue[T]) dependenciesMatch(objects ...client.Object) bool {
+	if len(v.Dependencies) != len(objects) {
+		return false
+	}
+
+	for _, object := range objects {
+		ref := newResourceCacheObjectRef(object)
+		if !slices.Contains(v.Dependencies, ref) {
+			return false
+		}
+	}
+
+	return true
+}
+
+func newResourceCacheObjectRef(object client.Object) ResourceCacheObjectRef {
+	resourceVersion := object.GetResourceVersion()
+	if resourceVersion == "" {
+		resourceVersion = hash.ComputeTemplateHash(object)
+	}
+
+	return ResourceCacheObjectRef{
+		ObjectType:      resourceCacheObjectType(object),
+		Namespace:       object.GetNamespace(),
+		Name:            resourceCacheObjectName(object),
+		UID:             object.GetUID(),
+		ResourceVersion: resourceVersion,
+	}
+}
+
+func compareResourceCacheObjectRefs(a, b ResourceCacheObjectRef) int {
+	if c := strings.Compare(a.ObjectType, b.ObjectType); c != 0 {
+		return c
+	}
+	if c := strings.Compare(a.Namespace, b.Namespace); c != 0 {
+		return c
+	}
+	if c := strings.Compare(a.Name, b.Name); c != 0 {
+		return c
+	}
+	if c := strings.Compare(string(a.UID), string(b.UID)); c != 0 {
+		return c
+	}
+	return strings.Compare(a.ResourceVersion, b.ResourceVersion)
+}
+
+func resourceCacheObjectType(object client.Object) string {
+	t := reflect.TypeOf(object)
+	if t.Kind() == reflect.Pointer {
+		t = t.Elem()
+	}
+	if objectType, ok := resourceCacheObjectTypes.Load(t); ok {
+		return objectType.(string)
+	}
+
+	objectType := t.PkgPath() + "." + t.Name()
+	actual, _ := resourceCacheObjectTypes.LoadOrStore(t, objectType)
+	return actual.(string)
+}
+
+func resourceCacheObjectName(object client.Object) string {
+	if object.GetName() != "" {
+		return object.GetName()
+	}
+	return object.GetGenerateName()
+}