-
Notifications
You must be signed in to change notification settings - Fork 238
/
count.go
142 lines (119 loc) · 3.55 KB
/
count.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
package scale
import (
"context"
"fmt"
"strconv"
"strings"
"github.com/spf13/cobra"
"github.com/superfly/flyctl/api"
"github.com/superfly/flyctl/client"
"github.com/superfly/flyctl/flaps"
"github.com/superfly/flyctl/internal/appconfig"
"github.com/superfly/flyctl/internal/command"
"github.com/superfly/flyctl/internal/flag"
"github.com/superfly/flyctl/iostreams"
"golang.org/x/exp/slices"
)
func newScaleCount() *cobra.Command {
const (
short = "Change an app's VM count to the given value"
long = `Change an app's VM count to the given value.
For pricing, see https://fly.io/docs/about/pricing/`
)
cmd := command.New("count [count]", short, long, runScaleCount,
command.RequireSession,
command.RequireAppName,
)
cmd.Args = cobra.MinimumNArgs(1)
flag.Add(cmd,
flag.App(),
flag.AppConfig(),
flag.Yes(),
flag.Int{Name: "max-per-region", Description: "Max number of VMs per region", Default: -1},
flag.String{Name: "region", Description: "Comma separated list of regions to act on. Defaults to all regions where there is at least one machine running for the app"},
flag.String{Name: "process-group", Description: "The process group to scale"},
)
return cmd
}
func runScaleCount(ctx context.Context) error {
appName := appconfig.NameFromContext(ctx)
flapsClient, err := flaps.NewFromAppName(ctx, appName)
if err != nil {
return err
}
ctx = flaps.NewContext(ctx, flapsClient)
appConfig, err := appconfig.FromRemoteApp(ctx, appName)
if err != nil {
return err
}
args := flag.Args(ctx)
processNames := appConfig.ProcessNames()
groupName := flag.GetString(ctx, "process-group")
if groupName == "" {
groupName = api.MachineProcessGroupApp
if !slices.Contains(processNames, groupName) {
return fmt.Errorf("--process-group flag is required when no group named 'app' is defined")
}
}
if !slices.Contains(processNames, groupName) {
return fmt.Errorf("process group '%s' not found", groupName)
}
groups, err := parseGroupCounts(args, groupName)
if err != nil {
return err
}
maxPerRegion := flag.GetInt(ctx, "max-per-region")
isV2, err := command.IsMachinesPlatform(ctx, appName)
if err != nil {
return err
}
if isV2 {
return runMachinesScaleCount(ctx, appName, appConfig, groups, maxPerRegion)
}
return runNomadScaleCount(ctx, appName, groups, maxPerRegion)
}
func parseGroupCounts(args []string, defaultGroupName string) (map[string]int, error) {
groups := make(map[string]int)
// single numeric arg: fly scale count 3
if len(args) == 1 {
count, err := strconv.Atoi(args[0])
if err == nil {
groups[defaultGroupName] = count
}
}
// group labels: fly scale web=X worker=Y
if len(groups) < 1 {
for _, arg := range args {
parts := strings.Split(arg, "=")
if len(parts) != 2 {
return nil, fmt.Errorf("'%s' is not a valid process=count option", arg)
}
count, err := strconv.Atoi(parts[1])
if err != nil {
return nil, err
}
groups[parts[0]] = count
}
}
return groups, nil
}
func runNomadScaleCount(ctx context.Context, appName string, groups map[string]int, maxPerRegion int) error {
io := iostreams.FromContext(ctx)
apiClient := client.FromContext(ctx).API()
var maxPerRegionPtr *int
if maxPerRegion >= 0 {
maxPerRegionPtr = &maxPerRegion
}
counts, warnings, err := apiClient.SetAppVMCount(ctx, appName, groups, maxPerRegionPtr)
if err != nil {
return err
}
if len(warnings) > 0 {
for _, warning := range warnings {
fmt.Fprintln(io.Out, "Warning:", warning)
}
fmt.Fprintln(io.Out)
}
fmt.Fprintf(io.Out, "Count changed to %s\n", countMessage(counts))
return nil
}