github.com/waynz0r/controller-tools@v0.4.1-0.20200916220028-16254aeef2d7/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  //  r.Register(MakeDefinition(name, target, obj))
    56  func (r *Registry) Define(name string, target TargetType, obj interface{}) error {
    57  	def, err := MakeDefinition(name, target, obj)
    58  	if err != nil {
    59  		return err
    60  	}
    61  	return r.Register(def)
    62  }
    63  
    64  // Register registers the given marker definition with this registry for later lookup.
    65  func (r *Registry) Register(def *Definition) error {
    66  	r.init()
    67  
    68  	r.mu.Lock()
    69  	defer r.mu.Unlock()
    70  
    71  	switch def.Target {
    72  	case DescribesPackage:
    73  		r.forPkg[def.Name] = def
    74  	case DescribesType:
    75  		r.forType[def.Name] = def
    76  	case DescribesField:
    77  		r.forField[def.Name] = def
    78  	default:
    79  		return fmt.Errorf("unknown target type %v", def.Target)
    80  	}
    81  	return nil
    82  }
    83  
    84  // AddHelp stores the given help in the registry, marking it as associated with
    85  // the given definition.
    86  func (r *Registry) AddHelp(def *Definition, help *DefinitionHelp) {
    87  	r.init()
    88  
    89  	r.mu.Lock()
    90  	defer r.mu.Unlock()
    91  
    92  	r.helpFor[def] = help
    93  }
    94  
    95  // Lookup fetches the definition corresponding to the given name and target type.
    96  func (r *Registry) Lookup(name string, target TargetType) *Definition {
    97  	r.init()
    98  
    99  	r.mu.RLock()
   100  	defer r.mu.RUnlock()
   101  
   102  	switch target {
   103  	case DescribesPackage:
   104  		return tryAnonLookup(name, r.forPkg)
   105  	case DescribesType:
   106  		return tryAnonLookup(name, r.forType)
   107  	case DescribesField:
   108  		return tryAnonLookup(name, r.forField)
   109  	default:
   110  		return nil
   111  	}
   112  }
   113  
   114  // HelpFor fetches the help for a given definition, if present.
   115  func (r *Registry) HelpFor(def *Definition) *DefinitionHelp {
   116  	r.init()
   117  
   118  	r.mu.RLock()
   119  	defer r.mu.RUnlock()
   120  
   121  	return r.helpFor[def]
   122  }
   123  
   124  // AllDefinitions returns all marker definitions known to this registry.
   125  func (r *Registry) AllDefinitions() []*Definition {
   126  	res := make([]*Definition, 0, len(r.forPkg)+len(r.forType)+len(r.forField))
   127  	for _, def := range r.forPkg {
   128  		res = append(res, def)
   129  	}
   130  	for _, def := range r.forType {
   131  		res = append(res, def)
   132  	}
   133  	for _, def := range r.forField {
   134  		res = append(res, def)
   135  	}
   136  	return res
   137  }
   138  
   139  // tryAnonLookup tries looking up the given marker as both an struct-based
   140  // marker and an anonymous marker, returning whichever format matches first,
   141  // preferring the longer (anonymous) name in case of conflicts.
   142  func tryAnonLookup(name string, defs map[string]*Definition) *Definition {
   143  	// NB(directxman12): we look up anonymous names first to work with
   144  	// legacy style marker definitions that have a namespaced approach
   145  	// (e.g. deepcopy-gen, which uses `+k8s:deepcopy-gen=foo,bar` *and*
   146  	// `+k8s.io:deepcopy-gen:interfaces=foo`).
   147  	name, anonName, _ := splitMarker(name)
   148  	if def, exists := defs[anonName]; exists {
   149  		return def
   150  	}
   151  
   152  	return defs[name]
   153  }