fix: get node in register

Signed-off-by: james <open4pd@4paradigm.com>

Author: DSFans2014

internal/pkg/api/device/ascend/device.go

@@ -101,7 +101,7 @@ func (dev *Devices) GetResource(n *corev1.Node) map[string]int {
 		resourceName: 0,
 	}
 	if !device.CheckHealthy(n, dev.config.ResourceName) {
-		klog.Infof("device %s is not healthy on this node", dev.CommonWord())
+		klog.Infof("device %s is unhealthy on this node", dev.CommonWord())
 		return resourceMap
 	}
 	devInfos, err := dev.GetNodeDevices(n)
@@ -121,9 +121,9 @@ func (dev *Devices) GetResource(n *corev1.Node) map[string]int {
 	return resourceMap
 }
 
-func (dev *Devices) RunManager(n *corev1.Node) {
+func (dev *Devices) RunManager() {
 	lmock := mock.NewMockLister(device.GetVendorName(dev.config.ResourceMemoryName))
-	go device.Register(n, lmock, dev)
+	go device.Register(lmock, dev)
 	mockmanager := dpm.NewManager(lmock)
 	klog.Infof("Running mocking dp: %s", dev.CommonWord())
 	mockmanager.Run()

internal/pkg/api/device/device.go

@@ -74,7 +74,7 @@ type Devices interface {
 	CommonWord() string
 	GetNodeDevices(n *corev1.Node) ([]*DeviceInfo, error)
 	GetResource(n *corev1.Node) map[string]int
-	RunManager(n *corev1.Node)
+	RunManager()
 }
 
 type ResourceNames struct {
@@ -101,27 +101,27 @@ func GetDevices() map[string]Devices {
 }
 
 func RunManagers() error {
-	nodeName := os.Getenv("NODE_NAME")
-	node, err := client.GetClient().CoreV1().Nodes().Get(context.Background(), nodeName, v1.GetOptions{})
-	if err != nil {
-		klog.Error("Get node error", err.Error())
-		return err
-	}
 	for name, dev := range DevicesMap {
 		klog.Infof("%s run manager", name)
 		ch[name] = make(chan int)
-		go dev.RunManager(node)
+		go dev.RunManager()
 	}
 	for _, val := range ch {
 		<-val
 	}
 	return nil
 }
 
-func Register(n *corev1.Node, l *mock.MockLister, dev Devices) {
+func Register(l *mock.MockLister, dev Devices) {
+	nodeName := os.Getenv("NODE_NAME")
 	for {
-		resourceMap := dev.GetResource(n)
-		l.SetResource(resourceMap)
+		node, err := client.GetClient().CoreV1().Nodes().Get(context.Background(), nodeName, v1.GetOptions{})
+		if err != nil {
+			klog.Error("Get node error", err.Error())
+		} else {
+			resourceMap := dev.GetResource(node)
+			l.SetResource(resourceMap)
+		}
 		time.Sleep(time.Second * 30)
 	}
 }

internal/pkg/api/device/hygon/device.go

@@ -90,7 +90,7 @@ func (dev *DCUDevices) GetResource(n *corev1.Node) map[string]int {
 		memoryResourceName: 0,
 	}
 	if !device.CheckHealthy(n, HygonResourceCount) {
-		klog.Infof("device %s is not healthy on this node", dev.CommonWord())
+		klog.Infof("device %s is unhealthy on this node", dev.CommonWord())
 		return resourceMap
 	}
 	devs, err := dev.GetNodeDevices(n)
@@ -110,9 +110,9 @@ func (dev *DCUDevices) GetResource(n *corev1.Node) map[string]int {
 	return resourceMap
 }
 
-func (dev *DCUDevices) RunManager(n *corev1.Node) {
+func (dev *DCUDevices) RunManager() {
 	lmock := mock.NewMockLister(device.GetVendorName(HygonResourceMemory))
-	device.Register(n, lmock, dev)
+	device.Register(lmock, dev)
 	mockmanager := dpm.NewManager(lmock)
 	klog.Infof("Running mocking dp: %s", dev.CommonWord())
 	mockmanager.Run()

internal/pkg/api/device/kunlun/device.go

@@ -100,9 +100,9 @@ func (dev *KunlunVDevices) GetResource(n *corev1.Node) map[string]int {
 	return resourceMap
 }
 
-func (dev *KunlunVDevices) RunManager(n *corev1.Node) {
+func (dev *KunlunVDevices) RunManager() {
 	lmock := mock.NewMockLister(device.GetVendorName(KunlunResourceVCount))
-	device.Register(n, lmock, dev)
+	device.Register(lmock, dev)
 	mockmanager := dpm.NewManager(lmock)
 	klog.Infof("Running mocking dp: %s", dev.CommonWord())
 	mockmanager.Run()

internal/pkg/api/device/nvidia/device.go

@@ -171,7 +171,7 @@ func (dev *NvidiaGPUDevices) GetResource(n *corev1.Node) map[string]int {
 		memoryPercentageName: 0,
 	}
 	if !device.CheckHealthy(n, dev.config.ResourceCountName) {
-		klog.Infof("device %s is not healthy on this node", dev.CommonWord())
+		klog.Infof("device %s is unhealthy on this node", dev.CommonWord())
 		return resourceMap
 	}
 	devs, err := dev.GetNodeDevices(n)
@@ -200,9 +200,9 @@ func (dev *NvidiaGPUDevices) GetResource(n *corev1.Node) map[string]int {
 	return resourceMap
 }
 
-func (dev *NvidiaGPUDevices) RunManager(n *corev1.Node) {
+func (dev *NvidiaGPUDevices) RunManager() {
 	lmock := mock.NewMockLister(Vendor)
-	go device.Register(n, lmock, dev)
+	go device.Register(lmock, dev)
 	mockmanager := dpm.NewManager(lmock)
 	klog.Infof("Running mocking dp: %s", dev.CommonWord())
 	mockmanager.Run()

internal/pkg/api/device/nvidia/device_test.go

@@ -176,7 +176,7 @@ func TestGetNodeDevices(t *testing.T) {
 			node := tt.setupNode()
 			dev := tt.setupDev()
 
-			devices, err := dev.GetNodeDevices(node)
+			devices, err := dev.GetNodeDevices(&node)
 
 			if (err != nil) != tt.wantErr {
 				t.Errorf("GetNodeDevices() error = %v, wantErr %v", err, tt.wantErr)

internal/pkg/mock/lister.go

@@ -84,9 +84,9 @@ func (l *MockLister) SetResource(resourceMap map[string]int) {
 		return
 	}
 	l.mutex.Lock()
+	defer l.mutex.Unlock()
 	l.counts = resourceMap
 	pluginNums := len(l.pluginsMap)
-	l.mutex.Unlock()
 
 	if pluginNums == 0 {
 		resourceNames := make([]string, 0, len(resourceMap))