github.com/alex123012/deckhouse-controller-tools@v0.0.0-20230510090815-d594daf1af8c/pkg/markers/reg.go (about)

     1  /*
     2  Copyright 2019 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 markers
    18  
    19  import (
    20  	"fmt"
    21  	"sync"
    22  )
    23  
    24  // Registry keeps track of registered definitions, and allows for easy lookup.
    25  // It's thread-safe, and the zero-value can be safely used.
    26  type Registry struct {
    27  	forPkg   map[string]*Definition
    28  	forType  map[string]*Definition
    29  	forField map[string]*Definition
    30  	helpFor  map[*Definition]*DefinitionHelp
    31  
    32  	mu       sync.RWMutex
    33  	initOnce sync.Once
    34  }
    35  
    36  func (r *Registry) init() {
    37  	r.initOnce.Do(func() {
    38  		if r.forPkg == nil {
    39  			r.forPkg = make(map[string]*Definition)
    40  		}
    41  		if r.forType == nil {
    42  			r.forType = make(map[string]*Definition)
    43  		}
    44  		if r.forField == nil {
    45  			r.forField = make(map[string]*Definition)
    46  		}
    47  		if r.helpFor == nil {
    48  			r.helpFor = make(map[*Definition]*DefinitionHelp)
    49  		}
    50  	})
    51  }
    52  
    53  // Define defines a new marker with the given name, target, and output type.
    54  // It's a shortcut around
    55  //
    56  //	r.Register(MakeDefinition(name, target, obj))
    57  func (r *Registry) Define(name string, target TargetType, obj interface{}) error {
    58  	def, err := MakeDefinition(name, target, obj)
    59  	if err != nil {
    60  		return err
    61  	}
    62  	return r.Register(def)
    63  }
    64  
    65  // Register registers the given marker definition with this registry for later lookup.
    66  func (r *Registry) Register(def *Definition) error {
    67  	r.init()
    68  
    69  	r.mu.Lock()
    70  	defer r.mu.Unlock()
    71  
    72  	switch def.Target {
    73  	case DescribesPackage:
    74  		r.forPkg[def.Name] = def
    75  	case DescribesType:
    76  		r.forType[def.Name] = def
    77  	case DescribesField:
    78  		r.forField[def.Name] = def
    79  	default:
    80  		return fmt.Errorf("unknown target type %v", def.Target)
    81  	}
    82  	return nil
    83  }
    84  
    85  // AddHelp stores the given help in the registry, marking it as associated with
    86  // the given definition.
    87  func (r *Registry) AddHelp(def *Definition, help *DefinitionHelp) {
    88  	r.init()
    89  
    90  	r.mu.Lock()
    91  	defer r.mu.Unlock()
    92  
    93  	r.helpFor[def] = help
    94  }
    95  
    96  // Lookup fetches the definition corresponding to the given name and target type.
    97  func (r *Registry) Lookup(name string, target TargetType) *Definition {
    98  	r.init()
    99  
   100  	r.mu.RLock()
   101  	defer r.mu.RUnlock()
   102  
   103  	switch target {
   104  	case DescribesPackage:
   105  		return tryAnonLookup(name, r.forPkg)
   106  	case DescribesType:
   107  		return tryAnonLookup(name, r.forType)
   108  	case DescribesField:
   109  		return tryAnonLookup(name, r.forField)
   110  	default:
   111  		return nil
   112  	}
   113  }
   114  
   115  // HelpFor fetches the help for a given definition, if present.
   116  func (r *Registry) HelpFor(def *Definition) *DefinitionHelp {
   117  	r.init()
   118  
   119  	r.mu.RLock()
   120  	defer r.mu.RUnlock()
   121  
   122  	return r.helpFor[def]
   123  }
   124  
   125  // AllDefinitions returns all marker definitions known to this registry.
   126  func (r *Registry) AllDefinitions() []*Definition {
   127  	res := make([]*Definition, 0, len(r.forPkg)+len(r.forType)+len(r.forField))
   128  	for _, def := range r.forPkg {
   129  		res = append(res, def)
   130  	}
   131  	for _, def := range r.forType {
   132  		res = append(res, def)
   133  	}
   134  	for _, def := range r.forField {
   135  		res = append(res, def)
   136  	}
   137  	return res
   138  }
   139  
   140  // tryAnonLookup tries looking up the given marker as both an struct-based
   141  // marker and an anonymous marker, returning whichever format matches first,
   142  // preferring the longer (anonymous) name in case of conflicts.
   143  func tryAnonLookup(name string, defs map[string]*Definition) *Definition {
   144  	// NB(directxman12): we look up anonymous names first to work with
   145  	// legacy style marker definitions that have a namespaced approach
   146  	// (e.g. deepcopy-gen, which uses `+k8s:deepcopy-gen=foo,bar` *and*
   147  	// `+k8s.io:deepcopy-gen:interfaces=foo`).
   148  	name, anonName, _ := splitMarker(name)
   149  	if def, exists := defs[anonName]; exists {
   150  		return def
   151  	}
   152  
   153  	return defs[name]
   154  }