refactor: simplify usb device plugin logic to 1-1 mapping

Signed-off-by: Jack Yu <jack.yu@suse.com>
harvester · Jul 5, 2024 · c41c51c · c41c51c
1 parent bd8f370
commit c41c51c
Show file tree

Hide file tree

Showing 6 changed files with 145 additions and 374 deletions.
diff --git a/pkg/controller/usbdevice/register.go b/pkg/controller/usbdevice/register.go
@@ -6,7 +6,6 @@ import (
 	"github.com/rancher/wrangler/pkg/relatedresource"
 
 	"github.com/harvester/pcidevices/pkg/config"
-	"github.com/harvester/pcidevices/pkg/deviceplugins"
 )
 
 const (
@@ -21,7 +20,7 @@ func Register(ctx context.Context, management *config.FactoryManager) error {
 	virtClient := management.KubevirtFactory.Kubevirt().V1().KubeVirt()
 
 	handler := NewHandler(usbDeviceCtrl, usbDeviceClaimCtrl, usbDeviceCtrl.Cache(), usbDeviceClaimCtrl.Cache())
-	usbDeviceClaimController := NewClaimHandler(usbDeviceCtrl.Cache(), usbDeviceClaimCtrl, usbDeviceCtrl, virtClient, deviceplugins.NewUSBDevicePlugin)
+	usbDeviceClaimController := NewClaimHandler(usbDeviceCtrl.Cache(), usbDeviceClaimCtrl, usbDeviceCtrl, virtClient)
 
 	usbDeviceClaimCtrl.OnChange(ctx, "usbClaimClient-device-claim", usbDeviceClaimController.OnUSBDeviceClaimChanged)
 	usbDeviceClaimCtrl.OnRemove(ctx, "usbClaimClient-device-claim-remove", usbDeviceClaimController.OnRemove)

diff --git a/pkg/controller/usbdevice/usbdevice_claim_controller.go b/pkg/controller/usbdevice/usbdevice_claim_controller.go
@@ -17,43 +17,34 @@ import (
 	ctlkubevirtv1 "github.com/harvester/pcidevices/pkg/generated/controllers/kubevirt.io/v1"
 )
 
-var (
-	discoverAllowedUSBDevices = deviceplugins.DiscoverAllowedUSBDevices
-)
-
 type DevClaimHandler struct {
-	usbClaimClient        ctldevicerv1beta1.USBDeviceClaimClient
-	usbClient             ctldevicerv1beta1.USBDeviceClient
-	virtClient            ctlkubevirtv1.KubeVirtClient
-	lock                  *sync.Mutex
-	usbDeviceCache        ctldevicerv1beta1.USBDeviceCache
-	devicePlugin          map[string]*deviceController
-	devicePluginConvertor devicePluginConvertor
+	usbClaimClient ctldevicerv1beta1.USBDeviceClaimClient
+	usbClient      ctldevicerv1beta1.USBDeviceClient
+	virtClient     ctlkubevirtv1.KubeVirtClient
+	lock           *sync.Mutex
+	usbDeviceCache ctldevicerv1beta1.USBDeviceCache
+	devicePlugin   map[string]*deviceController
 }
 
 type deviceController struct {
-	device  deviceplugins.USBDevicePluginInterface
+	device  *deviceplugins.USBDevicePlugin
 	stop    chan struct{}
 	started bool
 }
 
-type devicePluginConvertor func(resourceName string, devices []*deviceplugins.PluginDevices) deviceplugins.USBDevicePluginInterface
-
 func NewClaimHandler(
 	usbDeviceCache ctldevicerv1beta1.USBDeviceCache,
 	usbClaimClient ctldevicerv1beta1.USBDeviceClaimClient,
 	usbClient ctldevicerv1beta1.USBDeviceClient,
 	virtClient ctlkubevirtv1.KubeVirtClient,
-	devicePluginHelper devicePluginConvertor,
 ) *DevClaimHandler {
 	return &DevClaimHandler{
-		usbDeviceCache:        usbDeviceCache,
-		usbClaimClient:        usbClaimClient,
-		usbClient:             usbClient,
-		virtClient:            virtClient,
-		lock:                  &sync.Mutex{},
-		devicePlugin:          map[string]*deviceController{},
-		devicePluginConvertor: devicePluginHelper,
+		usbDeviceCache: usbDeviceCache,
+		usbClaimClient: usbClaimClient,
+		usbClient:      usbClient,
+		virtClient:     virtClient,
+		lock:           &sync.Mutex{},
+		devicePlugin:   map[string]*deviceController{},
 	}
 }
 
@@ -97,21 +88,25 @@ func (h *DevClaimHandler) OnUSBDeviceClaimChanged(_ string, usbDeviceClaim *v1be
 		return usbDeviceClaim, err
 	}
 
-	// start device plugin if it's not started yet.
-	if _, ok := h.devicePlugin[usbDeviceClaim.Name]; !ok {
-		pluginDevices := discoverAllowedUSBDevices(convertToKubeVirtUSBFormat(usbDevice))
+	deviceHandler, ok := h.devicePlugin[usbDeviceClaim.Name]
 
-		if pluginDevice := h.findDevicePlugin(pluginDevices, usbDevice); pluginDevice != nil {
-			usbDevicePlugin := h.devicePluginConvertor(usbDevice.Status.ResourceName, []*deviceplugins.PluginDevices{pluginDevice})
-			deviceHan := &deviceController{
-				device: usbDevicePlugin,
-			}
-			h.devicePlugin[usbDeviceClaim.Name] = deviceHan
-			h.startDevicePlugin(deviceHan, usbDeviceClaim.Name)
-		} else {
-			logrus.Errorf("failed to find device plugin for usb device %s", usbDevice.Name)
+	if !ok {
+		usbDevicePlugin, err := deviceplugins.NewUSBDevicePlugin(*usbDevice)
+
+		if err != nil {
+			logrus.Errorf("failed to create usb device plugin: %v", err)
 			return usbDeviceClaim, err
 		}
+
+		deviceHan := &deviceController{
+			device: usbDevicePlugin,
+		}
+		h.devicePlugin[usbDeviceClaim.Name] = deviceHan
+		deviceHandler = deviceHan
+	}
+
+	if !deviceHandler.started {
+		h.startDevicePlugin(deviceHandler, usbDeviceClaim.Name)
 	}
 
 	if !usbDevice.Status.Enabled {
@@ -167,25 +162,6 @@ func (h *DevClaimHandler) stopDevicePlugin(deviceHan *deviceController) {
 	deviceHan.started = false
 }
 
-func (h *DevClaimHandler) findDevicePlugin(pluginDevices map[string][]*deviceplugins.PluginDevices, usbDevice *v1beta1.USBDevice) *deviceplugins.PluginDevices {
-	var pluginDevice *deviceplugins.PluginDevices
-
-	for resourceName, devices := range pluginDevices {
-		for _, device := range devices {
-			device := device
-			for _, d := range device.Devices {
-				logrus.Debugf("resourceName: %s, device: %v", resourceName, d)
-				if usbDevice.Status.DevicePath == d.DevicePath {
-					pluginDevice = device
-					return pluginDevice
-				}
-			}
-		}
-	}
-
-	return pluginDevice
-}
-
 func (h *DevClaimHandler) OnRemove(_ string, claim *v1beta1.USBDeviceClaim) (*v1beta1.USBDeviceClaim, error) {
 	if claim == nil || claim.DeletionTimestamp == nil {
 		return claim, nil
@@ -296,18 +272,3 @@ func (h *DevClaimHandler) updateKubeVirt(virt *kubevirtv1.KubeVirt, usbDevice *v
 
 	return virt, nil
 }
-
-func convertToKubeVirtUSBFormat(ub *v1beta1.USBDevice) []kubevirtv1.USBHostDevice {
-	return []kubevirtv1.USBHostDevice{
-		{
-			Selectors: []kubevirtv1.USBSelector{
-				{
-					Vendor:  ub.Status.VendorID,
-					Product: ub.Status.ProductID,
-				},
-			},
-			ResourceName:             ub.Status.ResourceName,
-			ExternalResourceProvider: true,
-		},
-	}
-}
diff --git a/pkg/controller/usbdevice/usbdevice_claim_controller_test.go b/pkg/controller/usbdevice/usbdevice_claim_controller_test.go
@@ -10,20 +10,10 @@ import (
 	kubevirtv1 "kubevirt.io/api/core/v1"
 
 	"github.com/harvester/pcidevices/pkg/apis/devices.harvesterhci.io/v1beta1"
-	"github.com/harvester/pcidevices/pkg/deviceplugins"
 	"github.com/harvester/pcidevices/pkg/generated/clientset/versioned/fake"
 	"github.com/harvester/pcidevices/pkg/util/fakeclients"
 )
 
-type mockUSBDevicePlugin struct {
-	startTimes int // used to test how many startTimes the Start is called
-}
-
-func (m *mockUSBDevicePlugin) Start(_ <-chan struct{}) error {
-	m.startTimes++
-	return nil
-}
-
 var (
 	mockUsbDevice1 = &v1beta1.USBDevice{
 		ObjectMeta: metav1.ObjectMeta{
@@ -66,9 +56,6 @@ var (
 			NodeName: "test-node",
 		},
 	}
-	mockUSBDevicePluginHelper = func(_ string, _ []*deviceplugins.PluginDevices) deviceplugins.USBDevicePluginInterface {
-		return &mockUSBDevicePlugin{}
-	}
 )
 
 func Test_OnUSBDeviceClaimChanged(t *testing.T) {
@@ -79,23 +66,18 @@ func Test_OnUSBDeviceClaimChanged(t *testing.T) {
 	}{
 		{
 			fun: func(t *testing.T) {
-				client := generateClient()
-				mockObj := &mockUSBDevicePlugin{}
+				client := fake.NewSimpleClientset(mockUsbDevice1, mockUsbDeviceClaim1, mockKubeVirt)
 				handler := NewClaimHandler(
 					fakeclients.USBDeviceCache(client.DevicesV1beta1().USBDevices),
 					fakeclients.USBDeviceClaimsClient(client.DevicesV1beta1().USBDeviceClaims),
 					fakeclients.USBDevicesClient(client.DevicesV1beta1().USBDevices),
 					fakeclients.KubeVirtClient(client.KubevirtV1().KubeVirts),
-					func(_ string, _ []*deviceplugins.PluginDevices) deviceplugins.USBDevicePluginInterface {
-						return mockObj
-					},
 				)
 
 				// Test claim created
 				_, err := handler.OnUSBDeviceClaimChanged("", mockUsbDeviceClaim1)
 				assert.NoError(t, err)
 				time.Sleep(1 * time.Second)
-				assert.Equal(t, 1, mockObj.startTimes)
 
 				kubevirt, err := client.KubevirtV1().KubeVirts(mockKubeVirt.Namespace).Get(context.Background(), mockKubeVirt.Name, metav1.GetOptions{})
 				assert.NoError(t, err)
@@ -140,37 +122,30 @@ func Test_OnUSBDeviceClaimChanged(t *testing.T) {
 		},
 		{
 			fun: func(_ *testing.T) {
-				client := generateClient()
-				mockObj := &mockUSBDevicePlugin{startTimes: 0}
+				client := fake.NewSimpleClientset(mockUsbDevice1, mockUsbDeviceClaim1, mockKubeVirt)
 				handler := NewClaimHandler(
 					fakeclients.USBDeviceCache(client.DevicesV1beta1().USBDevices),
 					fakeclients.USBDeviceClaimsClient(client.DevicesV1beta1().USBDeviceClaims),
 					fakeclients.USBDevicesClient(client.DevicesV1beta1().USBDevices),
 					fakeclients.KubeVirtClient(client.KubevirtV1().KubeVirts),
-					func(_ string, _ []*deviceplugins.PluginDevices) deviceplugins.USBDevicePluginInterface {
-						return mockObj
-					},
 				)
 
 				// Test claim created
 				_, err := handler.OnUSBDeviceClaimChanged("", mockUsbDeviceClaim1)
 				assert.NoError(t, err)
 				_, err = handler.OnUSBDeviceClaimChanged("", mockUsbDeviceClaim1)
 				assert.NoError(t, err)
-				time.Sleep(1 * time.Second)
-				assert.Equal(t, 1, mockObj.startTimes)
 			},
 			description: "Case to create two identical claims",
 		},
 		{
 			fun: func(_ *testing.T) {
-				client := generateClient()
+				client := fake.NewSimpleClientset(mockUsbDevice1, mockUsbDeviceClaim1, mockKubeVirt)
 				handler := NewClaimHandler(
 					fakeclients.USBDeviceCache(client.DevicesV1beta1().USBDevices),
 					fakeclients.USBDeviceClaimsClient(client.DevicesV1beta1().USBDeviceClaims),
 					fakeclients.USBDevicesClient(client.DevicesV1beta1().USBDevices),
 					fakeclients.KubeVirtClient(client.KubevirtV1().KubeVirts),
-					mockUSBDevicePluginHelper,
 				)
 
 				// Test claim created
@@ -188,25 +163,3 @@ func Test_OnUSBDeviceClaimChanged(t *testing.T) {
 		})
 	}
 }
-
-func generateClient() *fake.Clientset {
-	client := fake.NewSimpleClientset(mockUsbDevice1, mockUsbDeviceClaim1, mockKubeVirt)
-	discoverAllowedUSBDevices = func(_ []kubevirtv1.USBHostDevice) map[string][]*deviceplugins.PluginDevices {
-		m := map[string][]*deviceplugins.PluginDevices{}
-		m[mockUsbDevice1.Status.ResourceName] = []*deviceplugins.PluginDevices{
-			{
-				ID: "test",
-				Devices: []*deviceplugins.USBDevice{
-					{
-						Vendor:     2385,
-						Product:    5734,
-						DevicePath: "/dev/bus/usb/001/002",
-					},
-				},
-			},
-		}
-		return m
-	}
-
-	return client
-}
diff --git a/pkg/deviceplugins/local_usb_device.go b/pkg/deviceplugins/local_usb_device.go