k3s/vendor/github.com/rancher/lasso/pkg/controller/sharedcontrollerfactory.go

package controller

import (
	"context"
	"sync"

	"github.com/rancher/lasso/pkg/cache"
	"github.com/rancher/lasso/pkg/client"
	"k8s.io/apimachinery/pkg/api/meta"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/util/workqueue"
)

type SharedControllerFactory interface {
	ForObject(obj runtime.Object) (SharedController, error)
	ForKind(gvk schema.GroupVersionKind) (SharedController, error)
	ForResource(gvr schema.GroupVersionResource, namespaced bool) SharedController
	ForResourceKind(gvr schema.GroupVersionResource, kind string, namespaced bool) SharedController
	SharedCacheFactory() cache.SharedCacheFactory
	Start(ctx context.Context, workers int) error
}

type SharedControllerFactoryOptions struct {
	DefaultRateLimiter workqueue.RateLimiter
	DefaultWorkers     int

	KindRateLimiter map[schema.GroupVersionKind]workqueue.RateLimiter
	KindWorkers     map[schema.GroupVersionKind]int
}

type sharedControllerFactory struct {
	controllerLock sync.RWMutex

	sharedCacheFactory cache.SharedCacheFactory
	controllers        map[schema.GroupVersionResource]*sharedController

	rateLimiter     workqueue.RateLimiter
	workers         int
	kindRateLimiter map[schema.GroupVersionKind]workqueue.RateLimiter
	kindWorkers     map[schema.GroupVersionKind]int
}

func NewSharedControllerFactoryFromConfig(config *rest.Config, scheme *runtime.Scheme) (SharedControllerFactory, error) {
	cf, err := client.NewSharedClientFactory(config, &client.SharedClientFactoryOptions{
		Scheme: scheme,
	})
	if err != nil {
		return nil, err
	}
	return NewSharedControllerFactory(cache.NewSharedCachedFactory(cf, nil), nil), nil
}

func NewSharedControllerFactory(cacheFactory cache.SharedCacheFactory, opts *SharedControllerFactoryOptions) SharedControllerFactory {
	opts = applyDefaultSharedOptions(opts)
	return &sharedControllerFactory{
		sharedCacheFactory: cacheFactory,
		controllers:        map[schema.GroupVersionResource]*sharedController{},
		workers:            opts.DefaultWorkers,
		kindWorkers:        opts.KindWorkers,
		rateLimiter:        opts.DefaultRateLimiter,
		kindRateLimiter:    opts.KindRateLimiter,
	}
}

func applyDefaultSharedOptions(opts *SharedControllerFactoryOptions) *SharedControllerFactoryOptions {
	var newOpts SharedControllerFactoryOptions
	if opts != nil {
		newOpts = *opts
	}
	if newOpts.DefaultWorkers == 0 {
		newOpts.DefaultWorkers = 5
	}
	return &newOpts
}

func (s *sharedControllerFactory) Start(ctx context.Context, defaultWorkers int) error {
	s.controllerLock.Lock()
	defer s.controllerLock.Unlock()

	if err := s.sharedCacheFactory.Start(ctx); err != nil {
		return err
	}

	s.sharedCacheFactory.WaitForCacheSync(ctx)

	for gvr, controller := range s.controllers {
		w, err := s.getWorkers(gvr, defaultWorkers)
		if err != nil {
			return err
		}
		if err := controller.Start(ctx, w); err != nil {
			return err
		}
	}

	return nil
}

func (s *sharedControllerFactory) ForObject(obj runtime.Object) (SharedController, error) {
	gvk, err := s.sharedCacheFactory.SharedClientFactory().GVKForObject(obj)
	if err != nil {
		return nil, err
	}
	return s.ForKind(gvk)
}

func (s *sharedControllerFactory) ForKind(gvk schema.GroupVersionKind) (SharedController, error) {
	gvr, nsed, err := s.sharedCacheFactory.SharedClientFactory().ResourceForGVK(gvk)
	if err != nil {
		return nil, err
	}

	return s.ForResourceKind(gvr, gvk.Kind, nsed), nil
}

func (s *sharedControllerFactory) ForResource(gvr schema.GroupVersionResource, namespaced bool) SharedController {
	return s.ForResourceKind(gvr, "", namespaced)
}

func (s *sharedControllerFactory) ForResourceKind(gvr schema.GroupVersionResource, kind string, namespaced bool) SharedController {
	controllerResult := s.byResource(gvr)
	if controllerResult != nil {
		return controllerResult
	}

	s.controllerLock.Lock()
	defer s.controllerLock.Unlock()

	controllerResult = s.controllers[gvr]
	if controllerResult != nil {
		return controllerResult
	}

	client := s.sharedCacheFactory.SharedClientFactory().ForResourceKind(gvr, kind, namespaced)

	handler := &SharedHandler{}

	controllerResult = &sharedController{
		deferredController: func() (Controller, error) {
			var (
				gvk schema.GroupVersionKind
				err error
			)

			if kind == "" {
				gvk, err = s.sharedCacheFactory.SharedClientFactory().GVKForResource(gvr)
				if err != nil {
					return nil, err
				}
			} else {
				gvk = gvr.GroupVersion().WithKind(kind)
			}

			cache, err := s.sharedCacheFactory.ForResourceKind(gvr, kind, namespaced)
			if err != nil {
				return nil, err
			}

			rateLimiter, ok := s.kindRateLimiter[gvk]
			if !ok {
				rateLimiter = s.rateLimiter
			}

			starter := func(ctx context.Context) error {
				return s.sharedCacheFactory.StartGVK(ctx, gvk)
			}

			c := New(gvk.String(), cache, starter, handler, &Options{
				RateLimiter: rateLimiter,
			})

			return c, err
		},
		handler: handler,
		client:  client,
	}

	s.controllers[gvr] = controllerResult
	return controllerResult
}

func (s *sharedControllerFactory) getWorkers(gvr schema.GroupVersionResource, workers int) (int, error) {
	gvk, err := s.sharedCacheFactory.SharedClientFactory().GVKForResource(gvr)
	if meta.IsNoMatchError(err) {
		return workers, nil
	} else if err != nil {
		return 0, err
	}

	w, ok := s.kindWorkers[gvk]
	if ok {
		return w, nil
	}
	if workers > 0 {
		return workers, nil
	}
	return s.workers, nil
}

func (s *sharedControllerFactory) byResource(gvr schema.GroupVersionResource) *sharedController {
	s.controllerLock.RLock()
	defer s.controllerLock.RUnlock()
	return s.controllers[gvr]
}

func (s *sharedControllerFactory) SharedCacheFactory() cache.SharedCacheFactory {
	return s.sharedCacheFactory
}