let cpumap has getnumanode func

projecteru2 · Jun 11, 2019 · 1212451 · 1212451
1 parent 132231f
commit 1212451
Show file tree

Hide file tree

Showing 7 changed files with 42 additions and 31 deletions.
diff --git a/cluster/calcium/create.go b/cluster/calcium/create.go
@@ -245,7 +245,7 @@ func (c *Calcium) doMakeContainerOptions(index int, cpumap types.CPUMap, opts *t
 	config.CPU = cpumap.Map()
 	config.Quota = opts.CPUQuota
 	config.Memory = opts.Memory
-	config.NUMANode = utils.GetNUMAMemoryNode(node, cpumap)
+	config.NUMANode = cpumap.GetNUMANode(node)
 	config.SoftLimit = opts.SoftLimit
 	entry := opts.Entrypoint
 

diff --git a/cluster/calcium/realloc.go b/cluster/calcium/realloc.go
@@ -109,7 +109,7 @@ func (c *Calcium) doReallocContainer(
 						node.CPU.Add(container.CPU)
 						node.SetCPUUsed(container.Quota, types.DecrUsage)
 						node.MemCap += container.Memory
-						if nodeID := utils.GetNUMAMemoryNode(node, container.CPU); nodeID != "" {
+						if nodeID := container.CPU.GetNUMANode(node); nodeID != "" {
 							if _, ok := node.NUMAMem[nodeID]; ok {
 								node.NUMAMem[nodeID] += container.Memory
 							}
@@ -143,7 +143,7 @@ func (c *Calcium) doReallocContainer(
 						newResource := &enginetypes.VirtualizationResource{Quota: newCPU, Memory: newMemory, SoftLimit: container.SoftLimit}
 						if len(container.CPU) > 0 {
 							newResource.CPU = cpusets[0]
-							newResource.NUMANode = utils.GetNUMAMemoryNode(node, newResource.CPU)
+							newResource.NUMANode = cpusets[0].GetNUMANode(node)
 							cpusets = cpusets[1:]
 						}
 						updateSuccess := false
@@ -162,7 +162,7 @@ func (c *Calcium) doReallocContainer(
 						node.CPU.Sub(container.CPU)
 						node.SetCPUUsed(container.Quota, types.IncrUsage)
 						node.MemCap -= container.Memory
-						if nodeID := utils.GetNUMAMemoryNode(node, container.CPU); nodeID != "" {
+						if nodeID := container.CPU.GetNUMANode(node); nodeID != "" {
 							if _, ok := node.NUMAMem[nodeID]; ok {
 								node.NUMAMem[nodeID] -= container.Memory
 							}

diff --git a/store/etcdv3/node.go b/store/etcdv3/node.go
@@ -173,7 +173,7 @@ func (m *Mercury) UpdateNodeResource(ctx context.Context, node *types.Node, cpu
 		node.CPU.Add(cpu)
 		node.SetCPUUsed(quota, types.DecrUsage)
 		node.MemCap += mem
-		if nodeID := utils.GetNUMAMemoryNode(node, cpu); nodeID != "" {
+		if nodeID := cpu.GetNUMANode(node); nodeID != "" {
 			if _, ok := node.NUMAMem[nodeID]; ok {
 				node.NUMAMem[nodeID] += mem
 			}
@@ -182,7 +182,7 @@ func (m *Mercury) UpdateNodeResource(ctx context.Context, node *types.Node, cpu
 		node.CPU.Sub(cpu)
 		node.SetCPUUsed(quota, types.IncrUsage)
 		node.MemCap -= mem
-		if nodeID := utils.GetNUMAMemoryNode(node, cpu); nodeID != "" {
+		if nodeID := cpu.GetNUMANode(node); nodeID != "" {
 			if _, ok := node.NUMAMem[nodeID]; ok {
 				node.NUMAMem[nodeID] -= mem
 			}

diff --git a/types/node.go b/types/node.go
@@ -53,6 +53,21 @@ func (c CPUMap) Map() map[string]int {
 	return map[string]int(c)
 }
 
+// GetNUMANode get numa node
+func (c CPUMap) GetNUMANode(node *Node) string {
+	nodeID := ""
+	for cpuID := range c {
+		if memoryNode, ok := node.NUMA[cpuID]; ok {
+			if nodeID == "" {
+				nodeID = memoryNode
+			} else if nodeID != memoryNode { // 如果跨 NODE 了，让系统决定 nodeID
+				nodeID = ""
+			}
+		}
+	}
+	return nodeID
+}
+
 // NUMA define NUMA cpuID->nodeID
 type NUMA map[string]string
 

diff --git a/types/node_test.go b/types/node_test.go
@@ -40,6 +40,27 @@ func TestCPUMap(t *testing.T) {
 	cpuMap.Add(CPUMap{"0": 20})
 	assert.Equal(t, cpuMap["0"], 70)
 
+	cpuMap.Add(CPUMap{"3": 100})
+	assert.Equal(t, cpuMap["3"], 100)
+
 	cpuMap.Sub(CPUMap{"1": 20})
 	assert.Equal(t, cpuMap["1"], 50)
+
+	m := cpuMap.Map()
+	assert.Equal(t, m["1"], 50)
+}
+
+func TestGetNUMANode(t *testing.T) {
+	node := &Node{
+		NUMA: NUMA{"1": "node1", "2": "node2", "3": "node1", "4": "node2"},
+	}
+	cpu := CPUMap{"1": 100, "2": 100}
+	nodeID := cpu.GetNUMANode(node)
+	assert.Equal(t, nodeID, "")
+	cpu = CPUMap{"1": 100, "3": 100}
+	nodeID = cpu.GetNUMANode(node)
+	assert.Equal(t, nodeID, "node1")
+	cpu = nil
+	nodeID = cpu.GetNUMANode(node)
+	assert.Equal(t, nodeID, "")
 }
diff --git a/utils/utils.go b/utils/utils.go
@@ -262,18 +262,3 @@ func safeSplit(s string) []string {
 
 	return result
 }
-
-// GetNUMAMemoryNode get numa memory node
-func GetNUMAMemoryNode(node *types.Node, cpu types.CPUMap) string {
-	nodeID := ""
-	for cpuID := range cpu {
-		if memoryNode, ok := node.NUMA[cpuID]; ok {
-			if nodeID == "" {
-				nodeID = memoryNode
-			} else if nodeID != memoryNode { // 如果跨 NODE 了，让系统决定 nodeID
-				nodeID = ""
-			}
-		}
-	}
-	return nodeID
-}
diff --git a/utils/utils_test.go b/utils/utils_test.go
@@ -201,13 +201,3 @@ func TestRound(t *testing.T) {
 	a = 19.99998
 	assert.Equal(t, f(Round(a)), "20")
 }
-
-func TestGetNUMAMemoryNode(t *testing.T) {
-	node := &types.Node{
-		NUMA: types.NUMA{"1": "node1", "2": "node2", "3": "node1", "4": "node2"},
-	}
-	nodeID := GetNUMAMemoryNode(node, types.CPUMap{"1": 100, "2": 100})
-	assert.Equal(t, nodeID, "")
-	nodeID = GetNUMAMemoryNode(node, types.CPUMap{"1": 100, "3": 100})
-	assert.Equal(t, nodeID, "node1")
-}