Go · 10728 bytes Raw Blame History
1 // SPDX-License-Identifier: AGPL-3.0-or-later
2
3 package main
4
5 import (
6 "context"
7 "encoding/json"
8 "errors"
9 "fmt"
10 "io"
11 "strconv"
12 "strings"
13 "text/tabwriter"
14 "time"
15
16 "github.com/jackc/pgx/v5/pgtype"
17 "github.com/jackc/pgx/v5/pgxpool"
18 "github.com/spf13/cobra"
19
20 "github.com/tenseleyFlow/shithub/internal/actions/runnerlabels"
21 "github.com/tenseleyFlow/shithub/internal/actions/runnertoken"
22 actionsdb "github.com/tenseleyFlow/shithub/internal/actions/sqlc"
23 "github.com/tenseleyFlow/shithub/internal/infra/config"
24 "github.com/tenseleyFlow/shithub/internal/infra/db"
25 "github.com/tenseleyFlow/shithub/internal/infra/metrics"
26 )
27
28 func newAdminRunnerCmd() *cobra.Command {
29 cmd := &cobra.Command{
30 Use: "runner",
31 Short: "Register, list, and revoke Actions runners",
32 }
33 cmd.AddCommand(newAdminRunnerRegisterCmd())
34 cmd.AddCommand(newAdminRunnerListCmd())
35 cmd.AddCommand(newAdminRunnerQueueCmd())
36 cmd.AddCommand(newAdminRunnerRevokeCmd())
37 return cmd
38 }
39
40 func newAdminRunnerRegisterCmd() *cobra.Command {
41 var name string
42 var labelsRaw string
43 var capacity int
44 var output string
45 var expiresIn time.Duration
46 cmd := &cobra.Command{
47 Use: "register --name <name> [--labels self-hosted,linux,ubuntu-latest,x64] [--capacity 1]",
48 Short: "Register an Actions runner and print its token once",
49 RunE: func(cmd *cobra.Command, _ []string) error {
50 name = strings.TrimSpace(name)
51 if name == "" {
52 return errors.New("admin runner register: --name is required")
53 }
54 labels, err := parseRunnerLabels(labelsRaw)
55 if err != nil {
56 return err
57 }
58 if capacity < 1 || capacity > 64 {
59 return errors.New("admin runner register: --capacity must be between 1 and 64")
60 }
61 output = strings.ToLower(strings.TrimSpace(output))
62 switch output {
63 case "", "text":
64 output = "text"
65 case "json":
66 default:
67 return errors.New("admin runner register: --output must be text or json")
68 }
69 if expiresIn < 0 {
70 return errors.New("admin runner register: --expires-in must be non-negative")
71 }
72
73 cfg, err := config.Load(nil)
74 if err != nil {
75 return err
76 }
77 ctx, cancel := context.WithTimeout(cmd.Context(), 30*time.Second)
78 defer cancel()
79 pool, err := openAdminRunnerPool(ctx, cfg, "register")
80 if err != nil {
81 return err
82 }
83 defer pool.Close()
84
85 token, tokenHash, err := runnertoken.New()
86 if err != nil {
87 return fmt.Errorf("admin runner register: mint token: %w", err)
88 }
89 var expiresAt pgtype.Timestamptz
90 var outputExpiresAt *time.Time
91 if expiresIn > 0 {
92 t := time.Now().UTC().Add(expiresIn)
93 expiresAt = pgtype.Timestamptz{Time: t, Valid: true}
94 outputExpiresAt = &t
95 }
96
97 q := actionsdb.New()
98 tx, err := pool.Begin(ctx)
99 if err != nil {
100 return fmt.Errorf("admin runner register: begin: %w", err)
101 }
102 committed := false
103 defer func() {
104 if !committed {
105 _ = tx.Rollback(ctx)
106 }
107 }()
108
109 runner, err := q.InsertRunner(ctx, tx, actionsdb.InsertRunnerParams{
110 Name: name,
111 Labels: labels,
112 Capacity: int32(capacity),
113 RegisteredByUserID: pgtype.Int8{},
114 })
115 if err != nil {
116 return fmt.Errorf("admin runner register: insert runner: %w", err)
117 }
118 if _, err := q.InsertRunnerToken(ctx, tx, actionsdb.InsertRunnerTokenParams{
119 RunnerID: runner.ID,
120 TokenHash: tokenHash,
121 ExpiresAt: expiresAt,
122 }); err != nil {
123 return fmt.Errorf("admin runner register: insert token: %w", err)
124 }
125 if err := tx.Commit(ctx); err != nil {
126 return fmt.Errorf("admin runner register: commit: %w", err)
127 }
128 committed = true
129 metrics.ActionsRunnerRegistrationsTotal.Inc()
130
131 return writeRunnerRegisterOutput(cmd.OutOrStdout(), output, runnerRegisterOutput{
132 ID: runner.ID,
133 Name: runner.Name,
134 Labels: runner.Labels,
135 Capacity: runner.Capacity,
136 Token: token,
137 TokenExpiresAt: outputExpiresAt,
138 })
139 },
140 }
141 cmd.Flags().StringVar(&name, "name", "", "Runner name (letters, numbers, underscore, dash)")
142 cmd.Flags().StringVar(&labelsRaw, "labels", strings.Join(runnerlabels.DefaultShared(), ","), "Comma-separated runner labels")
143 cmd.Flags().IntVar(&capacity, "capacity", 1, "Maximum concurrent jobs this runner may execute")
144 cmd.Flags().DurationVar(&expiresIn, "expires-in", 0, "Registration token lifetime (0 means no expiration)")
145 cmd.Flags().StringVar(&output, "output", "text", "Output format: text or json")
146 return cmd
147 }
148
149 type runnerRegisterOutput struct {
150 ID int64 `json:"id"`
151 Name string `json:"name"`
152 Labels []string `json:"labels"`
153 Capacity int32 `json:"capacity"`
154 Token string `json:"token"`
155 TokenExpiresAt *time.Time `json:"token_expires_at,omitempty"`
156 }
157
158 func writeRunnerRegisterOutput(w io.Writer, format string, out runnerRegisterOutput) error {
159 if format == "json" {
160 enc := json.NewEncoder(w)
161 enc.SetIndent("", " ")
162 return enc.Encode(out)
163 }
164 expires := "never"
165 if out.TokenExpiresAt != nil {
166 expires = out.TokenExpiresAt.Format(time.RFC3339)
167 }
168 _, err := fmt.Fprintf(w,
169 "runner registered\nid: %d\nname: %s\nlabels: %s\ncapacity: %d\ntoken_expires_at: %s\ntoken: %s\n\nStore this token now; shithub never shows it again.\n",
170 out.ID, out.Name, strings.Join(out.Labels, ","), out.Capacity, expires, out.Token)
171 return err
172 }
173
174 func newAdminRunnerListCmd() *cobra.Command {
175 return &cobra.Command{
176 Use: "list",
177 Short: "List registered Actions runners",
178 RunE: func(cmd *cobra.Command, _ []string) error {
179 cfg, err := config.Load(nil)
180 if err != nil {
181 return err
182 }
183 ctx, cancel := context.WithTimeout(cmd.Context(), 30*time.Second)
184 defer cancel()
185 pool, err := openAdminRunnerPool(ctx, cfg, "list")
186 if err != nil {
187 return err
188 }
189 defer pool.Close()
190
191 rows, err := actionsdb.New().ListRunners(ctx, pool)
192 if err != nil {
193 return fmt.Errorf("admin runner list: %w", err)
194 }
195 tw := tabwriter.NewWriter(cmd.OutOrStdout(), 0, 0, 2, ' ', 0)
196 _, _ = fmt.Fprintln(tw, "ID\tNAME\tSTATUS\tCAPACITY\tLABELS\tLAST_HEARTBEAT")
197 for _, r := range rows {
198 last := "never"
199 if r.LastHeartbeatAt.Valid {
200 last = r.LastHeartbeatAt.Time.Format(time.RFC3339)
201 }
202 _, _ = fmt.Fprintf(tw, "%d\t%s\t%s\t%d\t%s\t%s\n",
203 r.ID, r.Name, r.Status, r.Capacity, strings.Join(r.Labels, ","), last)
204 }
205 return tw.Flush()
206 },
207 }
208 }
209
210 func newAdminRunnerQueueCmd() *cobra.Command {
211 var output string
212 cmd := &cobra.Command{
213 Use: "queue",
214 Short: "Summarize queued Actions jobs by runs-on label",
215 RunE: func(cmd *cobra.Command, _ []string) error {
216 output = strings.ToLower(strings.TrimSpace(output))
217 switch output {
218 case "", "text":
219 output = "text"
220 case "json":
221 default:
222 return errors.New("admin runner queue: --output must be text or json")
223 }
224 cfg, err := config.Load(nil)
225 if err != nil {
226 return err
227 }
228 ctx, cancel := context.WithTimeout(cmd.Context(), 30*time.Second)
229 defer cancel()
230 pool, err := openAdminRunnerPool(ctx, cfg, "queue")
231 if err != nil {
232 return err
233 }
234 defer pool.Close()
235
236 rows, err := actionsdb.New().ListQueuedWorkflowJobRunsOn(ctx, pool)
237 if err != nil {
238 return fmt.Errorf("admin runner queue: %w", err)
239 }
240 return writeRunnerQueueOutput(cmd.OutOrStdout(), output, rows, time.Now().UTC())
241 },
242 }
243 cmd.Flags().StringVar(&output, "output", "text", "Output format: text or json")
244 return cmd
245 }
246
247 type runnerQueueOutputRow struct {
248 RunsOn string `json:"runs_on"`
249 QueuedJobs int32 `json:"queued_jobs"`
250 MatchingRunnerCount int32 `json:"matching_runner_count"`
251 OldestQueuedAt string `json:"oldest_queued_at,omitempty"`
252 OldestQueuedSeconds int64 `json:"oldest_queued_seconds,omitempty"`
253 }
254
255 func writeRunnerQueueOutput(w io.Writer, format string, rows []actionsdb.ListQueuedWorkflowJobRunsOnRow, now time.Time) error {
256 out := make([]runnerQueueOutputRow, 0, len(rows))
257 for _, row := range rows {
258 item := runnerQueueOutputRow{
259 RunsOn: row.RunsOn,
260 QueuedJobs: row.QueuedJobs,
261 MatchingRunnerCount: row.MatchingRunnerCount,
262 }
263 if row.OldestQueuedAt.Valid {
264 item.OldestQueuedAt = row.OldestQueuedAt.Time.UTC().Format(time.RFC3339)
265 if d := now.Sub(row.OldestQueuedAt.Time); d > 0 {
266 item.OldestQueuedSeconds = int64(d.Seconds())
267 }
268 }
269 out = append(out, item)
270 }
271 if format == "json" {
272 enc := json.NewEncoder(w)
273 enc.SetIndent("", " ")
274 return enc.Encode(out)
275 }
276
277 tw := tabwriter.NewWriter(w, 0, 0, 2, ' ', 0)
278 _, _ = fmt.Fprintln(tw, "RUNS_ON\tQUEUED_JOBS\tMATCHING_RUNNERS\tOLDEST_QUEUED")
279 for _, row := range out {
280 oldest := "-"
281 if row.OldestQueuedAt != "" {
282 oldest = row.OldestQueuedAt
283 }
284 _, _ = fmt.Fprintf(tw, "%s\t%d\t%d\t%s\n", row.RunsOn, row.QueuedJobs, row.MatchingRunnerCount, oldest)
285 }
286 return tw.Flush()
287 }
288
289 func newAdminRunnerRevokeCmd() *cobra.Command {
290 var idRaw string
291 cmd := &cobra.Command{
292 Use: "revoke --id <id>",
293 Short: "Revoke all registration tokens for an Actions runner",
294 RunE: func(cmd *cobra.Command, _ []string) error {
295 id, err := strconv.ParseInt(strings.TrimSpace(idRaw), 10, 64)
296 if err != nil || id <= 0 {
297 return errors.New("admin runner revoke: --id must be a positive integer")
298 }
299 cfg, err := config.Load(nil)
300 if err != nil {
301 return err
302 }
303 ctx, cancel := context.WithTimeout(cmd.Context(), 30*time.Second)
304 defer cancel()
305 pool, err := openAdminRunnerPool(ctx, cfg, "revoke")
306 if err != nil {
307 return err
308 }
309 defer pool.Close()
310
311 q := actionsdb.New()
312 runner, err := q.GetRunnerByID(ctx, pool, id)
313 if err != nil {
314 return fmt.Errorf("admin runner revoke: runner %d not found", id)
315 }
316 if err := q.RevokeAllTokensForRunner(ctx, pool, id); err != nil {
317 return fmt.Errorf("admin runner revoke: %w", err)
318 }
319 _, _ = fmt.Fprintf(cmd.OutOrStdout(), "runner revoked\nid: %d\nname: %s\n", runner.ID, runner.Name)
320 return nil
321 },
322 }
323 cmd.Flags().StringVar(&idRaw, "id", "", "Runner id")
324 return cmd
325 }
326
327 func openAdminRunnerPool(ctx context.Context, cfg config.Config, op string) (*pgxpool.Pool, error) {
328 if cfg.DB.URL == "" {
329 return nil, fmt.Errorf("admin runner %s: DB not configured (set SHITHUB_DATABASE_URL)", op)
330 }
331 pool, err := db.Open(ctx, db.Config{
332 URL: cfg.DB.URL, MaxConns: 2, MinConns: 0,
333 ConnectTimeout: cfg.DB.ConnectTimeout,
334 })
335 if err != nil {
336 return nil, fmt.Errorf("admin runner %s: db open: %w", op, err)
337 }
338 return pool, nil
339 }
340
341 func parseRunnerLabels(raw string) ([]string, error) {
342 return runnerlabels.ParseCSV(raw)
343 }
344
345 func init() {
346 adminCmd.AddCommand(newAdminRunnerCmd())
347 adminActionsCmd.AddCommand(newAdminRunnerCmd())
348 }
349