refactor: setup and register logic

Signed-off-by: Jack Yu <jack.yu@suse.com>

pkg/config/factory_magement.go

@@ -0,0 +1,29 @@
+package config
+
+import (
+	ctlnetwork "github.com/harvester/harvester-network-controller/pkg/generated/controllers/network.harvesterhci.io"
+	ctlcore "github.com/rancher/wrangler/pkg/generated/controllers/core"
+	"k8s.io/client-go/rest"
+	"kubevirt.io/client-go/kubecli"
+
+	ctldevices "github.com/harvester/pcidevices/pkg/generated/controllers/devices.harvesterhci.io"
+)
+
+type FactoryManager struct {
+	DeviceFactory  *ctldevices.Factory
+	CoreFactory    *ctlcore.Factory
+	NetworkFactory *ctlnetwork.Factory
+
+	KubevirtClient kubecli.KubevirtClient
+	Cfg            *rest.Config
+}
+
+func NewFactoryManager(deviceFactory *ctldevices.Factory, coreFactory *ctlcore.Factory, networkFactory *ctlnetwork.Factory, kubevirtClient kubecli.KubevirtClient, cfg *rest.Config) *FactoryManager {
+	return &FactoryManager{
+		DeviceFactory:  deviceFactory,
+		CoreFactory:    coreFactory,
+		NetworkFactory: networkFactory,
+		KubevirtClient: kubevirtClient,
+		Cfg:            cfg,
+	}
+}

pkg/controller/gpudevice/gpu_controller.go

