bimultimap.go

Documentation: k8s.io/kubernetes/pkg/controller/util/selectors

     1  /*
     2  Copyright 2022 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package selectors
    18  
    19  import (
    20  	"fmt"
    21  	"strings"
    22  	"sync"
    23  
    24  	pkglabels "k8s.io/apimachinery/pkg/labels"
    25  )
    26  
    27  // BiMultimap is an efficient, bi-directional mapping of object
    28  // keys. Associations are created by putting keys with a selector.
    29  type BiMultimap struct {
    30  	mux sync.RWMutex
    31  
    32  	// Objects.
    33  	labeledObjects   map[Key]*labeledObject
    34  	selectingObjects map[Key]*selectingObject
    35  
    36  	// Associations.
    37  	labeledBySelecting map[selectorKey]*labeledObjects
    38  	selectingByLabeled map[labelsKey]*selectingObjects
    39  }
    40  
    41  // NewBiMultimap creates a map.
    42  func NewBiMultimap() *BiMultimap {
    43  	return &BiMultimap{
    44  		labeledObjects:     make(map[Key]*labeledObject),
    45  		selectingObjects:   make(map[Key]*selectingObject),
    46  		labeledBySelecting: make(map[selectorKey]*labeledObjects),
    47  		selectingByLabeled: make(map[labelsKey]*selectingObjects),
    48  	}
    49  }
    50  
    51  // Key is a tuple of name and namespace.
    52  type Key struct {
    53  	Name      string
    54  	Namespace string
    55  }
    56  
    57  // Parse turns a string in the format namespace/name into a Key.
    58  func Parse(s string) (key Key) {
    59  	ns := strings.SplitN(s, "/", 2)
    60  	if len(ns) == 2 {
    61  		key.Namespace = ns[0]
    62  		key.Name = ns[1]
    63  	} else {
    64  		key.Name = ns[0]
    65  	}
    66  	return key
    67  }
    68  
    69  func (k Key) String() string {
    70  	return fmt.Sprintf("%v/%v", k.Namespace, k.Name)
    71  }
    72  
    73  type selectorKey struct {
    74  	key       string
    75  	namespace string
    76  }
    77  
    78  type selectingObject struct {
    79  	key      Key
    80  	selector pkglabels.Selector
    81  	// selectorKey is a stable serialization of selector for
    82  	// association caching.
    83  	selectorKey selectorKey
    84  }
    85  
    86  type selectingObjects struct {
    87  	objects  map[Key]*selectingObject
    88  	refCount int
    89  }
    90  
    91  type labelsKey struct {
    92  	key       string
    93  	namespace string
    94  }
    95  
    96  type labeledObject struct {
    97  	key    Key
    98  	labels map[string]string
    99  	// labelsKey is a stable serialization of labels for association
   100  	// caching.
   101  	labelsKey labelsKey
   102  }
   103  
   104  type labeledObjects struct {
   105  	objects  map[Key]*labeledObject
   106  	refCount int
   107  }
   108  
   109  // Put inserts or updates an object and the incoming associations
   110  // based on the object labels.
   111  func (m *BiMultimap) Put(key Key, labels map[string]string) {
   112  	m.mux.Lock()
   113  	defer m.mux.Unlock()
   114  
   115  	labelsKey := labelsKey{
   116  		key:       pkglabels.Set(labels).String(),
   117  		namespace: key.Namespace,
   118  	}
   119  	if l, ok := m.labeledObjects[key]; ok {
   120  		// Update labeled object.
   121  		if labelsKey == l.labelsKey {
   122  			// No change to labels.
   123  			return
   124  		}
   125  		// Delete before readding.
   126  		m.delete(key)
   127  	}
   128  	// Add labeled object.
   129  	labels = copyLabels(labels)
   130  	labeledObject := &labeledObject{
   131  		key:       key,
   132  		labels:    labels,
   133  		labelsKey: labelsKey,
   134  	}
   135  	m.labeledObjects[key] = labeledObject
   136  	// Add associations.
   137  	if _, ok := m.selectingByLabeled[labelsKey]; !ok {
   138  		// Cache miss. Scan selecting objects.
   139  		selecting := &selectingObjects{
   140  			objects: make(map[Key]*selectingObject),
   141  		}
   142  		set := pkglabels.Set(labels)
   143  		for _, s := range m.selectingObjects {
   144  			if s.key.Namespace != key.Namespace {
   145  				continue
   146  			}
   147  			if s.selector.Matches(set) {
   148  				selecting.objects[s.key] = s
   149  			}
   150  		}
   151  		// Associate selecting with labeled.
   152  		m.selectingByLabeled[labelsKey] = selecting
   153  	}
   154  	selecting := m.selectingByLabeled[labelsKey]
   155  	selecting.refCount++
   156  	for _, sObject := range selecting.objects {
   157  		// Associate labeled with selecting.
   158  		labeled := m.labeledBySelecting[sObject.selectorKey]
   159  		labeled.objects[labeledObject.key] = labeledObject
   160  	}
   161  }
   162  
   163  // Delete removes a labeled object and incoming associations.
   164  func (m *BiMultimap) Delete(key Key) {
   165  	m.mux.Lock()
   166  	defer m.mux.Unlock()
   167  	m.delete(key)
   168  }
   169  
   170  func (m *BiMultimap) delete(key Key) {
   171  	if _, ok := m.labeledObjects[key]; !ok {
   172  		// Does not exist.
   173  		return
   174  	}
   175  	labeledObject := m.labeledObjects[key]
   176  	labelsKey := labeledObject.labelsKey
   177  	defer delete(m.labeledObjects, key)
   178  	if _, ok := m.selectingByLabeled[labelsKey]; !ok {
   179  		// No associations.
   180  		return
   181  	}
   182  	// Remove associations.
   183  	for _, selectingObject := range m.selectingByLabeled[labelsKey].objects {
   184  		selectorKey := selectingObject.selectorKey
   185  		// Delete selectingObject to labeledObject association.
   186  		delete(m.labeledBySelecting[selectorKey].objects, key)
   187  	}
   188  	m.selectingByLabeled[labelsKey].refCount--
   189  	// Garbage collect labeledObject to selectingObject associations.
   190  	if m.selectingByLabeled[labelsKey].refCount == 0 {
   191  		delete(m.selectingByLabeled, labelsKey)
   192  	}
   193  }
   194  
   195  // Exists returns true if the labeled object is present in the map.
   196  func (m *BiMultimap) Exists(key Key) bool {
   197  	m.mux.RLock()
   198  	defer m.mux.RUnlock()
   199  
   200  	_, exists := m.labeledObjects[key]
   201  	return exists
   202  }
   203  
   204  // PutSelector inserts or updates an object with a selector. Associations
   205  // are created or updated based on the selector.
   206  func (m *BiMultimap) PutSelector(key Key, selector pkglabels.Selector) {
   207  	m.mux.Lock()
   208  	defer m.mux.Unlock()
   209  
   210  	selectorKey := selectorKey{
   211  		key:       selector.String(),
   212  		namespace: key.Namespace,
   213  	}
   214  	if s, ok := m.selectingObjects[key]; ok {
   215  		// Update selecting object.
   216  		if selectorKey == s.selectorKey {
   217  			// No change to selector.
   218  			return
   219  		}
   220  		// Delete before readding.
   221  		m.deleteSelector(key)
   222  	}
   223  	// Add selecting object.
   224  	selectingObject := &selectingObject{
   225  		key:         key,
   226  		selector:    selector,
   227  		selectorKey: selectorKey,
   228  	}
   229  	m.selectingObjects[key] = selectingObject
   230  	// Add associations.
   231  	if _, ok := m.labeledBySelecting[selectorKey]; !ok {
   232  		// Cache miss. Scan labeled objects.
   233  		labeled := &labeledObjects{
   234  			objects: make(map[Key]*labeledObject),
   235  		}
   236  		for _, l := range m.labeledObjects {
   237  			if l.key.Namespace != key.Namespace {
   238  				continue
   239  			}
   240  			set := pkglabels.Set(l.labels)
   241  			if selector.Matches(set) {
   242  				labeled.objects[l.key] = l
   243  			}
   244  		}
   245  		// Associate labeled with selecting.
   246  		m.labeledBySelecting[selectorKey] = labeled
   247  	}
   248  	labeled := m.labeledBySelecting[selectorKey]
   249  	labeled.refCount++
   250  	for _, labeledObject := range labeled.objects {
   251  		// Associate selecting with labeled.
   252  		selecting := m.selectingByLabeled[labeledObject.labelsKey]
   253  		selecting.objects[selectingObject.key] = selectingObject
   254  	}
   255  }
   256  
   257  // DeleteSelector deletes a selecting object and associations created by its
   258  // selector.
   259  func (m *BiMultimap) DeleteSelector(key Key) {
   260  	m.mux.Lock()
   261  	defer m.mux.Unlock()
   262  	m.deleteSelector(key)
   263  }
   264  
   265  func (m *BiMultimap) deleteSelector(key Key) {
   266  	if _, ok := m.selectingObjects[key]; !ok {
   267  		// Does not exist.
   268  		return
   269  	}
   270  	selectingObject := m.selectingObjects[key]
   271  	selectorKey := selectingObject.selectorKey
   272  	defer delete(m.selectingObjects, key)
   273  	if _, ok := m.labeledBySelecting[selectorKey]; !ok {
   274  		// No associations.
   275  		return
   276  	}
   277  	// Remove associations.
   278  	for _, labeledObject := range m.labeledBySelecting[selectorKey].objects {
   279  		labelsKey := labeledObject.labelsKey
   280  		// Delete labeledObject to selectingObject association.
   281  		delete(m.selectingByLabeled[labelsKey].objects, key)
   282  	}
   283  	m.labeledBySelecting[selectorKey].refCount--
   284  	// Garbage collect selectingObjects to labeledObject associations.
   285  	if m.labeledBySelecting[selectorKey].refCount == 0 {
   286  		delete(m.labeledBySelecting, selectorKey)
   287  	}
   288  }
   289  
   290  // SelectorExists returns true if the selecting object is present in the map.
   291  func (m *BiMultimap) SelectorExists(key Key) bool {
   292  	m.mux.RLock()
   293  	defer m.mux.RUnlock()
   294  
   295  	_, exists := m.selectingObjects[key]
   296  	return exists
   297  }
   298  
   299  // KeepOnly retains only the specified labeled objects and deletes the
   300  // rest. Like calling Delete for all keys not specified.
   301  func (m *BiMultimap) KeepOnly(keys []Key) {
   302  	m.mux.Lock()
   303  	defer m.mux.Unlock()
   304  
   305  	keyMap := make(map[Key]bool)
   306  	for _, k := range keys {
   307  		keyMap[k] = true
   308  	}
   309  	for k := range m.labeledObjects {
   310  		if !keyMap[k] {
   311  			m.delete(k)
   312  		}
   313  	}
   314  }
   315  
   316  // KeepOnlySelectors retains only the specified selecting objects and
   317  // deletes the rest. Like calling DeleteSelector for all keys not
   318  // specified.
   319  func (m *BiMultimap) KeepOnlySelectors(keys []Key) {
   320  	m.mux.Lock()
   321  	defer m.mux.Unlock()
   322  
   323  	keyMap := make(map[Key]bool)
   324  	for _, k := range keys {
   325  		keyMap[k] = true
   326  	}
   327  	for k := range m.selectingObjects {
   328  		if !keyMap[k] {
   329  			m.deleteSelector(k)
   330  		}
   331  	}
   332  }
   333  
   334  // Select finds objects associated with a selecting object. If the
   335  // given key was found in the map `ok` will be true. Otherwise false.
   336  func (m *BiMultimap) Select(key Key) (keys []Key, ok bool) {
   337  	m.mux.RLock()
   338  	defer m.mux.RUnlock()
   339  
   340  	selectingObject, ok := m.selectingObjects[key]
   341  	if !ok {
   342  		// Does not exist.
   343  		return nil, false
   344  	}
   345  	keys = make([]Key, 0)
   346  	if labeled, ok := m.labeledBySelecting[selectingObject.selectorKey]; ok {
   347  		for _, labeledObject := range labeled.objects {
   348  			keys = append(keys, labeledObject.key)
   349  		}
   350  	}
   351  	return keys, true
   352  }
   353  
   354  // ReverseSelect finds objects selecting the given object. If the
   355  // given key was found in the map `ok` will be true. Otherwise false.
   356  func (m *BiMultimap) ReverseSelect(key Key) (keys []Key, ok bool) {
   357  	m.mux.RLock()
   358  	defer m.mux.RUnlock()
   359  
   360  	labeledObject, ok := m.labeledObjects[key]
   361  	if !ok {
   362  		// Does not exist.
   363  		return []Key{}, false
   364  	}
   365  	keys = make([]Key, 0)
   366  	if selecting, ok := m.selectingByLabeled[labeledObject.labelsKey]; ok {
   367  		for _, selectingObject := range selecting.objects {
   368  			keys = append(keys, selectingObject.key)
   369  		}
   370  	}
   371  	return keys, true
   372  }
   373  
   374  func copyLabels(labels map[string]string) map[string]string {
   375  	l := make(map[string]string)
   376  	for k, v := range labels {
   377  		l[k] = v
   378  	}
   379  	return l
   380  }
   381
View as plain text