github.com/tilt-dev/tilt@v0.33.15-0.20240515162809-0a22ed45d8a0/internal/controllers/indexer/indexer.go (about)

     1  package indexer
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync"
     7  
     8  	"k8s.io/apimachinery/pkg/runtime"
     9  	"k8s.io/apimachinery/pkg/runtime/schema"
    10  	"k8s.io/apimachinery/pkg/types"
    11  	"sigs.k8s.io/controller-runtime/pkg/client"
    12  	"sigs.k8s.io/controller-runtime/pkg/client/apiutil"
    13  	"sigs.k8s.io/controller-runtime/pkg/reconcile"
    14  )
    15  
    16  // A key to help index objects we watch.
    17  type Key struct {
    18  	Name types.NamespacedName
    19  	GVK  schema.GroupVersionKind
    20  }
    21  
    22  type KeyFunc func(obj client.Object) []Key
    23  
    24  // Helper struct to help reconcilers determine when
    25  // to start their objects when a dependency triggers.
    26  type Indexer struct {
    27  	scheme *runtime.Scheme
    28  
    29  	indexFuncs []KeyFunc
    30  
    31  	// A map to help determine which Objects to reconcile when one of the objects
    32  	// they're watching change.
    33  	//
    34  	// The first key is the name and type of the object being watched.
    35  	//
    36  	// The second key is the name of the main object being reconciled.
    37  	//
    38  	// For example, if a Cmd is triggered by a FileWatch, the first
    39  	// key is the FileWatch name and GVK, while the second key is the Cmd name.
    40  	indexByWatchedObjects map[Key]map[types.NamespacedName]bool
    41  
    42  	mu sync.Mutex
    43  }
    44  
    45  func NewIndexer(scheme *runtime.Scheme, keyFuncs ...KeyFunc) *Indexer {
    46  	return &Indexer{
    47  		scheme:                scheme,
    48  		indexFuncs:            keyFuncs,
    49  		indexByWatchedObjects: make(map[Key]map[types.NamespacedName]bool),
    50  	}
    51  }
    52  
    53  // Register the watched object for the given primary object.
    54  func (m *Indexer) OnReconcile(name types.NamespacedName, obj client.Object) {
    55  	m.mu.Lock()
    56  	defer m.mu.Unlock()
    57  
    58  	// Delete all the mappings for this object.
    59  	for _, index := range m.indexByWatchedObjects {
    60  		delete(index, name)
    61  	}
    62  
    63  	// Re-add all the mappings.
    64  	for _, indexFunc := range m.indexFuncs {
    65  		for _, key := range indexFunc(obj) {
    66  			index, ok := m.indexByWatchedObjects[key]
    67  			if !ok {
    68  				index = make(map[types.NamespacedName]bool)
    69  				m.indexByWatchedObjects[key] = index
    70  			}
    71  
    72  			index[name] = true
    73  		}
    74  	}
    75  }
    76  
    77  // Given an update of a watched object, return the names of objects watching it
    78  // that we need to reconcile.
    79  func (m *Indexer) Enqueue(ctx context.Context, obj client.Object) []reconcile.Request {
    80  	gvk, err := apiutil.GVKForObject(obj, m.scheme)
    81  	if err != nil {
    82  		panic(fmt.Sprintf("Unrecognized object: %v", err))
    83  	}
    84  	key := Key{
    85  		Name: types.NamespacedName{Name: obj.GetName(), Namespace: obj.GetNamespace()},
    86  		GVK:  gvk,
    87  	}
    88  	return m.EnqueueKey(key)
    89  }
    90  
    91  // EnqueueKey() when we don't have the full object, only the name and kind.
    92  func (m *Indexer) EnqueueKey(key Key) []reconcile.Request {
    93  	m.mu.Lock()
    94  	defer m.mu.Unlock()
    95  
    96  	result := make([]reconcile.Request, 0, len(m.indexByWatchedObjects[key]))
    97  	for watchingName := range m.indexByWatchedObjects[key] {
    98  		result = append(result, reconcile.Request{NamespacedName: watchingName})
    99  	}
   100  	return result
   101  }
   102  
   103  // AddKeyFunc registers a new indexer function.
   104  //
   105  // In practice, all KeyFunc indexer functions should be added before or during controller initialization
   106  // to avoid missed updates.
   107  func (m *Indexer) AddKeyFunc(fn KeyFunc) {
   108  	m.mu.Lock()
   109  	defer m.mu.Unlock()
   110  
   111  	m.indexFuncs = append(m.indexFuncs, fn)
   112  }