istio.io/istio@v0.0.0-20240520182934-d79c90f27776/pkg/kube/krt/index.go (about)

     1  // Copyright Istio Authors
     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 krt
    16  
    17  import (
    18  	"sync"
    19  
    20  	"istio.io/istio/pkg/ptr"
    21  	"istio.io/istio/pkg/util/sets"
    22  )
    23  
    24  // Index maintains a simple index over an informer
    25  type Index[I any, K comparable] struct {
    26  	mu      sync.RWMutex
    27  	objects map[K]sets.Set[Key[I]]
    28  	c       Collection[I]
    29  	extract func(o I) []K
    30  }
    31  
    32  // Lookup finds all objects matching a given key
    33  func (i *Index[I, K]) Lookup(k K) []I {
    34  	i.mu.RLock()
    35  	defer i.mu.RUnlock()
    36  	var res []I
    37  	for obj := range i.objects[k] {
    38  		item := i.c.GetKey(obj)
    39  		if item == nil {
    40  			// This should be extremely rare, but possible. While we have a mutex here, the underlying collection
    41  			// is not locked and maybe have changed in the meantime.
    42  			log.Debugf("missing item for %v", obj)
    43  			continue
    44  		}
    45  		res = append(res, *item)
    46  	}
    47  	return res
    48  }
    49  
    50  func (i *Index[I, K]) objectHasKey(obj I, k K) bool {
    51  	for _, got := range i.extract(obj) {
    52  		if got == k {
    53  			return true
    54  		}
    55  	}
    56  	return false
    57  }
    58  
    59  func (i *Index[I, K]) Dump() {
    60  	i.mu.RLock()
    61  	defer i.mu.RUnlock()
    62  	log.Errorf("> BEGIN DUMP (index %v[%T])", i.c.(internalCollection[I]).name(), ptr.TypeName[K]())
    63  	for k, v := range i.objects {
    64  		log.Errorf("key %v: %v", k, v.UnsortedList())
    65  	}
    66  	log.Errorf("< END DUMP (index %v[%T]", i.c.(internalCollection[I]).name(), ptr.TypeName[K]())
    67  }
    68  
    69  // NewNamespaceIndex is a small helper to index a collection by namespace
    70  func NewNamespaceIndex[I Namespacer](c Collection[I]) *Index[I, string] {
    71  	return NewIndex(c, func(o I) []string {
    72  		return []string{o.GetNamespace()}
    73  	})
    74  }
    75  
    76  // NewIndex creates a simple index, keyed by key K, over an informer for O. This is similar to
    77  // NewInformer.AddIndex, but is easier to use and can be added after an informer has already started.
    78  func NewIndex[I any, K comparable](
    79  	c Collection[I],
    80  	extract func(o I) []K,
    81  ) *Index[I, K] {
    82  	idx := Index[I, K]{
    83  		objects: make(map[K]sets.Set[Key[I]]),
    84  		c:       c,
    85  		mu:      sync.RWMutex{},
    86  		extract: extract,
    87  	}
    88  	c.Register(func(o Event[I]) {
    89  		idx.mu.Lock()
    90  		defer idx.mu.Unlock()
    91  
    92  		if o.Old != nil {
    93  			obj := *o.Old
    94  			key := GetKey(obj)
    95  			for _, indexKey := range extract(obj) {
    96  				sets.DeleteCleanupLast(idx.objects, indexKey, key)
    97  			}
    98  		}
    99  		if o.New != nil {
   100  			obj := *o.New
   101  			key := GetKey(obj)
   102  			for _, indexKey := range extract(obj) {
   103  				sets.InsertOrNew(idx.objects, indexKey, key)
   104  			}
   105  		}
   106  	})
   107  
   108  	return &idx
   109  }