Add support for passing in a virtual function index to assign pci device

internal/devices/assigned_devices.go

@@ -30,7 +30,7 @@ import (
 // this function, `vpci` is released and nil is returned for that value.
 //
 // Returns a slice of strings representing the resulting location path(s) for the specified device.
-func AddDevice(ctx context.Context, vm *uvm.UtilityVM, idType, deviceID, deviceUtilPath string) (vpci *uvm.VPCIDevice, locationPaths []string, err error) {
+func AddDevice(ctx context.Context, vm *uvm.UtilityVM, idType, deviceID string, index uint16, deviceUtilPath string) (vpci *uvm.VPCIDevice, locationPaths []string, err error) {
 	defer func() {
 		if err != nil && vpci != nil {
 			// best effort clean up allocated resource on failure
@@ -41,7 +41,7 @@ func AddDevice(ctx context.Context, vm *uvm.UtilityVM, idType, deviceID, deviceU
 		}
 	}()
 	if idType == uvm.VPCIDeviceIDType || idType == uvm.VPCIDeviceIDTypeLegacy {
-		vpci, err = vm.AssignDevice(ctx, deviceID)
+		vpci, err = vm.AssignDevice(ctx, deviceID, index)
 		if err != nil {
 			return vpci, nil, errors.Wrapf(err, "failed to assign device %s of type %s to pod %s", deviceID, idType, vm.ID())
 		}

internal/hcsoci/devices.go

@@ -7,6 +7,7 @@ import (
 	"io/ioutil"
 	"os"
 	"path/filepath"
+	"strconv"
 
 	"github.com/Microsoft/hcsshim/internal/devices"
 	hcsschema "github.com/Microsoft/hcsshim/internal/hcs/schema2"
@@ -123,7 +124,8 @@ func handleAssignedDevicesWindows(ctx context.Context, vm *uvm.UtilityVM, annota
 
 	// assign device into UVM and create corresponding spec windows devices
 	for _, d := range specDevs {
-		vpciCloser, locationPaths, err := devices.AddDevice(ctx, vm, d.IDType, d.ID, deviceUtilPath)
+		pciID, index := getDeviceInfoFromPath(d.ID)
+		vpciCloser, locationPaths, err := devices.AddDevice(ctx, vm, d.IDType, pciID, index, deviceUtilPath)
 		if err != nil {
 			return nil, nil, err
 		}
@@ -140,3 +142,14 @@ func handleAssignedDevicesWindows(ctx context.Context, vm *uvm.UtilityVM, annota
 
 	return resultDevs, closers, nil
 }
+
+func getDeviceInfoFromPath(rawDevicePath string) (string, uint16) {
+	indexString := filepath.Base(rawDevicePath)
+	index, err := strconv.ParseUint(indexString, 10, 16)
+	if err == nil {
+		// we have a vf index
+		return filepath.Dir(rawDevicePath), uint16(index)
+	}
+	// otherwise, just use default index and full device ID given
+	return rawDevicePath, 0
+}

internal/hcsoci/resources_lcow.go

@@ -174,7 +174,8 @@ func allocateLinuxResources(ctx context.Context, coi *createOptionsInternal, r *
 		switch d.IDType {
 		case uvm.GPUDeviceIDType:
 			addGPUVHD = true
-			vpci, err := coi.HostingSystem.AssignDevice(ctx, d.ID)
+			pciID, index := getDeviceInfoFromPath(d.ID)
+			vpci, err := coi.HostingSystem.AssignDevice(ctx, pciID, index)
 			if err != nil {
 				return errors.Wrapf(err, "failed to assign gpu device %s to pod %s", d.ID, coi.HostingSystem.ID())
 			}

internal/uvm/create_lcow.go

@@ -173,7 +173,7 @@ func CreateLCOW(ctx context.Context, opts *OptionsLCOW) (_ *UtilityVM, err error
 		scsiControllerCount:     opts.SCSIControllerCount,
 		vpmemMaxCount:           opts.VPMemDeviceCount,
 		vpmemMaxSizeBytes:       opts.VPMemSizeBytes,
-		vpciDevices:             make(map[string]*VPCIDevice),
+		vpciDevices:             make(map[VPCIDeviceKey]*VPCIDevice),
 		physicallyBacked:        !opts.AllowOvercommit,
 		devicesPhysicallyBacked: opts.FullyPhysicallyBacked,
 		createOpts:              opts,

internal/uvm/create_wcow.go

@@ -221,7 +221,7 @@ func CreateWCOW(ctx context.Context, opts *OptionsWCOW) (_ *UtilityVM, err error
 		scsiControllerCount:     1,
 		vsmbDirShares:           make(map[string]*VSMBShare),
 		vsmbFileShares:          make(map[string]*VSMBShare),
-		vpciDevices:             make(map[string]*VPCIDevice),
+		vpciDevices:             make(map[VPCIDeviceKey]*VPCIDevice),
 		physicallyBacked:        !opts.AllowOvercommit,
 		devicesPhysicallyBacked: opts.FullyPhysicallyBacked,
 		vsmbNoDirectMap:         opts.NoDirectMap,

internal/uvm/types.go

@@ -85,7 +85,7 @@ type UtilityVM struct {
 	scsiControllerCount uint32            // Number of SCSI controllers in the utility VM
 	encryptScratch      bool              // Enable scratch encryption
 
-	vpciDevices map[string]*VPCIDevice // map of device instance id to vpci device
+	vpciDevices map[VPCIDeviceKey]*VPCIDevice // map of device instance id to vpci device
 
 	// Plan9 are directories mapped into a Linux utility VM
 	plan9Counter uint64 // Each newly-added plan9 share has a counter used as its ID in the ResourceURI and for the name

internal/uvm/virtual_device.go

@@ -24,6 +24,11 @@ const (
 const vmbusChannelTypeGUIDFormatted = "{44c4f61d-4444-4400-9d52-802e27ede19f}"
 const assignedDeviceEnumerator = "VMBUS"
 
+type VPCIDeviceKey struct {
+	deviceInstanceID     string
+	virtualFunctionIndex uint16
+}
+
 // VPCIDevice represents a vpci device. Holds its guid and a handle to the uvm it
 // belongs to.
 type VPCIDevice struct {
@@ -33,6 +38,8 @@ type VPCIDevice struct {
 	VMBusGUID string
 	// deviceInstanceID is the instance ID of the device on the host
 	deviceInstanceID string
+	// virtualFunctionIndex is the function index for the pci device to assign
+	virtualFunctionIndex uint16
 	// refCount stores the number of references to this device in the UVM
 	refCount uint32
 }
@@ -56,7 +63,7 @@ func (uvm *UtilityVM) GetAssignedDeviceVMBUSInstanceID(vmBusChannelGUID string)
 
 // Release frees the resources of the corresponding vpci device
 func (vpci *VPCIDevice) Release(ctx context.Context) error {
-	if err := vpci.vm.removeDevice(ctx, vpci.deviceInstanceID); err != nil {
+	if err := vpci.vm.removeDevice(ctx, vpci.deviceInstanceID, vpci.virtualFunctionIndex); err != nil {
 		return fmt.Errorf("failed to remove VPCI device: %s", err)
 	}
 	return nil
@@ -68,17 +75,22 @@ func (vpci *VPCIDevice) Release(ctx context.Context) error {
 // Otherwise, a new request is made to assign the target device indicated by the deviceID
 // onto the UVM. A new VPCIDevice entry is made on the UVM and the VPCIDevice is returned
 // to the caller
-func (uvm *UtilityVM) AssignDevice(ctx context.Context, deviceID string) (*VPCIDevice, error) {
+func (uvm *UtilityVM) AssignDevice(ctx context.Context, deviceID string, index uint16) (*VPCIDevice, error) {
 	guid, err := guid.NewV4()
 	if err != nil {
 		return nil, err
 	}
 	vmBusGUID := guid.String()
 
+	key := VPCIDeviceKey{
+		deviceInstanceID:     deviceID,
+		virtualFunctionIndex: index,
+	}
+
 	uvm.m.Lock()
 	defer uvm.m.Unlock()
 
-	existingVPCIDevice := uvm.vpciDevices[deviceID]
+	existingVPCIDevice := uvm.vpciDevices[key]
 	if existingVPCIDevice != nil {
 		existingVPCIDevice.refCount++
 		return existingVPCIDevice, nil
@@ -88,6 +100,7 @@ func (uvm *UtilityVM) AssignDevice(ctx context.Context, deviceID string) (*VPCID
 		Functions: []hcsschema.VirtualPciFunction{
 			{
 				DeviceInstancePath: deviceID,
+				VirtualFunction:    index,
 			},
 		},
 	}
@@ -117,30 +130,36 @@ func (uvm *UtilityVM) AssignDevice(ctx context.Context, deviceID string) (*VPCID
 		return nil, err
 	}
 	result := &VPCIDevice{
-		vm:               uvm,
-		VMBusGUID:        vmBusGUID,
-		deviceInstanceID: deviceID,
-		refCount:         1,
+		vm:                   uvm,
+		VMBusGUID:            vmBusGUID,
+		deviceInstanceID:     deviceID,
+		virtualFunctionIndex: index,
+		refCount:             1,
 	}
-	uvm.vpciDevices[deviceID] = result
+	uvm.vpciDevices[key] = result
 	return result, nil
 }
 
 // removeDevice removes a vpci device from a uvm when there are
 // no more references to a given VPCIDevice. Otherwise, decrements
 // the reference count of the stored VPCIDevice and returns nil.
-func (uvm *UtilityVM) removeDevice(ctx context.Context, deviceInstanceID string) error {
+func (uvm *UtilityVM) removeDevice(ctx context.Context, deviceInstanceID string, index uint16) error {
+	key := VPCIDeviceKey{
+		deviceInstanceID:     deviceInstanceID,
+		virtualFunctionIndex: index,
+	}
+
 	uvm.m.Lock()
 	defer uvm.m.Unlock()
 
-	vpci := uvm.vpciDevices[deviceInstanceID]
+	vpci := uvm.vpciDevices[key]
 	if vpci == nil {
-		return fmt.Errorf("no device with ID %s is present on the uvm %s", deviceInstanceID, uvm.ID())
+		return fmt.Errorf("no device with ID %s and index %d is present on the uvm %s", deviceInstanceID, index, uvm.ID())
 	}
 
 	vpci.refCount--
 	if vpci.refCount == 0 {
-		delete(uvm.vpciDevices, deviceInstanceID)
+		delete(uvm.vpciDevices, key)
 		return uvm.modify(ctx, &hcsschema.ModifySettingRequest{
 			ResourcePath: fmt.Sprintf(resourcepaths.VirtualPCIResourceFormat, vpci.VMBusGUID),
 			RequestType:  requesttype.Remove,

test/functional/uvm_virtualdevice_test.go

@@ -55,7 +55,7 @@ func TestVirtualDevice(t *testing.T) {
 	// create test uvm and ensure we can assign and remove the device
 	vm := testutilities.CreateLCOWUVMFromOpts(ctx, t, opts)
 	defer vm.Close()
-	vpci, err := vm.AssignDevice(ctx, testDeviceInstanceID)
+	vpci, err := vm.AssignDevice(ctx, testDeviceInstanceID, 0)
 	if err != nil {
 		t.Fatalf("failed to assign device %s with %v", testDeviceInstanceID, err)
 	}