k3s/vendor/github.com/rancher/wrangler/pkg/generic/controllerfactory.go

package generic

import (
	"context"
	"sync"

	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/util/workqueue"
)

type ControllerManager struct {
	lock        sync.Mutex
	generation  int
	started     map[schema.GroupVersionKind]bool
	controllers map[schema.GroupVersionKind]*Controller
	handlers    map[schema.GroupVersionKind]*Handlers
}

func (g *ControllerManager) Start(ctx context.Context, defaultThreadiness int, threadiness map[schema.GroupVersionKind]int) error {
	g.lock.Lock()
	defer g.lock.Unlock()

	for gvk, controller := range g.controllers {
		if g.started[gvk] {
			continue
		}

		threadiness, ok := threadiness[gvk]
		if !ok {
			threadiness = defaultThreadiness
		}
		if err := controller.Run(threadiness, ctx.Done()); err != nil {
			return err
		}

		if g.started == nil {
			g.started = map[schema.GroupVersionKind]bool{}
		}
		g.started[gvk] = true
	}

	return nil
}

func (g *ControllerManager) Enqueue(gvk schema.GroupVersionKind, informer cache.SharedIndexInformer, namespace, name string) {
	_, controller, _ := g.getController(gvk, informer, true)
	controller.Enqueue(namespace, name)
}

func (g *ControllerManager) removeHandler(gvk schema.GroupVersionKind, generation int) {
	g.lock.Lock()
	defer g.lock.Unlock()

	handlers, ok := g.handlers[gvk]
	if !ok {
		return
	}

	var newHandlers []handlerEntry
	for _, h := range handlers.handlers {
		if h.generation == generation {
			continue
		}
		newHandlers = append(newHandlers, h)
	}

	handlers.handlers = newHandlers
}

func (g *ControllerManager) getController(gvk schema.GroupVersionKind, informer cache.SharedIndexInformer, lock bool) (*Handlers, *Controller, bool) {
	if lock {
		g.lock.Lock()
		defer g.lock.Unlock()
	}

	if controller, ok := g.controllers[gvk]; ok {
		return g.handlers[gvk], controller, true
	}

	handlers := &Handlers{}

	queue := workqueue.NewNamedRateLimitingQueue(workqueue.DefaultControllerRateLimiter(), gvk.String())
	controller := NewController(gvk.String(), informer, queue, handlers.Handle)

	if g.handlers == nil {
		g.handlers = map[schema.GroupVersionKind]*Handlers{}
	}

	if g.controllers == nil {
		g.controllers = map[schema.GroupVersionKind]*Controller{}
	}

	g.handlers[gvk] = handlers
	g.controllers[gvk] = controller

	return handlers, controller, false
}

func (g *ControllerManager) AddHandler(ctx context.Context, gvk schema.GroupVersionKind, informer cache.SharedIndexInformer, name string, handler Handler) {
	g.lock.Lock()
	defer g.lock.Unlock()

	g.generation++
	entry := handlerEntry{
		generation: g.generation,
		name:       name,
		handler:    handler,
	}

	go func() {
		<-ctx.Done()
		g.removeHandler(gvk, entry.generation)
	}()

	handlers, controller, ok := g.getController(gvk, informer, false)
	handlers.handlers = append(handlers.handlers, entry)

	if ok {
		for _, key := range controller.informer.GetStore().ListKeys() {
			controller.workqueue.Add(key)
		}
	}
}