add deviceMapCache

Signed-off-by: Sunyanan Choochotkaew <[email protected]>
foundation-model-stack · Sep 25, 2023 · f2f957e · f2f957e
1 parent 2ad16af
commit f2f957e
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 5 deletions.
diff --git a/daemon/src/iface/iface.go b/daemon/src/iface/iface.go
@@ -162,3 +162,28 @@ func GetInterfaces() []InterfaceInfoType {
 	}
 	return interfaces
 }
+
+func getNetAddressFromDevice(devName string) (string, error) {
+	devLink, err := netlink.LinkByName(devName)
+	if err != nil {
+		log.Printf("cannot find link %s: %v", devName, err)
+		return "", err
+	}
+	addrs, err := netlink.AddrList(devLink, netlink.FAMILY_V4)
+	if err != nil || len(addrs) == 0 {
+		log.Printf("cannot list address on %s: %v", devName, err)
+		return "", err
+	}
+	addr := addrs[0].IPNet
+	if addr == nil {
+		log.Printf("no address set on %s", devName)
+		return "", err
+	}
+	if devLink.Attrs().Flags&net.FlagUp == 0 {
+		// interface down
+		log.Printf("%s down", devName)
+		return "", err
+	}
+	netAddress := getNetAddress(addr)
+	return netAddress, nil
+}
diff --git a/daemon/src/iface/pci.go b/daemon/src/iface/pci.go
@@ -26,6 +26,8 @@ const (
 
 var CheckPointfile string = "/var/lib/kubelet/device-plugins/kubelet_internal_checkpoint"
 
+var deviceMapCache = InitSafeCache()
+
 // modify from https://github.com/k8snetworkplumbingwg/multus-cni/blob/9b45d4b211728aa0db44a1624aac8e61843390cf/pkg/checkpoint/checkpoint.go#L72
 // DeviceIDs can map[string]string or []string
 type PodDevicesEntry struct {
@@ -121,10 +123,32 @@ func GetDeviceMap(resourceMap map[string][]string, resourceName string) map[stri
 
 	if deviceIDs, exist := resourceMap[resourceName]; exist {
 		for _, deviceID := range deviceIDs {
-			masterName, err := GetPfName(deviceID)
-			if err == nil {
-				if netAddress, exist := nameNetMap[masterName]; exist {
+			var masterName string
+			deviceNameInterface := deviceMapCache.GetCache(deviceID)
+			if deviceNameInterface != nil {
+				masterName = deviceNameInterface.(string)
+			} else {
+				var err error
+				masterName, err = GetPfName(deviceID)
+				if err != nil {
+					log.Printf("cannot get physical device %s: %v\n", deviceID, err)
+				} else {
+					log.Printf("set deviceMapCache %s=%s\n", deviceID, masterName)
+					deviceMapCache.SetCache(deviceID, masterName)
+				}
+			}
+			if netAddress, exist := nameNetMap[masterName]; exist {
+				deviceMap[netAddress] = deviceID
+			} else {
+				netAddress, err := getNetAddressFromDevice(masterName)
+				if err != nil {
+					log.Printf("cannot get network address of device %s: %v\n", masterName, err)
+				} else {
 					deviceMap[netAddress] = deviceID
+					// found new device, update interfaces
+					GetInterfaces()
+					nameNetMap = GetNameNetMap()
+					log.Printf("updated nameNetMap map: %v\n", nameNetMap)
 				}
 			}
 		}

diff --git a/daemon/src/selector/selector.go b/daemon/src/selector/selector.go
@@ -71,8 +71,6 @@ func getDefaultResponse(req NICSelectRequest, masterNameMap map[string]string, n
 }
 
 func Select(req NICSelectRequest) NICSelectResponse {
-	masterNameMap := iface.GetInterfaceNameMap()
-	nameNetMap := iface.GetNameNetMap()
 	deviceMap := make(map[string]string)
 	resourceMap := make(map[string][]string)
 
@@ -93,6 +91,8 @@ func Select(req NICSelectRequest) NICSelectResponse {
 		log.Printf("Cannot get pod: %v\n", err)
 	}
 
+	masterNameMap := iface.GetInterfaceNameMap()
+	nameNetMap := iface.GetNameNetMap()
 	netSpec, err := MultinicnetHandler.Get(req.NetAttachDefName, req.PodNamespace)
 	if err != nil {
 		return getDefaultResponse(req, masterNameMap, nameNetMap, deviceMap, resourceMap)