1
2
3 package main
4
5 import (
6 "context"
7 "errors"
8 "fmt"
9 "os"
10 "strconv"
11
12 "github.com/Microsoft/hcsshim/internal/jobobject"
13 "github.com/urfave/cli"
14 )
15
16 const (
17 cpuLimitFlag = "cpu-limit"
18 cpuWeightFlag = "cpu-weight"
19 memoryLimitFlag = "memory-limit"
20 affinityFlag = "cpu-affinity"
21 useNTVariantFlag = "use-nt"
22
23 usage = `jobobject-util is a command line tool for getting and setting job object limits`
24 )
25
26 func main() {
27 app := cli.NewApp()
28 app.Name = "jobobject-util"
29 app.Commands = []cli.Command{
30 getJobObjectLimitsCommand,
31 setJobObjectLimitsCommand,
32 }
33 app.Usage = usage
34
35 if err := app.Run(os.Args); err != nil {
36 fmt.Fprintln(os.Stderr, err)
37 os.Exit(1)
38 }
39 }
40
41 var getJobObjectLimitsCommand = cli.Command{
42 Name: "get",
43 Usage: "gets the job object's resource limits",
44 ArgsUsage: "get [flags] <name>",
45 Flags: []cli.Flag{
46 cli.BoolFlag{
47 Name: useNTVariantFlag,
48 Usage: `Optional: indicates if the command should use the NT variant of job object Open/Create calls. `,
49 },
50 cli.BoolFlag{
51 Name: cpuLimitFlag,
52 Usage: "Optional: get job object's CPU limit",
53 },
54 cli.BoolFlag{
55 Name: cpuWeightFlag,
56 Usage: "Optional: get job object's CPU weight.",
57 },
58 cli.BoolFlag{
59 Name: memoryLimitFlag,
60 Usage: "Optional: get job object's memory limit in bytes.",
61 },
62 cli.BoolFlag{
63 Name: affinityFlag,
64 Usage: "Optional: get job object's CPU affinity as a bitmask.",
65 },
66 },
67 Action: func(cli *cli.Context) error {
68 ctx := context.Background()
69 name := cli.Args().First()
70 if name == "" {
71 return errors.New("`get` command must specify a target job object name")
72 }
73 options := &jobobject.Options{
74 Name: name,
75 Notifications: false,
76 UseNTVariant: cli.Bool(useNTVariantFlag),
77 }
78 job, err := jobobject.Open(ctx, options)
79 if err != nil {
80 return err
81 }
82 defer job.Close()
83
84 output := ""
85
86
87
88 if cli.IsSet(cpuLimitFlag) && cli.IsSet(cpuWeightFlag) {
89 return errors.New("cpu limit and weight are mutually exclusive")
90 }
91 if cli.IsSet(cpuLimitFlag) {
92 cpuRate, err := job.GetCPULimit(jobobject.RateBased)
93 if err != nil {
94 return err
95 }
96 output += fmt.Sprintf("%s: %d\n", cpuLimitFlag, cpuRate)
97 } else if cli.IsSet(cpuWeightFlag) {
98 cpuWeight, err := job.GetCPULimit(jobobject.WeightBased)
99 if err != nil {
100 return err
101 }
102 output += fmt.Sprintf("%s: %d\n", cpuWeightFlag, cpuWeight)
103 }
104
105 if cli.IsSet(memoryLimitFlag) {
106 jobObjMemLimit, err := job.GetMemoryLimit()
107 if err != nil {
108 return err
109 }
110 output += fmt.Sprintf("%s: %d\n", memoryLimitFlag, jobObjMemLimit)
111 }
112
113 if cli.IsSet(affinityFlag) {
114 affinity, err := job.GetCPUAffinity()
115 if err != nil {
116 return err
117 }
118 affinityString := strconv.FormatUint(affinity, 2)
119 output += fmt.Sprintf("%s: %s\n", affinityFlag, affinityString)
120 }
121 fmt.Fprintln(os.Stdout, output)
122 return nil
123 },
124 }
125
126 var setJobObjectLimitsCommand = cli.Command{
127 Name: "set",
128 Usage: "tool used to set resource limits on job objects",
129 ArgsUsage: "set [flags] <name>",
130 Flags: []cli.Flag{
131 cli.BoolFlag{
132 Name: useNTVariantFlag,
133 Usage: `Optional: indicates if the command should use the NT variant of job object Open/Create calls. `,
134 },
135 cli.Uint64Flag{
136 Name: cpuLimitFlag,
137 Usage: "Optional: set job object's CPU limit",
138 },
139 cli.Uint64Flag{
140 Name: cpuWeightFlag,
141 Usage: "Optional: set job object's CPU weight.",
142 },
143 cli.Uint64Flag{
144 Name: memoryLimitFlag,
145 Usage: "Optional: set job object's memory limit in bytes.",
146 },
147 cli.StringFlag{
148 Name: affinityFlag,
149 Usage: "Optional: set job object's CPU affinity given a bitmask",
150 },
151 },
152 Action: func(cli *cli.Context) error {
153 ctx := context.Background()
154 name := cli.Args().First()
155 if name == "" {
156 return errors.New("`set` command must specify a job object name")
157 }
158
159 options := &jobobject.Options{
160 Name: name,
161 Notifications: false,
162 UseNTVariant: cli.Bool(useNTVariantFlag),
163 }
164 job, err := jobobject.Open(ctx, options)
165 if err != nil {
166 return err
167 }
168 defer job.Close()
169
170
171
172 if cli.IsSet(cpuLimitFlag) && cli.IsSet(cpuWeightFlag) {
173 return errors.New("cpu limit and weight are mutually exclusive")
174 }
175 if cli.IsSet(cpuLimitFlag) {
176 cpuRate := uint32(cli.Uint64(cpuLimitFlag))
177 if err := job.SetCPULimit(jobobject.RateBased, cpuRate); err != nil {
178 return err
179 }
180 } else if cli.IsSet(cpuWeightFlag) {
181 cpuWeight := uint32(cli.Uint64(cpuWeightFlag))
182 if err := job.SetCPULimit(jobobject.WeightBased, cpuWeight); err != nil {
183 return err
184 }
185 }
186
187 if cli.IsSet(memoryLimitFlag) {
188 memLimitInBytes := cli.Uint64(memoryLimitFlag)
189 if err := job.SetMemoryLimit(memLimitInBytes); err != nil {
190 return err
191 }
192 }
193
194 if cli.IsSet(affinityFlag) {
195 affinityString := cli.String(affinityFlag)
196 affinity, err := strconv.ParseUint(affinityString, 2, 64)
197 if err != nil {
198 return err
199 }
200 if err := job.SetCPUAffinity(affinity); err != nil {
201 return err
202 }
203 }
204
205 return nil
206 },
207 }
208
View as plain text