@@ -8,7 +8,6 @@ import (
 	"reflect"
 	"strings"
 
-	ctlcorev1 "github.com/rancher/wrangler/pkg/generated/controllers/core/v1"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/pflag"
 	"gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvpci"
@@ -20,6 +19,7 @@ import (
 	"kubevirt.io/client-go/kubecli"
 
 	"github.com/harvester/pcidevices/pkg/apis/devices.harvesterhci.io/v1beta1"
+	"github.com/harvester/pcidevices/pkg/config"
 	"github.com/harvester/pcidevices/pkg/deviceplugins"
 	ctl "github.com/harvester/pcidevices/pkg/generated/controllers/devices.harvesterhci.io/v1beta1"
 	"github.com/harvester/pcidevices/pkg/util/executor"
@@ -72,13 +72,18 @@ func NewHandler(ctx context.Context, sriovGPUController ctl.SRIOVGPUDeviceContro
 }
 
 // Register setups up handlers for SRIOVGPUDevices and VGPUDevices
-func Register(ctx context.Context, sriovGPUController ctl.SRIOVGPUDeviceController, vGPUController ctl.VGPUDeviceController, pciDeviceClaimController ctl.PCIDeviceClaimController, podController ctlcorev1.PodController, cfg *rest.Config) error {
+func Register(ctx context.Context, management *config.FactoryManager) error {
+	sriovGPUController := management.DeviceFactory.Devices().V1beta1().SRIOVGPUDevice()
+	vGPUController := management.DeviceFactory.Devices().V1beta1().VGPUDevice()
+	pciDeviceClaimController := management.DeviceFactory.Devices().V1beta1().PCIDeviceClaim()
+	podController := management.CoreFactory.Core().V1().Pod()
+
 	clientConfig := kubecli.DefaultClientConfig(&pflag.FlagSet{})
 	virtClient, err := kubecli.GetKubevirtClientFromClientConfig(clientConfig)
 	if err != nil {
 		return err
 	}
-	h, err := NewHandler(ctx, sriovGPUController, vGPUController, pciDeviceClaimController, virtClient, nil, cfg)
+	h, err := NewHandler(ctx, sriovGPUController, vGPUController, pciDeviceClaimController, virtClient, nil, management.Cfg)
 	if err != nil {
 		return err
 	}

pkg/controller/nodecleanup/nodecleanup_controller.go

@@ -9,6 +9,7 @@ import (
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
+	"github.com/harvester/pcidevices/pkg/config"
 	"github.com/harvester/pcidevices/pkg/generated/controllers/devices.harvesterhci.io/v1beta1"
 )
 
@@ -70,11 +71,11 @@ func (h *Handler) OnRemove(_ string, node *v1.Node) (*v1.Node, error) {
 	return node, nil
 }
 
-func Register(
-	ctx context.Context,
-	pdcClient v1beta1.PCIDeviceClaimController,
-	pdClient v1beta1.PCIDeviceController,
-	nodeClient corecontrollers.NodeController) error {
+func Register(ctx context.Context, management *config.FactoryManager) error {
+	pdcClient := management.DeviceFactory.Devices().V1beta1().PCIDeviceClaim()
+	pdClient := management.DeviceFactory.Devices().V1beta1().PCIDevice()
+	nodeClient := management.CoreFactory.Core().V1().Node()
+
 	handler := &Handler{
 		pdcClient:  pdcClient,
 		pdClient:   pdClient,

pkg/controller/nodes/node_controller.go

@@ -14,6 +14,7 @@ import (
 	"k8s.io/apimachinery/pkg/labels"
 	"kubevirt.io/client-go/kubecli"
 
+	"github.com/harvester/pcidevices/pkg/config"
 	"github.com/harvester/pcidevices/pkg/controller/gpudevice"
 	"github.com/harvester/pcidevices/pkg/controller/usbdevice"
 
@@ -54,8 +55,20 @@ const (
 	reconcilePCIDevices = "reconcile-pcidevices"
 )
 
-func Register(ctx context.Context, sriovCtl ctl.SRIOVNetworkDeviceController, pciDeviceCtl ctl.PCIDeviceController, nodeCtl ctl.NodeController, coreNodeCtl ctlcorev1.NodeController, vlanConfigCache ctlnetworkv1beta1.VlanConfigCache, sriovNetworkDeviceCache ctl.SRIOVNetworkDeviceCache, pciDeviceClaimController ctl.PCIDeviceClaimController, vGPUController ctl.VGPUDeviceController, sriovGPUController ctl.SRIOVGPUDeviceController, usbCtl ctl.USBDeviceController, usbClaimCtl ctl.USBDeviceClaimController, virtClient kubecli.KubevirtClient) error {
+func Register(ctx context.Context, management *config.FactoryManager) error {
+	sriovCtl := management.DeviceFactory.Devices().V1beta1().SRIOVNetworkDevice()
+	pciDeviceCtl := management.DeviceFactory.Devices().V1beta1().PCIDevice()
+	coreNodeCtl := management.CoreFactory.Core().V1().Node()
+	vlanConfigCache := management.NetworkFactory.Network().V1beta1().VlanConfig().Cache()
+	nodeCtl := management.DeviceFactory.Devices().V1beta1().Node()
+	vGPUController := management.DeviceFactory.Devices().V1beta1().VGPUDevice()
+	pciDeviceClaimController := management.DeviceFactory.Devices().V1beta1().PCIDeviceClaim()
+	sriovGPUController := management.DeviceFactory.Devices().V1beta1().SRIOVGPUDevice()
+	usbCtl := management.DeviceFactory.Devices().V1beta1().USBDevice()
+	usbClaimCtl := management.DeviceFactory.Devices().V1beta1().USBDeviceClaim()
+	virtClient := management.KubevirtClient
 	nodeName := os.Getenv(v1beta1.NodeEnvVarName)
+
 	h := &handler{
 		ctx:                      ctx,
 		sriovCache:               sriovCtl.Cache(),
@@ -67,7 +80,7 @@ func Register(ctx context.Context, sriovCtl ctl.SRIOVNetworkDeviceController, pc
 		coreNodeCtl:              coreNodeCtl,
 		vlanConfigCache:          vlanConfigCache,
 		nodeCtl:                  nodeCtl,
-		sriovNetworkDeviceCache:  sriovNetworkDeviceCache,
+		sriovNetworkDeviceCache:  sriovCtl.Cache(),
 		vGPUController:           vGPUController,
 		pciDeviceClaimController: pciDeviceClaimController,
 		sriovGPUController:       sriovGPUController,

pkg/controller/pcideviceclaim/pcideviceclaim_controller.go

@@ -21,6 +21,7 @@ import (
 	"kubevirt.io/client-go/kubecli"
 
 	"github.com/harvester/pcidevices/pkg/apis/devices.harvesterhci.io/v1beta1"
+	"github.com/harvester/pcidevices/pkg/config"
 	"github.com/harvester/pcidevices/pkg/deviceplugins"
 	v1beta1gen "github.com/harvester/pcidevices/pkg/generated/controllers/devices.harvesterhci.io/v1beta1"
 )
@@ -59,11 +60,10 @@ type Handler struct {
 	devicePlugins map[string]*deviceplugins.PCIDevicePlugin
 }
 
-func Register(
-	ctx context.Context,
-	pdcClient v1beta1gen.PCIDeviceClaimController,
-	pdClient v1beta1gen.PCIDeviceController,
-) error {
+func Register(ctx context.Context, management *config.FactoryManager) error {
+	pdcClient := management.DeviceFactory.Devices().V1beta1().PCIDeviceClaim()
+	pdClient := management.DeviceFactory.Devices().V1beta1().PCIDevice()
+
 	logrus.Info("Registering PCI Device Claims controller")
 	nodeName := os.Getenv(v1beta1.NodeEnvVarName)
 	clientConfig := kubecli.DefaultClientConfig(&pflag.FlagSet{})

pkg/controller/setup.go

@@ -18,14 +18,15 @@ import (
 	"k8s.io/client-go/util/workqueue"
 	"kubevirt.io/client-go/kubecli"
 
+	"github.com/harvester/pcidevices/pkg/config"
 	"github.com/harvester/pcidevices/pkg/controller/gpudevice"
 	"github.com/harvester/pcidevices/pkg/controller/nodecleanup"
 	"github.com/harvester/pcidevices/pkg/controller/nodes"
 	"github.com/harvester/pcidevices/pkg/controller/pcideviceclaim"
 	"github.com/harvester/pcidevices/pkg/controller/sriovdevice"
 	"github.com/harvester/pcidevices/pkg/controller/usbdevice"
 	"github.com/harvester/pcidevices/pkg/crd"
-	ctl "github.com/harvester/pcidevices/pkg/generated/controllers/devices.harvesterhci.io"
+	ctldevices "github.com/harvester/pcidevices/pkg/generated/controllers/devices.harvesterhci.io"
 	"github.com/harvester/pcidevices/pkg/webhook"
 )
 
@@ -49,11 +50,8 @@ func Setup(ctx context.Context, cfg *rest.Config, _ *runtime.Scheme) error {
 		DefaultRateLimiter: rateLimit,
 		DefaultWorkers:     2,
 	})
-	if err != nil {
-		return err
-	}
 
-	deviceFactory, err := ctl.NewFactoryFromConfigWithOptions(cfg, &generic.FactoryOptions{
+	deviceFactory, err := ctldevices.NewFactoryFromConfigWithOptions(cfg, &generic.FactoryOptions{
 		SharedControllerFactory: factory,
 	})
 
@@ -77,46 +75,37 @@ func Setup(ctx context.Context, cfg *rest.Config, _ *runtime.Scheme) error {
 		return fmt.Errorf("error building network controllers: %v", err)
 	}
 
-	pdCtl := deviceFactory.Devices().V1beta1().PCIDevice()
-	pdcCtl := deviceFactory.Devices().V1beta1().PCIDeviceClaim()
-	usbDeviceCtrl := deviceFactory.Devices().V1beta1().USBDevice()
-	usbDeviceClaimCtrl := deviceFactory.Devices().V1beta1().USBDeviceClaim()
-	sriovCtl := deviceFactory.Devices().V1beta1().SRIOVNetworkDevice()
-	nodeCtl := deviceFactory.Devices().V1beta1().Node()
-	coreNodeCtl := coreFactory.Core().V1().Node()
-	vlanCtl := networkFactory.Network().V1beta1().VlanConfig()
-	sriovNetworkDeviceCache := sriovCtl.Cache()
-	sriovGPUCtl := deviceFactory.Devices().V1beta1().SRIOVGPUDevice()
-	vGPUCtl := deviceFactory.Devices().V1beta1().VGPUDevice()
-	podCtl := coreFactory.Core().V1().Pod()
 	clientConfig := kubecli.DefaultClientConfig(&pflag.FlagSet{})
 	virtClient, err := kubecli.GetKubevirtClientFromClientConfig(clientConfig)
-	RegisterIndexers(sriovNetworkDeviceCache)
 
-	if err := pcideviceclaim.Register(ctx, pdcCtl, pdCtl); err != nil {
-		return fmt.Errorf("error registering pcidevicclaim controllers :%v", err)
-	}
+	management := config.NewFactoryManager(
+		deviceFactory,
+		coreFactory,
+		networkFactory,
+		virtClient,
+		cfg,
+	)
 
-	if err := usbdevice.Register(ctx, usbDeviceCtrl, usbDeviceClaimCtrl); err != nil {
-		return fmt.Errorf("error registering usbdevice controllers :%v", err)
-	}
+	nodeCtl := deviceFactory.Devices().V1beta1().Node()
+	sriovNetworkDeviceCache := deviceFactory.Devices().V1beta1().SRIOVNetworkDevice().Cache()
 
-	if err := nodes.Register(ctx, sriovCtl, pdCtl, nodeCtl, coreNodeCtl, vlanCtl.Cache(),
-		sriovNetworkDeviceCache, pdcCtl, vGPUCtl, sriovGPUCtl, usbDeviceCtrl, usbDeviceClaimCtrl, virtClient); err != nil {
-		return fmt.Errorf("error registering node controller: %v", err)
-	}
+	RegisterIndexers(sriovNetworkDeviceCache)
 
-	if err := sriovdevice.Register(ctx, sriovCtl, coreNodeCtl.Cache(), vlanCtl.Cache()); err != nil {
-		return fmt.Errorf("error registering sriovdevice controller: %v", err)
+	registers := []func(context.Context, *config.FactoryManager) error{
+		pcideviceclaim.Register,
+		usbdevice.Register,
+		nodes.Register,
+		sriovdevice.Register,
+		nodecleanup.Register,
+		gpudevice.Register,
 	}
 
-	if err := nodecleanup.Register(ctx, pdcCtl, pdCtl, coreNodeCtl); err != nil {
-		return fmt.Errorf("error registering nodecleanup controller: %v", err)
+	for _, register := range registers {
+		if err := register(ctx, management); err != nil {
+			return fmt.Errorf("error registering controller: %v", err)
+		}
 	}
 
-	if err := gpudevice.Register(ctx, sriovGPUCtl, vGPUCtl, pdcCtl, podCtl, cfg); err != nil {
-		return fmt.Errorf("error registering gpudevice controller :%v", err)
-	}
 	if err := start.All(ctx, 2, coreFactory, networkFactory, deviceFactory); err != nil {
 		return fmt.Errorf("error starting controllers :%v", err)
 	}

pkg/controller/sriovdevice/sriovdevice_controller.go

@@ -13,6 +13,7 @@ import (
 	"k8s.io/apimachinery/pkg/labels"
 
 	"github.com/harvester/pcidevices/pkg/apis/devices.harvesterhci.io/v1beta1"
+	"github.com/harvester/pcidevices/pkg/config"
 	ctl "github.com/harvester/pcidevices/pkg/generated/controllers/devices.harvesterhci.io/v1beta1"
 	"github.com/harvester/pcidevices/pkg/util/nichelper"
 )
@@ -42,10 +43,12 @@ func NewHandler(ctx context.Context, sriovCache ctl.SRIOVNetworkDeviceCache, sri
 	}
 }
 
-func Register(ctx context.Context, sriovDeviceController ctl.SRIOVNetworkDeviceController,
-	nodeCache ctlcorev1.NodeCache, vlanConfigCache ctlnetworkv1beta1.VlanConfigCache) error {
-
+func Register(ctx context.Context, management *config.FactoryManager) error {
+	sriovDeviceController := management.DeviceFactory.Devices().V1beta1().SRIOVNetworkDevice()
+	nodeCache := management.CoreFactory.Core().V1().Node().Cache()
+	vlanConfigCache := management.NetworkFactory.Network().V1beta1().VlanConfig().Cache()
 	nodeName := os.Getenv(v1beta1.NodeEnvVarName)
+
 	h := NewHandler(ctx, sriovDeviceController.Cache(), sriovDeviceController, nodeName,
 		nodeCache, vlanConfigCache)
 	sriovDeviceController.OnChange(ctx, reconcileSriovDevice, h.reconcileSriovDevice)

pkg/controller/usbdevice/register.go

@@ -4,11 +4,8 @@ import (
 	"context"
 
 	"github.com/rancher/wrangler/pkg/relatedresource"
-	"github.com/sirupsen/logrus"
-	"github.com/spf13/pflag"
-	"kubevirt.io/client-go/kubecli"
 
-	v1beta1gen "github.com/harvester/pcidevices/pkg/generated/controllers/devices.harvesterhci.io/v1beta1"
+	"github.com/harvester/pcidevices/pkg/config"
 )
 
 const (
@@ -17,13 +14,10 @@ const (
 	KubeVirtResourcePrefix = "kubevirt.io/"
 )
 
-func Register(ctx context.Context, usbDeviceCtrl v1beta1gen.USBDeviceController, usbDeviceClaimCtrl v1beta1gen.USBDeviceClaimController) error {
-	clientConfig := kubecli.DefaultClientConfig(&pflag.FlagSet{})
-	virtClient, err := kubecli.GetKubevirtClientFromClientConfig(clientConfig)
-	if err != nil {
-		logrus.Errorf("failed to get kubevirt client: %v", err)
-		return err
-	}
+func Register(ctx context.Context, management *config.FactoryManager) error {
+	usbDeviceCtrl := management.DeviceFactory.Devices().V1beta1().USBDevice()
+	usbDeviceClaimCtrl := management.DeviceFactory.Devices().V1beta1().USBDeviceClaim()
+	virtClient := management.KubevirtClient
 
 	setupCommonLabels()