Merge pull request #101 from trickest/fix/machine-backward-compatibiltiy

Support the small-medium-large machine format for backward compatibility

cmd/execute/execute.go

@@ -86,51 +86,32 @@ var ExecuteCmd = &cobra.Command{
 		}
 
 		allNodes, roots = CreateTrees(version, false)
+
 		if maxMachines {
 			executionMachines = version.MaxMachines
 		} else if machineConfiguration != "" {
-			machines := &types.Machines{}
+			machines, err := parseMachineConfiguration(machineConfiguration)
+			if err != nil {
+				fmt.Printf("Error: %s\n", err)
+				os.Exit(1)
+			}
 
-			// Managed enterprise fleet, 3 types of machines, small-medium-large
 			if len(fleet.Machines) == 3 {
-				pattern := `^\d+-\d+-\d+$`
-				regex := regexp.MustCompile(pattern)
-
-				if regex.MatchString(machineConfiguration) {
-					parts := strings.Split(machineConfiguration, "-")
-
-					if small, err := strconv.Atoi(parts[0]); err == nil && small != 0 {
-						machines.Small = &small
-					}
+				// 3 types of machines: small, medium, and large
+				executionMachines = machines
 
-					if medium, err := strconv.Atoi(parts[1]); err == nil && medium != 0 {
-						machines.Medium = &medium
-					}
-
-					if large, err := strconv.Atoi(parts[2]); err == nil && large != 0 {
-						machines.Large = &large
-					}
-					executionMachines = *machines
-				} else {
-					fmt.Printf("Invalid machine configuration \"%s\".\n", machineConfiguration)
-					fmt.Println("Please use the format: small-medium-large (e.g., 0-0-3)")
+				if machines.Default != nil {
+					fmt.Printf("Error: you need to use the small-medium-large format to specify the numbers of machines (e.g. 1-2-3)")
 					os.Exit(1)
 				}
 			} else {
-				defaultOrSelfHosted, err := strconv.Atoi(machineConfiguration)
+				// 1 type of machine
+				executionMachines, err = handleSingleMachineType(*fleet, machines)
 				if err != nil {
-					fmt.Printf("Invalid machine configuration \"%s\".\n", machineConfiguration)
+					fmt.Printf("Error: %s\n", err)
 					os.Exit(1)
 				}
-
-				if fleet.Type == "MANAGED" {
-					machines.Default = &defaultOrSelfHosted
-				} else if fleet.Type == "HOSTED" {
-					machines.SelfHosted = &defaultOrSelfHosted
-				}
 			}
-			executionMachines = *machines
-
 		} else {
 			executionMachines = setMachinesToMinimum(version.MaxMachines)
 		}
@@ -151,6 +132,70 @@ var ExecuteCmd = &cobra.Command{
 	},
 }
 
+func parseMachineConfiguration(config string) (types.Machines, error) {
+	pattern := `^\d+-\d+-\d+$`
+	regex := regexp.MustCompile(pattern)
+
+	if regex.MatchString(config) {
+		// 3 types of machines, 3 hyphen-delimited inputs
+		parts := strings.Split(config, "-")
+
+		if len(parts) != 3 {
+			return types.Machines{}, fmt.Errorf("invalid number of machines in machine configuration \"%s\"", config)
+		}
+
+		sizes := make([]int, 3)
+		for index, part := range parts {
+			if size, err := strconv.Atoi(part); err == nil {
+				sizes[index] = size
+			} else {
+				return types.Machines{}, fmt.Errorf("invalid machine configuration \"%s\"", config)
+			}
+		}
+
+		return types.Machines{
+			// sizes = [small, medium, large]
+			Small:  &sizes[0],
+			Medium: &sizes[1],
+			Large:  &sizes[2],
+		}, nil
+	}
+
+	// One type of machine
+	val, err := strconv.Atoi(config)
+	if err != nil {
+		return types.Machines{}, fmt.Errorf("invalid machine configuration \"%s\"", config)
+	}
+
+	return types.Machines{Default: &val}, nil
+}
+
+func handleSingleMachineType(fleet types.Fleet, machines types.Machines) (types.Machines, error) {
+	var configMachines types.Machines
+
+	var defaultOrSelfHosted int
+	if machines.Default != nil {
+		defaultOrSelfHosted = *machines.Default
+	} else {
+		// Backward-compatibility with the small-medium-large format
+		defaultOrSelfHosted = *machines.Small + *machines.Medium + *machines.Large
+		fmt.Printf("Warning: You have one type of machine in your fleet. %d identical or self-hosted machines will be used.\n", defaultOrSelfHosted)
+	}
+
+	if defaultOrSelfHosted == 0 {
+		return types.Machines{}, fmt.Errorf("cannot run the workflow on %d machines", defaultOrSelfHosted)
+	}
+
+	if fleet.Type == "MANAGED" {
+		configMachines.Default = &defaultOrSelfHosted
+	} else if fleet.Type == "HOSTED" {
+		configMachines.SelfHosted = &defaultOrSelfHosted
+	} else {
+		return types.Machines{}, fmt.Errorf("unsupported format. Use small-medium-large (e.g., 0-0-3)")
+	}
+	return configMachines, nil
+}
+
 func init() {
 	ExecuteCmd.Flags().StringVar(&newWorkflowName, "set-name", "", "Set workflow name")
 	ExecuteCmd.Flags().StringVar(&configFile, "config", "", "YAML file for run configuration")

cmd/execute/helpers.go

@@ -592,34 +592,19 @@ func uploadFilesIfNeeded(primitiveNodes map[string]*types.PrimitiveNode) {
 }
 
 func maxMachinesTypeCompatible(machines, maxMachines types.Machines) bool {
-	if machines.Default != nil && *machines.Default > *maxMachines.Default {
-		return false
-	}
-	if machines.SelfHosted != nil && *machines.SelfHosted > *maxMachines.SelfHosted {
-		return false
-	}
-
-	if machines.Small != nil && *machines.Small > *maxMachines.Small {
-		return false
-	}
-
-	if machines.Medium != nil && *machines.Medium > *maxMachines.Medium {
-		return false
-	}
-
-	if machines.Large != nil && *machines.Large > *maxMachines.Large {
-		return false
-	}
+	return verifyMachineType(machines.Default, maxMachines.Default) &&
+		verifyMachineType(machines.SelfHosted, maxMachines.SelfHosted) &&
+		verifyMachineType(machines.Small, maxMachines.Small) &&
+		verifyMachineType(machines.Medium, maxMachines.Medium) &&
+		verifyMachineType(machines.Large, maxMachines.Large)
+}
 
-	if (machines.Small != nil && maxMachines.Small == nil) ||
-		(machines.Medium != nil && maxMachines.Medium == nil) ||
-		(machines.Large != nil && maxMachines.Large == nil) {
+func verifyMachineType(machine, maxMachine *int) bool {
+	if machine != nil && maxMachine != nil && *machine > *maxMachine {
 		return false
 	}
 
-	if (machines.Small == nil && maxMachines.Small != nil) ||
-		(machines.Medium == nil && maxMachines.Medium != nil) ||
-		(machines.Large == nil && maxMachines.Large != nil) {
+	if (machine != nil && maxMachine == nil) || (machine == nil && maxMachine != nil) {
 		return false
 	}