fix node Add/Set with extra resources

projecteru2 · May 15, 2023 · 701f6f7 · 701f6f7
1 parent 1987591
commit 701f6f7
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 6 deletions.
diff --git a/cmd/node/add.go b/cmd/node/add.go
@@ -172,20 +172,19 @@ func generateAddNodeOptions(c *cli.Context) (*corepb.AddNodeOptions, error) {
 		"cpumem":  cb,
 		"storage": sb,
 	}
-	extraResources := c.String("extra-resources")
-	if extraResources != "" {
-		extraResourcesMap := make(map[string]any)
-		if err := json.Unmarshal([]byte(extraResources), &extraResourcesMap); err != nil {
-			return nil, fmt.Errorf("Invalid value for extra-resources: %v", err)
-		}
+
+	if extraResourcesMap, err := utils.ParseExtraResources(c); err != nil {
 		for k, v := range extraResourcesMap {
 			if _, ok := resources[k]; ok {
 				continue
 			}
 			eb, _ := json.Marshal(v)
 			resources[k] = eb
 		}
+	} else {
+		return nil, fmt.Errorf("[generateAddNodeOptions] get extra resources failed %v", err)
 	}
+
 	labels := utils.SplitEquality(c.StringSlice("label"))
 	return &corepb.AddNodeOptions{
 		Nodename:  nodename,

diff --git a/cmd/node/cmd.go b/cmd/node/cmd.go
@@ -184,6 +184,11 @@ func Command() *cli.Command {
 						Usage: "key file, like /etc/docker/tls/client.key",
 						Value: "",
 					},
+					&cli.StringFlag{
+						Name:  "extra-resources",
+						Usage: "add extra resource requests",
+						Value: "",
+					},
 				},
 			},
 			{
@@ -261,6 +266,11 @@ func Command() *cli.Command {
 						Name:  "test",
 						Usage: "mark node for testing, maybe no health check and status report",
 					},
+					&cli.StringFlag{
+						Name:  "extra-resources",
+						Usage: "add extra resource requests",
+						Value: "",
+					},
 				},
 			},
 		},

diff --git a/cmd/node/set.go b/cmd/node/set.go
@@ -3,6 +3,7 @@ package node
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 
 	"github.com/juju/errors"
 	"github.com/sirupsen/logrus"
@@ -98,6 +99,18 @@ func generateSetNodeOptions(c *cli.Context, _ corepb.CoreRPCClient) (*corepb.Set
 		"storage": sb,
 	}
 
+	if extraResourcesMap, err := utils.ParseExtraResources(c); err != nil {
+		for k, v := range extraResourcesMap {
+			if _, ok := resources[k]; ok {
+				continue
+			}
+			eb, _ := json.Marshal(v)
+			resources[k] = eb
+		}
+	} else {
+		return nil, fmt.Errorf("[generateSetNodeOptions] get extra resources failed %v", err)
+	}
+
 	return &corepb.SetNodeOptions{
 		Nodename:      name,
 		Resources:     resources,