From 7c6ce52074d272b1f378cb34e21efa4db8e5a570 Mon Sep 17 00:00:00 2001 From: "K.Milchev" <49519924+KristiforMilchev@users.noreply.github.com> Date: Tue, 3 Nov 2026 14:15:58 +0200 Subject: [PATCH 1/2] Controller: Connect stream, worker sessions and the worker client - Connect requires the worker's current client certificate; the worker is identified by the certificate, never by what it says in Hello - Sessions with heartbeats, a silence timeout, a registry the scheduler can send through, and replacement when the same worker reconnects - Revocation: controller worker revoke disconnects at once; workers revoked elsewhere are dropped at the next periodic check - pkg/worker: Join saves the identity to disk; Client connects, sends heartbeats and reconnects with jittered backoff, stopping for good when the controller rejects the worker permanently - internal/controller/api: handler and TLS config shared by the binary and tests - Integration tests: real controller, PostgreSQL and worker client over TLS - controller worker list and controller worker revoke Co-Authored-By: Claude Opus 5.5 --- README.md | 27 +- cmd/controller/main.go | 136 ++++--- cmd/controller/main_test.go | 11 +- docs/ROADMAP.md | 9 +- internal/controller/api/api.go | 66 ++++ internal/controller/store/db/workers.sql.go | 36 +- internal/controller/store/queries/workers.sql | 12 +- internal/controller/workers/connect.go | 300 +++++++++++++++ internal/controller/workers/join.go | 15 + internal/controller/workers/registry.go | 169 ++++++++ internal/controller/workers/registry_test.go | 72 ++++ internal/integration/connect_test.go | 361 ++++++++++++++++++ pkg/worker/client.go | 343 +++++++++++++++++ 13 files changed, 1485 insertions(+), 72 deletions(-) create mode 100644 internal/controller/api/api.go create mode 100644 internal/controller/workers/connect.go create mode 100644 internal/controller/workers/registry.go create mode 100644 internal/controller/workers/registry_test.go create mode 100644 internal/integration/connect_test.go create mode 100644 pkg/worker/client.go diff --git a/README.md b/README.md index 4bfa51c..65a6439 100644 --- a/README.md +++ b/README.md @@ -3,8 +3,9 @@ Core of [Opifica](https://github.com/opifica), an open-source CI/CD system built in the EU where every job runs in a fresh virtual machine. -> **Status:** early development. The pipeline language, the API definitions and the worker kit's warm pool -> work and are tested. The controller, agent and signer are skeletons. +> **Status:** early development. Working and tested: the pipeline language, the worker kit's warm pool, and +> workers joining and staying connected to the controller. Scheduling and the agent come next; see +> [docs/ROADMAP.md](docs/ROADMAP.md). > See the [organisation page](https://github.com/opifica) for the architecture. ## What lives here @@ -15,9 +16,10 @@ where every job runs in a fresh virtual machine. | `schema/` | JSON Schema for pipeline files, for editor autocompletion | Working | | `cmd/opifica` | CLI: `validate`, `plan`, `schema` | Working | | `proto/` | API definitions: worker ↔ controller, worker ↔ guest agent (Protobuf, ConnectRPC) | Draft | -| `pkg/worker` | Shared worker kit: `Hypervisor` interface, warm VM pool, image verification | Working | +| `pkg/worker` | Shared worker kit: `Hypervisor` interface, warm VM pool, image verification, controller client | Working | | `pkg/worker/workertest` | Fake hypervisor for testing worker code without VMs | Working | -| `cmd/controller` | API, Git providers, pipeline engine, scheduler, secrets broker | Skeleton | +| `cmd/controller` | Join tokens, worker joining over TLS, the Connect stream, worker revocation | Working | +| | Scheduler, runs, logs, secrets, Git providers | Not started | | `cmd/agent` | Runs inside each VM, executes steps, streams logs and results | Skeleton | | `cmd/signer` | Signing service for release artifacts | Not started | @@ -111,9 +113,22 @@ broken.yaml:11:5: unknown field "secret" Requires Go 1.27 or later. ```sh -make tools # install buf and the protoc plugins (once) +make tools # install buf, the protoc plugins and sqlc (once) make # lint, test and build into bin/ -make generate # after changing proto/; generated code in gen/ is committed +make generate # after changing proto/ or SQL queries; generated code is committed +``` + +Database tests need PostgreSQL. Either set `OPIFICA_TEST_DATABASE_URL` to a server the tests may create databases +on, or install the PostgreSQL server binaries (e.g. `apt install postgresql-17`): the tests then start their own +throwaway server. Without either, those tests are skipped. + +### Trying the controller + +```sh +export OPIFICA_DATABASE_URL=postgres://localhost/opifica +go run ./cmd/controller token create -pool trusted # prints a join token and the CA fingerprint +go run ./cmd/controller serve # TLS on :8443 +go run ./cmd/controller worker list ``` ## Stack diff --git a/cmd/controller/main.go b/cmd/controller/main.go index b11d822..90ba172 100644 --- a/cmd/controller/main.go +++ b/cmd/controller/main.go @@ -2,28 +2,27 @@ // integration, pipeline engine, scheduler and secrets broker. // // Implemented so far: the database, the built-in certificate authority, join -// tokens and Join. Every other RPC returns "unimplemented". +// tokens, Join and the Connect stream (sessions, heartbeats, revocation). +// Scheduling comes next. package main import ( "context" - "crypto/tls" "errors" "flag" "fmt" "io" "log/slog" - "net/http" "os" "os/signal" "strings" "syscall" + "text/tabwriter" "time" - "connectrpc.com/connect" + "github.com/google/uuid" - opificav1 "github.com/opifica/opifica/gen/opifica/v1" - "github.com/opifica/opifica/gen/opifica/v1/opificav1connect" + "github.com/opifica/opifica/internal/controller/api" "github.com/opifica/opifica/internal/controller/ca" "github.com/opifica/opifica/internal/controller/store" "github.com/opifica/opifica/internal/controller/workers" @@ -36,6 +35,8 @@ const usage = `Usage: controller [options] Commands: serve Run the controller token create Create a one-time join token for a new build machine + worker list List workers + worker revoke Revoke a worker; it is disconnected and can't reconnect version Print the version Common options (or environment variables): @@ -57,6 +58,10 @@ func run(args []string, stdout, stderr io.Writer) int { return serve(args[1:], stderr) case args[0] == "token" && len(args) > 1 && args[1] == "create": return tokenCreate(args[2:], stdout, stderr) + case args[0] == "worker" && len(args) > 1 && args[1] == "list": + return workerList(args[2:], stdout, stderr) + case args[0] == "worker" && len(args) > 1 && args[1] == "revoke": + return workerRevoke(args[2:], stdout, stderr) case args[0] == "version": fmt.Fprintln(stdout, version.Version) return 0 @@ -121,42 +126,19 @@ func serve(args []string, stderr io.Writer) int { defer closeFn() svc.Log = log - tlsCfg, err := tlsConfig(svc.CA, splitList(*hosts)) + tlsCfg, err := api.TLSConfig(svc.CA, splitList(*hosts)) if err != nil { log.Error("issuing server certificate", "err", err) return 1 } log.Info("certificate authority ready", "fingerprint", svc.CA.Fingerprint()) - if err := server.RunTLS(ctx, *addr, newHandler(svc), tlsCfg, log); err != nil { + if err := server.RunTLS(ctx, *addr, api.NewHandler(svc), tlsCfg, log); err != nil { log.Error("server stopped", "err", err) return 1 } return 0 } -// tlsConfig issues a fresh server certificate and accepts, but doesn't yet -// require, worker client certificates: Join is called before a worker has -// one. RPCs that need a worker identity check it themselves. -func tlsConfig(authority *ca.CA, hosts []string) (*tls.Config, error) { - certPEM, keyPEM, err := authority.IssueServerCert(hosts) - if err != nil { - return nil, err - } - // Send the CA too, so a joining worker can check it against the pinned - // fingerprint before it has the CA bundle. - certPEM = append(certPEM, authority.CertPEM()...) - pair, err := tls.X509KeyPair(certPEM, keyPEM) - if err != nil { - return nil, err - } - return &tls.Config{ - Certificates: []tls.Certificate{pair}, - ClientCAs: authority.Pool(), - ClientAuth: tls.VerifyClientCertIfGiven, - MinVersion: tls.VersionTLS13, - }, nil -} - func defaultHosts() string { hosts := []string{"localhost", "127.0.0.1", "::1"} if h, err := os.Hostname(); err == nil && h != "" { @@ -175,27 +157,6 @@ func splitList(s string) []string { return out } -func newHandler(svc *workers.Service) http.Handler { - mux := http.NewServeMux() - mux.HandleFunc("GET /healthz", server.Health) - mux.Handle(opificav1connect.NewWorkerServiceHandler(&workerHandler{svc: svc})) - return mux -} - -// workerHandler implements the API workers connect to. -type workerHandler struct { - opificav1connect.UnimplementedWorkerServiceHandler - svc *workers.Service -} - -func (h *workerHandler) Join(ctx context.Context, req *connect.Request[opificav1.JoinRequest]) (*connect.Response[opificav1.JoinResponse], error) { - resp, err := h.svc.Join(ctx, req.Msg) - if err != nil { - return nil, err - } - return connect.NewResponse(resp), nil -} - type labelFlag map[string]string func (l labelFlag) String() string { return "" } @@ -252,3 +213,74 @@ The token is shown only once. `, *pool, tok.ExpiresAt.Format(time.RFC3339), tok.Secret, tok.CAFingerprint) return 0 } + +func workerList(args []string, stdout, stderr io.Writer) int { + fs := flag.NewFlagSet("worker list", flag.ContinueOnError) + fs.SetOutput(stderr) + var c common + c.register(fs) + if err := fs.Parse(args); err != nil { + return 2 + } + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + svc, closeFn, err := c.open(ctx) + if err != nil { + fmt.Fprintf(stderr, "controller: %v\n", err) + return 1 + } + defer closeFn() + ws, err := svc.Store.ListWorkers(ctx) + if err != nil { + fmt.Fprintf(stderr, "controller: %v\n", err) + return 1 + } + tw := tabwriter.NewWriter(stdout, 0, 4, 2, ' ', 0) + fmt.Fprintln(tw, "ID\tHOSTNAME\tOS\tARCH\tPOOL\tJOBS\tLAST SEEN\tSTATUS") + for _, w := range ws { + seen, status := "never", "active" + if w.LastSeenAt != nil { + seen = time.Since(*w.LastSeenAt).Round(time.Second).String() + " ago" + } + if w.RevokedAt != nil { + status = "revoked" + } + fmt.Fprintf(tw, "%s\t%s\t%s\t%s\t%s\t%d\t%s\t%s\n", w.ID, w.Hostname, w.Os, w.Arch, w.Pool, w.MaxConcurrentJobs, seen, status) + } + tw.Flush() + return 0 +} + +func workerRevoke(args []string, stdout, stderr io.Writer) int { + fs := flag.NewFlagSet("worker revoke", flag.ContinueOnError) + fs.SetOutput(stderr) + var c common + c.register(fs) + fs.Usage = func() { fmt.Fprintln(stderr, "Usage: controller worker revoke [options] ") } + if err := fs.Parse(args); err != nil { + return 2 + } + if fs.NArg() != 1 { + fs.Usage() + return 2 + } + id, err := uuid.Parse(fs.Arg(0)) + if err != nil { + fmt.Fprintf(stderr, "controller: invalid worker id %q\n", fs.Arg(0)) + return 2 + } + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + svc, closeFn, err := c.open(ctx) + if err != nil { + fmt.Fprintf(stderr, "controller: %v\n", err) + return 1 + } + defer closeFn() + if err := svc.Revoke(ctx, id); err != nil { + fmt.Fprintf(stderr, "controller: %v\n", err) + return 1 + } + fmt.Fprintf(stdout, "Worker %s revoked. A connected controller disconnects it within %v.\n", id, workers.DefaultTouchInterval) + return 0 +} diff --git a/cmd/controller/main_test.go b/cmd/controller/main_test.go index 9399a12..3648b9d 100644 --- a/cmd/controller/main_test.go +++ b/cmd/controller/main_test.go @@ -22,6 +22,7 @@ import ( opificav1 "github.com/opifica/opifica/gen/opifica/v1" "github.com/opifica/opifica/gen/opifica/v1/opificav1connect" + "github.com/opifica/opifica/internal/controller/api" "github.com/opifica/opifica/internal/server" "github.com/opifica/opifica/internal/testutil/pgtest" "github.com/opifica/opifica/pkg/worker" @@ -56,7 +57,7 @@ func TestTokenCreateAndJoinOverTLS(t *testing.T) { } defer closeFn() svc.Log = slog.New(slog.NewTextHandler(io.Discard, nil)) - tlsCfg, err := tlsConfig(svc.CA, []string{"127.0.0.1"}) + tlsCfg, err := api.TLSConfig(svc.CA, []string{"127.0.0.1"}) if err != nil { t.Fatal(err) } @@ -64,7 +65,7 @@ func TestTokenCreateAndJoinOverTLS(t *testing.T) { if err != nil { t.Fatal(err) } - go server.ServeTLS(ctx, ln, newHandler(svc), tlsCfg, svc.Log) + go server.ServeTLS(ctx, ln, api.NewHandler(svc), tlsCfg, svc.Log) // The worker has no CA yet: it accepts the controller only if the // certificate chains to a CA with the pinned fingerprint. @@ -94,11 +95,11 @@ func TestTokenCreateAndJoinOverTLS(t *testing.T) { t.Errorf("Join response = %+v", resp.Msg) } - // Still unimplemented: the stream comes next. + // Connect without a client certificate is refused. stream := wc.Connect(ctx) stream.Send(&opificav1.ConnectRequest{}) - if _, err := stream.Receive(); connect.CodeOf(err) != connect.CodeUnimplemented { - t.Errorf("Connect: err = %v, want unimplemented", err) + if _, err := stream.Receive(); connect.CodeOf(err) != connect.CodeUnauthenticated { + t.Errorf("Connect without a certificate: err = %v, want unauthenticated", err) } stream.CloseRequest() stream.CloseResponse() diff --git a/docs/ROADMAP.md b/docs/ROADMAP.md index 2c63900..4fdf977 100644 --- a/docs/ROADMAP.md +++ b/docs/ROADMAP.md @@ -29,7 +29,8 @@ result are stored. Tested in ordinary CI, using the fake hypervisor with a real - [x] **Certificate authority:** built-in CA that issues worker client certificates and the server certificate - [x] **Join:** admins create one-time join tokens (`controller token create`); `Join` checks the token and signs the worker's CSR; first contact is verified against the CA fingerprint printed with the token -- [ ] **Connect stream:** Hello, heartbeats, worker registry, disconnect and reconnect handling +- [x] **Connect stream:** Hello, heartbeats, worker registry, disconnect and reconnect handling, revocation + (`controller worker list`, `controller worker revoke`) - [ ] **Leases:** job assignment with a lease renewed by heartbeats; expired leases put the job back in the queue - [ ] **Scheduler:** match planned jobs to workers by OS, architecture, image, labels and pool; respect `needs` - [ ] **Runs:** create a run from a pipeline file and an event; track job and run state @@ -37,12 +38,14 @@ result are stored. Tested in ordinary CI, using the fake hypervisor with a real - [ ] **Secrets:** encrypted store; requests checked against the job token and policy - [ ] **Admin config:** static file for pools, routing (event → pool) and secret policy - [x] **TLS** on the worker API, with worker client certificates accepted -- [ ] **mTLS required** for `Connect`, and certificate renewal +- [x] **mTLS required** for `Connect`; the certificate must be the worker's current one +- [ ] Certificate renewal before the 30-day expiry - [ ] **Admin API + CLI:** `opifica worker token create`, `opifica runs list`, `opifica logs ` ### Worker kit and agent -- [ ] Controller client: join, connect, reconnect with backoff, heartbeats from pool stats +- [x] Controller client: join (saving the identity to disk), connect, reconnect with backoff, heartbeats +- [ ] Heartbeats filled from the warm pool and job runner - [ ] Job runner: lease a VM → checkout → run steps through the agent → stream masked logs → collect reports and artifacts → destroy the VM - [ ] Local Git mirror per repository diff --git a/internal/controller/api/api.go b/internal/controller/api/api.go new file mode 100644 index 0000000..a2dba33 --- /dev/null +++ b/internal/controller/api/api.go @@ -0,0 +1,66 @@ +// Package api assembles the controller's HTTP handler and TLS settings. +package api + +import ( + "context" + "crypto/tls" + "net/http" + + "connectrpc.com/connect" + + opificav1 "github.com/opifica/opifica/gen/opifica/v1" + "github.com/opifica/opifica/gen/opifica/v1/opificav1connect" + "github.com/opifica/opifica/internal/controller/ca" + "github.com/opifica/opifica/internal/controller/workers" + "github.com/opifica/opifica/internal/server" +) + +// NewHandler returns the controller's HTTP handler. +func NewHandler(svc *workers.Service) http.Handler { + mux := http.NewServeMux() + mux.HandleFunc("GET /healthz", server.Health) + mux.Handle(opificav1connect.NewWorkerServiceHandler(&workerHandler{svc: svc})) + return workers.WithPeerCert(mux) +} + +// TLSConfig issues a fresh server certificate for hosts. Worker client +// certificates are verified when presented but not required at the TLS +// layer, because Join is called before a worker has one; Connect requires +// one itself. +func TLSConfig(authority *ca.CA, hosts []string) (*tls.Config, error) { + certPEM, keyPEM, err := authority.IssueServerCert(hosts) + if err != nil { + return nil, err + } + // Send the CA too, so a joining worker can check it against the pinned + // fingerprint before it has the CA bundle. + certPEM = append(certPEM, authority.CertPEM()...) + pair, err := tls.X509KeyPair(certPEM, keyPEM) + if err != nil { + return nil, err + } + return &tls.Config{ + Certificates: []tls.Certificate{pair}, + ClientCAs: authority.Pool(), + ClientAuth: tls.VerifyClientCertIfGiven, + MinVersion: tls.VersionTLS13, + }, nil +} + +// workerHandler implements the API workers connect to. +type workerHandler struct { + opificav1connect.UnimplementedWorkerServiceHandler + svc *workers.Service +} + +func (h *workerHandler) Join(ctx context.Context, req *connect.Request[opificav1.JoinRequest]) (*connect.Response[opificav1.JoinResponse], error) { + resp, err := h.svc.Join(ctx, req.Msg) + if err != nil { + return nil, err + } + return connect.NewResponse(resp), nil +} + +func (h *workerHandler) Connect(ctx context.Context, stream *connect.BidiStream[opificav1.ConnectRequest, opificav1.ConnectResponse]) error { + return h.svc.Connect(ctx, workers.PeerCert(ctx), stream) +} diff --git a/internal/controller/store/db/workers.sql.go b/internal/controller/store/db/workers.sql.go index dfc2409..fed158f 100644 --- a/internal/controller/store/db/workers.sql.go +++ b/internal/controller/store/db/workers.sql.go @@ -267,11 +267,39 @@ func (q *Queries) SetJoinTokenWorker(ctx context.Context, arg SetJoinTokenWorker return err } -const touchWorker = `-- name: TouchWorker :exec -UPDATE workers SET last_seen_at = now() WHERE id = $1 +const touchWorker = `-- name: TouchWorker :one +UPDATE workers SET last_seen_at = now() WHERE id = $1 RETURNING revoked_at ` -func (q *Queries) TouchWorker(ctx context.Context, id uuid.UUID) error { - _, err := q.db.Exec(ctx, touchWorker, id) +// Records that the worker is alive and returns its revocation time, so a +// connected worker that has been revoked can be disconnected. +func (q *Queries) TouchWorker(ctx context.Context, id uuid.UUID) (*time.Time, error) { + row := q.db.QueryRow(ctx, touchWorker, id) + var revoked_at *time.Time + err := row.Scan(&revoked_at) + return revoked_at, err +} + +const updateWorkerHello = `-- name: UpdateWorkerHello :exec +UPDATE workers +SET hostname = $2, version = $3, max_concurrent_jobs = $4, last_seen_at = now() +WHERE id = $1 +` + +type UpdateWorkerHelloParams struct { + ID uuid.UUID + Hostname string + Version string + MaxConcurrentJobs int32 +} + +// Facts a worker reports each time it connects. +func (q *Queries) UpdateWorkerHello(ctx context.Context, arg UpdateWorkerHelloParams) error { + _, err := q.db.Exec(ctx, updateWorkerHello, + arg.ID, + arg.Hostname, + arg.Version, + arg.MaxConcurrentJobs, + ) return err } diff --git a/internal/controller/store/queries/workers.sql b/internal/controller/store/queries/workers.sql index 5419fbf..83a4031 100644 --- a/internal/controller/store/queries/workers.sql +++ b/internal/controller/store/queries/workers.sql @@ -30,8 +30,16 @@ SELECT * FROM workers WHERE id = $1; -- name: ListWorkers :many SELECT * FROM workers ORDER BY created_at; --- name: TouchWorker :exec -UPDATE workers SET last_seen_at = now() WHERE id = $1; +-- name: TouchWorker :one +-- Records that the worker is alive and returns its revocation time, so a +-- connected worker that has been revoked can be disconnected. +UPDATE workers SET last_seen_at = now() WHERE id = $1 RETURNING revoked_at; + +-- name: UpdateWorkerHello :exec +-- Facts a worker reports each time it connects. +UPDATE workers +SET hostname = $2, version = $3, max_concurrent_jobs = $4, last_seen_at = now() +WHERE id = $1; -- name: RevokeWorker :one UPDATE workers SET revoked_at = now() WHERE id = $1 AND revoked_at IS NULL RETURNING *; diff --git a/internal/controller/workers/connect.go b/internal/controller/workers/connect.go new file mode 100644 index 0000000..99151d1 --- /dev/null +++ b/internal/controller/workers/connect.go @@ -0,0 +1,300 @@ +package workers + +import ( + "context" + "crypto/x509" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "sync" + "time" + + "connectrpc.com/connect" + "github.com/google/uuid" + + opificav1 "github.com/opifica/opifica/gen/opifica/v1" + "github.com/opifica/opifica/internal/controller/ca" + "github.com/opifica/opifica/internal/controller/store" + "github.com/opifica/opifica/internal/controller/store/db" +) + +// Defaults for ConnectConfig. +const ( + DefaultHelloTimeout = 10 * time.Second + DefaultHeartbeatTimeout = 30 * time.Second + DefaultTouchInterval = 30 * time.Second +) + +// ConnectConfig tunes the Connect stream. Zero values use the defaults. +type ConnectConfig struct { + // HelloTimeout is how long a new stream may take to send Hello. + HelloTimeout time.Duration + // HeartbeatTimeout ends a session that has sent nothing for this long. + // Workers send heartbeats every 10s by default. + HeartbeatTimeout time.Duration + // TouchInterval is how often last_seen is written to the database and + // revocation is checked. + TouchInterval time.Duration +} + +func (c ConnectConfig) withDefaults() ConnectConfig { + if c.HelloTimeout <= 0 { + c.HelloTimeout = DefaultHelloTimeout + } + if c.HeartbeatTimeout <= 0 { + c.HeartbeatTimeout = DefaultHeartbeatTimeout + } + if c.TouchInterval <= 0 { + c.TouchInterval = DefaultTouchInterval + } + return c +} + +// Stream is the server side of a Connect stream. +type Stream interface { + Receive() (*opificav1.ConnectRequest, error) + Send(*opificav1.ConnectResponse) error +} + +// MessageHandler receives what workers send about jobs. It is nil until the +// scheduler exists; until then those messages are logged and dropped. +type MessageHandler interface { + HandleWorkerMessage(ctx context.Context, worker uuid.UUID, msg *opificav1.ConnectRequest) error +} + +type peerCertKey struct{} + +// WithPeerCert stores the verified TLS client certificate, if any, in the +// request context so RPC handlers can identify the worker. +func WithPeerCert(h http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.TLS != nil && len(r.TLS.VerifiedChains) > 0 { + r = r.WithContext(context.WithValue(r.Context(), peerCertKey{}, r.TLS.VerifiedChains[0][0])) + } + h.ServeHTTP(w, r) + }) +} + +// PeerCert returns the verified client certificate stored by WithPeerCert. +func PeerCert(ctx context.Context) *x509.Certificate { + c, _ := ctx.Value(peerCertKey{}).(*x509.Certificate) + return c +} + +var ( + errNoCert = connect.NewError(connect.CodeUnauthenticated, errors.New("a worker client certificate is required")) + errRevoked = connect.NewError(connect.CodePermissionDenied, errors.New("worker has been revoked")) + errUnknown = connect.NewError(connect.CodeUnauthenticated, errors.New("unknown worker")) + errSuperseded = connect.NewError(connect.CodeUnauthenticated, errors.New("certificate has been superseded")) + + errRevokedCause = errors.New("revoked") +) + +// Connect runs one worker's stream until it ends. The worker is identified +// by its client certificate, never by what it says in Hello. +func (s *Service) Connect(ctx context.Context, cert *x509.Certificate, stream Stream) error { + cfg := s.StreamConfig.withDefaults() + if cert == nil { + return errNoCert + } + idStr, err := ca.WorkerID(cert) + if err != nil { + return errNoCert + } + id, err := uuid.Parse(idStr) + if err != nil { + return errNoCert + } + w, err := s.Store.GetWorker(ctx, id) + switch { + case store.IsNotFound(err): + return errUnknown + case err != nil: + s.log().Error("loading worker", "worker", id, "err", err) + return connect.NewError(connect.CodeUnavailable, errors.New("try again later")) + case w.RevokedAt != nil: + return errRevoked + case w.CertSerial != cert.SerialNumber.Text(16): + return errSuperseded + } + + // Deferred first so it runs last: the writer must be done with the stream + // before the handler returns. + var writer sync.WaitGroup + defer writer.Wait() + ctx, cancel := context.WithCancelCause(ctx) + defer cancel(nil) + + // Receive in the background so timeouts and cancellation can interrupt + // waiting for the next message. + msgs := make(chan *opificav1.ConnectRequest) + recvErr := make(chan error, 1) + go func() { + for { + m, err := stream.Receive() + if err != nil { + recvErr <- err + return + } + select { + case msgs <- m: + case <-ctx.Done(): + return + } + } + }() + + hello, err := s.awaitHello(ctx, msgs, recvErr, cfg.HelloTimeout) + if err != nil { + return err + } + if hello.GetWorkerId() != id.String() { + return connect.NewError(connect.CodeInvalidArgument, + fmt.Errorf("hello names worker %q but the certificate is for %s", hello.GetWorkerId(), id)) + } + info := hello.GetInfo() + maxJobs := int(info.GetMaxConcurrentJobs()) + if maxJobs < 1 || maxJobs > maxJobsPerHost { + return connect.NewError(connect.CodeInvalidArgument, + fmt.Errorf("max_concurrent_jobs must be between 1 and %d", maxJobsPerHost)) + } + if err := s.Store.UpdateWorkerHello(ctx, db.UpdateWorkerHelloParams{ + ID: id, Hostname: info.GetHostname(), Version: info.GetVersion(), MaxConcurrentJobs: int32(maxJobs), + }); err != nil { + s.log().Error("recording hello", "worker", id, "err", err) + return connect.NewError(connect.CodeUnavailable, errors.New("try again later")) + } + + var labels map[string]string + if err := json.Unmarshal(w.Labels, &labels); err != nil { + return connect.NewError(connect.CodeInternal, errors.New("stored worker labels are invalid")) + } + now := s.now() + sess := &Session{ + WorkerID: id, Pool: w.Pool, OS: w.Os, Arch: w.Arch, Labels: labels, MaxJobs: maxJobs, + ConnectedAt: now, lastSeen: now, running: hello.GetRunningJobIds(), + out: make(chan *opificav1.ConnectResponse, 16), ctx: ctx, cancel: cancel, + } + s.Sessions().register(sess) + defer s.Sessions().unregister(sess) + log := s.log().With("worker", id, "hostname", info.GetHostname()) + log.Info("worker connected", "pool", w.Pool, "os", w.Os, "arch", w.Arch, "max_jobs", maxJobs) + + // Writer: the only goroutine that sends on the stream. + writer.Add(1) + go func() { + defer writer.Done() + for { + select { + case m := <-sess.out: + if err := stream.Send(m); err != nil { + cancel(err) + return + } + case <-ctx.Done(): + return + } + } + }() + + idle := time.NewTimer(cfg.HeartbeatTimeout) + defer idle.Stop() + touch := time.NewTicker(cfg.TouchInterval) + defer touch.Stop() + + for { + select { + case m := <-msgs: + idle.Reset(cfg.HeartbeatTimeout) + sess.seen(s.now()) + if err := s.handle(ctx, sess, m); err != nil { + log.Info("worker disconnected", "reason", err) + return err + } + case err := <-recvErr: + if errors.Is(err, io.EOF) { + log.Info("worker disconnected", "reason", "closed by worker") + return nil + } + log.Info("worker disconnected", "reason", err) + return err + case <-idle.C: + log.Warn("worker disconnected", "reason", "no heartbeat") + return connect.NewError(connect.CodeDeadlineExceeded, + fmt.Errorf("no message for %v", cfg.HeartbeatTimeout)) + case <-touch.C: + revoked, err := s.Store.TouchWorker(ctx, id) + if err != nil && ctx.Err() == nil { + log.Error("recording heartbeat", "err", err) + continue + } + if revoked != nil { + log.Warn("worker disconnected", "reason", "revoked") + return errRevoked + } + case <-ctx.Done(): + cause := context.Cause(ctx) + log.Info("worker disconnected", "reason", cause) + if errors.Is(cause, errReplaced) { + return connect.NewError(connect.CodeAborted, errReplaced) + } + if errors.Is(cause, errRevokedCause) { + return errRevoked + } + return connect.NewError(connect.CodeCanceled, cause) + } + } +} + +func (s *Service) awaitHello(ctx context.Context, msgs <-chan *opificav1.ConnectRequest, recvErr <-chan error, timeout time.Duration) (*opificav1.Hello, error) { + t := time.NewTimer(timeout) + defer t.Stop() + select { + case m := <-msgs: + h := m.GetHello() + if h == nil { + return nil, connect.NewError(connect.CodeInvalidArgument, errors.New("the first message must be hello")) + } + return h, nil + case err := <-recvErr: + return nil, err + case <-t.C: + return nil, connect.NewError(connect.CodeDeadlineExceeded, errors.New("no hello received")) + case <-ctx.Done(): + return nil, context.Cause(ctx) + } +} + +func (s *Service) handle(ctx context.Context, sess *Session, m *opificav1.ConnectRequest) error { + switch msg := m.GetMessage().(type) { + case *opificav1.ConnectRequest_Heartbeat: + sess.heartbeat(msg.Heartbeat, s.now()) + return nil + case *opificav1.ConnectRequest_Hello: + return connect.NewError(connect.CodeInvalidArgument, errors.New("hello may only be sent once")) + case *opificav1.ConnectRequest_JobUpdate, *opificav1.ConnectRequest_LogChunk, *opificav1.ConnectRequest_SecretRequest: + if s.Messages == nil { + s.log().Debug("dropping job message: no scheduler yet", "worker", sess.WorkerID) + return nil + } + return s.Messages.HandleWorkerMessage(ctx, sess.WorkerID, m) + case nil: + return connect.NewError(connect.CodeInvalidArgument, errors.New("empty message")) + } + return connect.NewError(connect.CodeInvalidArgument, fmt.Errorf("unknown message %T", m.GetMessage())) +} + +// Revoke marks a worker as revoked and disconnects it if this controller +// holds its session. Other controllers disconnect it at their next touch. +func (s *Service) Revoke(ctx context.Context, id uuid.UUID) error { + if _, err := s.Store.RevokeWorker(ctx, id); err != nil { + if store.IsNotFound(err) { + return fmt.Errorf("worker %s not found or already revoked", id) + } + return err + } + s.Sessions().Disconnect(id, errRevokedCause) + return nil +} diff --git a/internal/controller/workers/join.go b/internal/controller/workers/join.go index 256b5ed..06b60b7 100644 --- a/internal/controller/workers/join.go +++ b/internal/controller/workers/join.go @@ -14,6 +14,7 @@ import ( "maps" "regexp" "strings" + "sync" "time" "connectrpc.com/connect" @@ -49,6 +50,20 @@ type Service struct { CA *ca.CA Log *slog.Logger Now func() time.Time + + // StreamConfig tunes the Connect stream. + StreamConfig ConnectConfig + // Messages receives job messages from workers; nil drops them. + Messages MessageHandler + + sessionsOnce sync.Once + sessions *Registry +} + +// Sessions returns the registry of connected workers. +func (s *Service) Sessions() *Registry { + s.sessionsOnce.Do(func() { s.sessions = NewRegistry() }) + return s.sessions } func (s *Service) now() time.Time { diff --git a/internal/controller/workers/registry.go b/internal/controller/workers/registry.go new file mode 100644 index 0000000..f41863c --- /dev/null +++ b/internal/controller/workers/registry.go @@ -0,0 +1,169 @@ +package workers + +import ( + "context" + "errors" + "maps" + "slices" + "sync" + "time" + + "github.com/google/uuid" + + opificav1 "github.com/opifica/opifica/gen/opifica/v1" +) + +// ErrNotConnected is returned when sending to a worker without a live session. +var ErrNotConnected = errors.New("worker is not connected") + +// Session is one live Connect stream from a worker. +type Session struct { + WorkerID uuid.UUID + Pool string + OS string + Arch string + Labels map[string]string + MaxJobs int + ConnectedAt time.Time + + out chan *opificav1.ConnectResponse + ctx context.Context + cancel context.CancelCauseFunc + + mu sync.Mutex + lastSeen time.Time + warmVMs uint32 + running []string +} + +// Send queues a message for the worker. It fails if the session ends first. +func (s *Session) Send(ctx context.Context, msg *opificav1.ConnectResponse) error { + select { + case s.out <- msg: + return nil + case <-s.ctx.Done(): + return ErrNotConnected + case <-ctx.Done(): + return ctx.Err() + } +} + +func (s *Session) heartbeat(hb *opificav1.Heartbeat, now time.Time) { + s.mu.Lock() + defer s.mu.Unlock() + s.lastSeen = now + s.warmVMs = hb.GetWarmVms() + s.running = slices.Clone(hb.GetRunningJobIds()) +} + +func (s *Session) seen(now time.Time) { + s.mu.Lock() + s.lastSeen = now + s.mu.Unlock() +} + +// SessionInfo is a snapshot of a session, for the scheduler and admin views. +type SessionInfo struct { + WorkerID uuid.UUID + Pool string + OS string + Arch string + Labels map[string]string + MaxJobs int + ConnectedAt time.Time + LastSeen time.Time + WarmVMs uint32 + RunningJobs []string +} + +func (s *Session) info() SessionInfo { + s.mu.Lock() + defer s.mu.Unlock() + return SessionInfo{ + WorkerID: s.WorkerID, Pool: s.Pool, OS: s.OS, Arch: s.Arch, + Labels: maps.Clone(s.Labels), MaxJobs: s.MaxJobs, ConnectedAt: s.ConnectedAt, + LastSeen: s.lastSeen, WarmVMs: s.warmVMs, RunningJobs: slices.Clone(s.running), + } +} + +// Registry holds the live session of every connected worker. +type Registry struct { + mu sync.Mutex + sessions map[uuid.UUID]*Session +} + +// NewRegistry returns an empty registry. +func NewRegistry() *Registry { + return &Registry{sessions: map[uuid.UUID]*Session{}} +} + +// errReplaced ends a session when the same worker connects again, e.g. after +// a network change left the old stream half-open. +var errReplaced = errors.New("replaced by a newer connection from the same worker") + +func (r *Registry) register(s *Session) { + r.mu.Lock() + old := r.sessions[s.WorkerID] + r.sessions[s.WorkerID] = s + r.mu.Unlock() + if old != nil { + old.cancel(errReplaced) + } +} + +func (r *Registry) unregister(s *Session) { + r.mu.Lock() + defer r.mu.Unlock() + if r.sessions[s.WorkerID] == s { + delete(r.sessions, s.WorkerID) + } +} + +// Get returns the live session for a worker. +func (r *Registry) Get(id uuid.UUID) (*Session, bool) { + r.mu.Lock() + defer r.mu.Unlock() + s, ok := r.sessions[id] + return s, ok +} + +// Send queues a message for a connected worker. +func (r *Registry) Send(ctx context.Context, id uuid.UUID, msg *opificav1.ConnectResponse) error { + s, ok := r.Get(id) + if !ok { + return ErrNotConnected + } + return s.Send(ctx, msg) +} + +// Disconnect ends a worker's session, if it has one. +func (r *Registry) Disconnect(id uuid.UUID, cause error) { + if s, ok := r.Get(id); ok { + s.cancel(cause) + } +} + +// List returns a snapshot of every connected worker, ordered by ID. +func (r *Registry) List() []SessionInfo { + r.mu.Lock() + ss := make([]*Session, 0, len(r.sessions)) + for _, s := range r.sessions { + ss = append(ss, s) + } + r.mu.Unlock() + out := make([]SessionInfo, len(ss)) + for i, s := range ss { + out[i] = s.info() + } + slices.SortFunc(out, func(a, b SessionInfo) int { return compareUUID(a.WorkerID, b.WorkerID) }) + return out +} + +func compareUUID(a, b uuid.UUID) int { + for i := range a { + if a[i] != b[i] { + return int(a[i]) - int(b[i]) + } + } + return 0 +} diff --git a/internal/controller/workers/registry_test.go b/internal/controller/workers/registry_test.go new file mode 100644 index 0000000..6fd45f5 --- /dev/null +++ b/internal/controller/workers/registry_test.go @@ -0,0 +1,72 @@ +package workers + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/google/uuid" + + opificav1 "github.com/opifica/opifica/gen/opifica/v1" +) + +func newSession(id uuid.UUID) *Session { + ctx, cancel := context.WithCancelCause(context.Background()) + return &Session{WorkerID: id, out: make(chan *opificav1.ConnectResponse, 1), ctx: ctx, cancel: cancel} +} + +func TestRegistryReplacesOldSession(t *testing.T) { + r := NewRegistry() + id := uuid.New() + first, second := newSession(id), newSession(id) + + r.register(first) + r.register(second) + if !errors.Is(context.Cause(first.ctx), errReplaced) { + t.Errorf("old session cause = %v, want errReplaced", context.Cause(first.ctx)) + } + if second.ctx.Err() != nil { + t.Error("new session was cancelled") + } + + // The old session's handler unregisters on its way out; that must not + // remove the new session. + r.unregister(first) + if s, ok := r.Get(id); !ok || s != second { + t.Error("unregistering the old session removed the new one") + } + r.unregister(second) + if _, ok := r.Get(id); ok { + t.Error("session still registered") + } +} + +func TestSessionSend(t *testing.T) { + r := NewRegistry() + id := uuid.New() + s := newSession(id) + r.register(s) + + msg := &opificav1.ConnectResponse{} + if err := r.Send(context.Background(), id, msg); err != nil { + t.Fatal(err) + } + if got := <-s.out; got != msg { + t.Error("wrong message queued") + } + + // Queue full and the caller gives up. + s.out <- msg + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancel() + if err := r.Send(ctx, id, msg); !errors.Is(err, context.DeadlineExceeded) { + t.Errorf("send to full queue: err = %v", err) + } + + // Session ended while the sender waits. + r.Disconnect(id, errors.New("gone")) + if err := s.Send(context.Background(), msg); !errors.Is(err, ErrNotConnected) { + t.Errorf("send to ended session: err = %v", err) + } +} diff --git a/internal/integration/connect_test.go b/internal/integration/connect_test.go new file mode 100644 index 0000000..aa8060f --- /dev/null +++ b/internal/integration/connect_test.go @@ -0,0 +1,361 @@ +// Package integration runs the controller and the worker kit against each +// other over real TLS connections and a real PostgreSQL database. +package integration + +import ( + "context" + "errors" + "io" + "log/slog" + "net" + "net/http" + "os" + "sync" + "testing" + "time" + + "connectrpc.com/connect" + "github.com/google/uuid" + + opificav1 "github.com/opifica/opifica/gen/opifica/v1" + "github.com/opifica/opifica/gen/opifica/v1/opificav1connect" + "github.com/opifica/opifica/internal/controller/api" + "github.com/opifica/opifica/internal/controller/ca" + "github.com/opifica/opifica/internal/controller/store" + "github.com/opifica/opifica/internal/controller/workers" + "github.com/opifica/opifica/internal/server" + "github.com/opifica/opifica/internal/testutil/pgtest" + "github.com/opifica/opifica/pkg/worker" +) + +func TestMain(m *testing.M) { os.Exit(pgtest.Run(m)) } + +func quiet() *slog.Logger { return slog.New(slog.NewTextHandler(io.Discard, nil)) } + +type env struct { + t *testing.T + svc *workers.Service + url string +} + +// newEnv starts a controller with short timeouts on a random port. +func newEnv(t *testing.T) *env { + t.Helper() + ctx, cancel := context.WithCancel(context.Background()) + st, err := store.Open(ctx, pgtest.New(t)) + if err != nil { + t.Fatal(err) + } + authority, err := ca.LoadOrCreate(t.TempDir()) + if err != nil { + t.Fatal(err) + } + svc := &workers.Service{ + Store: st, CA: authority, Log: quiet(), + StreamConfig: workers.ConnectConfig{ + HelloTimeout: 2 * time.Second, + HeartbeatTimeout: 500 * time.Millisecond, + TouchInterval: 50 * time.Millisecond, + }, + } + tlsCfg, err := api.TLSConfig(authority, []string{"127.0.0.1"}) + if err != nil { + t.Fatal(err) + } + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + done := make(chan struct{}) + go func() { + server.ServeTLS(ctx, ln, api.NewHandler(svc), tlsCfg, quiet()) + close(done) + }() + t.Cleanup(func() { + cancel() + <-done + st.Close() + }) + return &env{t: t, svc: svc, url: "https://" + ln.Addr().String()} +} + +func info() *opificav1.WorkerInfo { + return &opificav1.WorkerInfo{ + Hostname: "build-1", Os: opificav1.OperatingSystem_OPERATING_SYSTEM_LINUX, + Arch: opificav1.Architecture_ARCHITECTURE_AMD64, MaxConcurrentJobs: 4, Version: "test", + Labels: map[string]string{"kvm": "true"}, + } +} + +// join creates a token for pool and joins a worker with it. +func (e *env) join(pool string) *worker.Identity { + e.t.Helper() + tok, err := e.svc.CreateToken(context.Background(), workers.TokenRequest{Pool: pool}) + if err != nil { + e.t.Fatal(err) + } + id, err := worker.Join(context.Background(), worker.JoinConfig{ + ControllerURL: e.url, Token: tok.Secret, CAFingerprint: tok.CAFingerprint, + StateDir: e.t.TempDir(), Info: info(), + }) + if err != nil { + e.t.Fatalf("Join: %v", err) + } + return id +} + +type runningClient struct { + *worker.Client + cancel context.CancelFunc + done chan error + msgs chan *opificav1.ConnectResponse +} + +func (e *env) start(id *worker.Identity, mutate func(*worker.ClientConfig)) *runningClient { + e.t.Helper() + rc := &runningClient{done: make(chan error, 1), msgs: make(chan *opificav1.ConnectResponse, 16)} + cfg := worker.ClientConfig{ + ControllerURL: e.url, Identity: id, Info: info(), + HeartbeatInterval: 50 * time.Millisecond, MaxBackoff: 100 * time.Millisecond, Logger: quiet(), + Status: func() (uint32, []string) { return 3, []string{"job-1"} }, + OnMessage: func(_ context.Context, m *opificav1.ConnectResponse) { rc.msgs <- m }, + } + if mutate != nil { + mutate(&cfg) + } + c, err := worker.NewClient(cfg) + if err != nil { + e.t.Fatal(err) + } + rc.Client = c + ctx, cancel := context.WithCancel(context.Background()) + rc.cancel = cancel + go func() { rc.done <- c.Run(ctx) }() + e.t.Cleanup(func() { + cancel() + <-rc.done + }) + return rc +} + +func workerHTTPClient(id *worker.Identity) *http.Client { + return &http.Client{Transport: &http.Transport{TLSClientConfig: id.TLSConfig(), ForceAttemptHTTP2: true}} +} + +func eventually(t *testing.T, what string, cond func() bool) { + t.Helper() + deadline := time.Now().Add(10 * time.Second) + for !cond() { + if time.Now().After(deadline) { + t.Fatalf("timed out waiting for %s", what) + } + time.Sleep(5 * time.Millisecond) + } +} + +func (e *env) session(id *worker.Identity) (workers.SessionInfo, bool) { + for _, s := range e.svc.Sessions().List() { + if s.WorkerID.String() == id.WorkerID { + return s, true + } + } + return workers.SessionInfo{}, false +} + +func TestWorkerConnectsAndReportsStatus(t *testing.T) { + e := newEnv(t) + id := e.join("trusted") + e.start(id, nil) + + eventually(t, "a session with heartbeat status", func() bool { + s, ok := e.session(id) + return ok && s.WarmVMs == 3 && len(s.RunningJobs) == 1 + }) + s, _ := e.session(id) + if s.Pool != "trusted" || s.OS != "linux" || s.Arch != "amd64" || s.MaxJobs != 4 || s.Labels["kvm"] != "true" { + t.Errorf("session = %+v", s) + } + + eventually(t, "last_seen recorded", func() bool { + w, err := e.svc.Store.GetWorker(context.Background(), uuid.MustParse(id.WorkerID)) + return err == nil && w.LastSeenAt != nil && w.Version == "test" + }) +} + +func TestControllerSendsToWorker(t *testing.T) { + e := newEnv(t) + id := e.join("trusted") + rc := e.start(id, nil) + eventually(t, "session", func() bool { _, ok := e.session(id); return ok }) + + err := e.svc.Sessions().Send(context.Background(), uuid.MustParse(id.WorkerID), &opificav1.ConnectResponse{ + Message: &opificav1.ConnectResponse_Drain{Drain: &opificav1.Drain{Reason: "maintenance"}}, + }) + if err != nil { + t.Fatal(err) + } + select { + case m := <-rc.msgs: + if m.GetDrain().GetReason() != "maintenance" { + t.Errorf("worker got %v", m) + } + case <-time.After(5 * time.Second): + t.Fatal("worker never received the message") + } + + if err := e.svc.Sessions().Send(context.Background(), uuid.New(), &opificav1.ConnectResponse{}); !errors.Is(err, workers.ErrNotConnected) { + t.Errorf("send to unknown worker: err = %v", err) + } +} + +func TestWorkerReconnects(t *testing.T) { + e := newEnv(t) + id := e.join("trusted") + e.start(id, nil) + eventually(t, "session", func() bool { _, ok := e.session(id); return ok }) + first, _ := e.session(id) + + e.svc.Sessions().Disconnect(first.WorkerID, errors.New("network blip")) + eventually(t, "a new session", func() bool { + s, ok := e.session(id) + return ok && s.ConnectedAt.After(first.ConnectedAt) + }) +} + +func TestSilentWorkerIsDropped(t *testing.T) { + e := newEnv(t) + id := e.join("trusted") + // Heartbeats far apart: the controller gives up after 500ms of silence. + e.start(id, func(c *worker.ClientConfig) { c.HeartbeatInterval = time.Hour }) + eventually(t, "session", func() bool { _, ok := e.session(id); return ok }) + first, _ := e.session(id) + + // The worker reconnects, is dropped again, and so on: each time a new session. + eventually(t, "the silent session to be replaced", func() bool { + s, ok := e.session(id) + return ok && s.ConnectedAt.After(first.ConnectedAt) + }) +} + +func TestRevokedWorkerIsDisconnectedAndStaysOut(t *testing.T) { + e := newEnv(t) + id := e.join("trusted") + rc := e.start(id, nil) + eventually(t, "session", func() bool { _, ok := e.session(id); return ok }) + + if err := e.svc.Revoke(context.Background(), uuid.MustParse(id.WorkerID)); err != nil { + t.Fatal(err) + } + select { + case err := <-rc.done: + if connect.CodeOf(err) != connect.CodePermissionDenied || !worker.IsPermanent(err) { + t.Errorf("Run returned %v, want permission denied", err) + } + rc.done <- err // for the cleanup + case <-time.After(5 * time.Second): + t.Fatal("revoked worker is still running") + } + if _, ok := e.session(id); ok { + t.Error("revoked worker still has a session") + } + + // Revoked on another controller (directly in the database): the + // periodic touch finds it. + id2 := e.join("trusted") + rc2 := e.start(id2, nil) + eventually(t, "second session", func() bool { _, ok := e.session(id2); return ok }) + if _, err := e.svc.Store.RevokeWorker(context.Background(), uuid.MustParse(id2.WorkerID)); err != nil { + t.Fatal(err) + } + select { + case err := <-rc2.done: + if !worker.IsPermanent(err) { + t.Errorf("Run returned %v, want a permanent error", err) + } + rc2.done <- err + case <-time.After(5 * time.Second): + t.Fatal("worker revoked in the database is still connected") + } +} + +func TestConnectRejects(t *testing.T) { + e := newEnv(t) + id := e.join("trusted") + ctx := context.Background() + + connectWith := func(t *testing.T, hello *opificav1.Hello) error { + t.Helper() + client := opificav1connect.NewWorkerServiceClient(workerHTTPClient(id), e.url) + stream := client.Connect(ctx) + defer stream.CloseResponse() + stream.Send(&opificav1.ConnectRequest{Message: &opificav1.ConnectRequest_Hello{Hello: hello}}) + _, err := stream.Receive() + return err + } + + t.Run("hello for another worker", func(t *testing.T) { + err := connectWith(t, &opificav1.Hello{WorkerId: uuid.NewString(), Info: info()}) + if connect.CodeOf(err) != connect.CodeInvalidArgument { + t.Errorf("err = %v, want invalid argument", err) + } + }) + t.Run("bad capacity", func(t *testing.T) { + bad := info() + bad.MaxConcurrentJobs = 0 + err := connectWith(t, &opificav1.Hello{WorkerId: id.WorkerID, Info: bad}) + if connect.CodeOf(err) != connect.CodeInvalidArgument { + t.Errorf("err = %v, want invalid argument", err) + } + }) + t.Run("superseded certificate", func(t *testing.T) { + if _, err := e.svc.Store.Pool.Exec(ctx, "UPDATE workers SET cert_serial = 'other' WHERE id = $1", id.WorkerID); err != nil { + t.Fatal(err) + } + err := connectWith(t, &opificav1.Hello{WorkerId: id.WorkerID, Info: info()}) + if connect.CodeOf(err) != connect.CodeUnauthenticated { + t.Errorf("err = %v, want unauthenticated", err) + } + }) +} + +// A worker whose certificate is from a different CA can't connect. +func TestForeignCertificateRejected(t *testing.T) { + e := newEnv(t) + other := newEnv(t) + foreign := other.join("trusted") + + c, err := worker.NewClient(worker.ClientConfig{ + ControllerURL: e.url, Identity: foreign, Info: info(), Logger: quiet(), MaxBackoff: 50 * time.Millisecond, + }) + if err != nil { + t.Fatal(err) + } + // Neither side trusts the other, so the TLS handshake fails on every + // attempt and Run keeps retrying until the deadline. + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if err := c.Run(ctx); !errors.Is(err, context.DeadlineExceeded) { + t.Errorf("Run returned %v, want it to keep retrying until the deadline", err) + } + if n := len(e.svc.Sessions().List()); n != 0 { + t.Errorf("%d sessions, want 0", n) + } +} + +func TestManyWorkers(t *testing.T) { + e := newEnv(t) + var wg sync.WaitGroup + ids := make([]*worker.Identity, 10) + for i := range ids { + wg.Add(1) + go func() { + defer wg.Done() + ids[i] = e.join("untrusted") + }() + } + wg.Wait() + for _, id := range ids { + e.start(id, nil) + } + eventually(t, "10 sessions", func() bool { return len(e.svc.Sessions().List()) == 10 }) +} diff --git a/pkg/worker/client.go b/pkg/worker/client.go new file mode 100644 index 0000000..c3c09b8 --- /dev/null +++ b/pkg/worker/client.go @@ -0,0 +1,343 @@ +package worker + +import ( + "context" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/tls" + "crypto/x509" + "encoding/pem" + "errors" + "fmt" + "log/slog" + "math" + mrand "math/rand/v2" + "net/http" + "net/url" + "os" + "path/filepath" + "strings" + "sync" + "time" + + "connectrpc.com/connect" + + opificav1 "github.com/opifica/opifica/gen/opifica/v1" + "github.com/opifica/opifica/gen/opifica/v1/opificav1connect" +) + +// Files in the worker's state directory. +const ( + idFile = "worker.id" + keyFile = "worker.key" + certFile = "worker.crt" + caFile = "ca.crt" +) + +// Identity is what a worker keeps after joining: its ID, private key, +// certificate and the controller's CA. +type Identity struct { + WorkerID string + cert tls.Certificate + caPool *x509.CertPool +} + +// JoinConfig is what a worker needs to join a controller. +type JoinConfig struct { + // ControllerURL is e.g. "https://controller.example:8443". + ControllerURL string + // Token and CAFingerprint are printed by `controller token create`. + Token string + CAFingerprint string + // StateDir is where the identity is stored. It must be private to the + // worker: it holds the worker's private key. + StateDir string + Info *opificav1.WorkerInfo +} + +// Join exchanges a join token for a worker identity and saves it in +// cfg.StateDir. The private key is generated here and never leaves the +// machine; only a certificate request is sent. +func Join(ctx context.Context, cfg JoinConfig) (*Identity, error) { + if _, err := os.Stat(filepath.Join(cfg.StateDir, idFile)); err == nil { + return nil, fmt.Errorf("%s already holds a worker identity; remove it to join again", cfg.StateDir) + } + u, err := url.Parse(cfg.ControllerURL) + if err != nil || u.Scheme != "https" || u.Hostname() == "" { + return nil, fmt.Errorf("controller URL must be https://host[:port], got %q", cfg.ControllerURL) + } + tlsCfg, err := PinnedCAConfig(cfg.CAFingerprint, u.Hostname()) + if err != nil { + return nil, err + } + + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + return nil, err + } + csrDER, err := x509.CreateCertificateRequest(rand.Reader, &x509.CertificateRequest{}, key) + if err != nil { + return nil, err + } + client := opificav1connect.NewWorkerServiceClient(httpClient(tlsCfg), cfg.ControllerURL) + resp, err := client.Join(ctx, connect.NewRequest(&opificav1.JoinRequest{ + JoinToken: cfg.Token, + CsrPem: pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE REQUEST", Bytes: csrDER}), + Info: cfg.Info, + })) + if err != nil { + return nil, fmt.Errorf("joining: %w", err) + } + + keyDER, err := x509.MarshalPKCS8PrivateKey(key) + if err != nil { + return nil, err + } + keyPEM := pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: keyDER}) + if err := os.MkdirAll(cfg.StateDir, 0o700); err != nil { + return nil, err + } + // Write the ID last: its presence marks a complete identity. + for _, f := range []struct { + name string + data []byte + perm os.FileMode + }{ + {keyFile, keyPEM, 0o600}, + {certFile, resp.Msg.GetCertificatePem(), 0o644}, + {caFile, resp.Msg.GetCaBundlePem(), 0o644}, + {idFile, []byte(resp.Msg.GetWorkerId() + "\n"), 0o644}, + } { + if err := os.WriteFile(filepath.Join(cfg.StateDir, f.name), f.data, f.perm); err != nil { + return nil, err + } + } + return LoadIdentity(cfg.StateDir) +} + +// LoadIdentity reads an identity saved by Join. +func LoadIdentity(dir string) (*Identity, error) { + id, err := os.ReadFile(filepath.Join(dir, idFile)) + if err != nil { + return nil, fmt.Errorf("no worker identity in %s (join first): %w", dir, err) + } + cert, err := tls.LoadX509KeyPair(filepath.Join(dir, certFile), filepath.Join(dir, keyFile)) + if err != nil { + return nil, fmt.Errorf("loading worker certificate: %w", err) + } + caPEM, err := os.ReadFile(filepath.Join(dir, caFile)) + if err != nil { + return nil, err + } + pool := x509.NewCertPool() + if !pool.AppendCertsFromPEM(caPEM) { + return nil, errors.New("ca.crt holds no certificates") + } + return &Identity{WorkerID: strings.TrimSpace(string(id)), cert: cert, caPool: pool}, nil +} + +// TLSConfig returns the mTLS config for talking to the controller. +func (id *Identity) TLSConfig() *tls.Config { + return &tls.Config{ + Certificates: []tls.Certificate{id.cert}, + RootCAs: id.caPool, + MinVersion: tls.VersionTLS13, + } +} + +func httpClient(cfg *tls.Config) *http.Client { + return &http.Client{Transport: &http.Transport{ + TLSClientConfig: cfg, + ForceAttemptHTTP2: true, + IdleConnTimeout: 90 * time.Second, + }} +} + +// ClientConfig configures a Client. +type ClientConfig struct { + ControllerURL string + Identity *Identity + Info *opificav1.WorkerInfo + + // HeartbeatInterval defaults to 10s; the controller drops workers it + // hasn't heard from for 30s. + HeartbeatInterval time.Duration + // Status reports what goes into heartbeats. Optional. + Status func() (warmVMs uint32, runningJobs []string) + // OnMessage receives every message from the controller, in order, on + // one goroutine. It should return quickly. + OnMessage func(ctx context.Context, msg *opificav1.ConnectResponse) + // MaxBackoff caps the wait between reconnect attempts. Defaults to 1m. + MaxBackoff time.Duration + Logger *slog.Logger +} + +// Client keeps a worker connected to the controller. +type Client struct { + cfg ClientConfig + rpc opificav1connect.WorkerServiceClient + log *slog.Logger + + mu sync.Mutex + stream *connect.BidiStreamForClient[opificav1.ConnectRequest, opificav1.ConnectResponse] + ready chan struct{} // closed while connected +} + +// NewClient creates a client. Call Run to connect. +func NewClient(cfg ClientConfig) (*Client, error) { + if cfg.Identity == nil { + return nil, errors.New("identity is required") + } + if cfg.Info == nil { + return nil, errors.New("info is required") + } + if cfg.HeartbeatInterval <= 0 { + cfg.HeartbeatInterval = 10 * time.Second + } + if cfg.MaxBackoff <= 0 { + cfg.MaxBackoff = time.Minute + } + log := cfg.Logger + if log == nil { + log = slog.Default() + } + return &Client{ + cfg: cfg, + rpc: opificav1connect.NewWorkerServiceClient(httpClient(cfg.Identity.TLSConfig()), cfg.ControllerURL), + log: log.With("worker", cfg.Identity.WorkerID), + ready: make(chan struct{}), + }, nil +} + +// ErrNotConnected is returned by Send while the client is between connections. +var ErrNotConnected = errors.New("not connected to the controller") + +// Send sends a message on the current stream. +func (c *Client) Send(msg *opificav1.ConnectRequest) error { + c.mu.Lock() + defer c.mu.Unlock() + if c.stream == nil { + return ErrNotConnected + } + return c.stream.Send(msg) +} + +// Connected returns a channel that is closed while the client is connected. +// After a disconnect a new channel is returned. +func (c *Client) Connected() <-chan struct{} { + c.mu.Lock() + defer c.mu.Unlock() + return c.ready +} + +// IsPermanent reports whether a Run error means retrying is pointless: the +// worker was revoked or its identity is no longer accepted. +func IsPermanent(err error) bool { + switch connect.CodeOf(err) { + case connect.CodePermissionDenied, connect.CodeUnauthenticated: + return true + } + return false +} + +// Run connects and stays connected, reconnecting with backoff, until ctx is +// cancelled or the controller rejects the worker permanently. +func (c *Client) Run(ctx context.Context) error { + attempt := 0 + for { + start := time.Now() + err := c.session(ctx) + if ctx.Err() != nil { + return ctx.Err() + } + if IsPermanent(err) { + c.log.Error("controller rejected this worker", "err", err) + return err + } + if time.Since(start) > c.cfg.MaxBackoff { + attempt = 0 // the last connection was healthy + } + wait := backoff(attempt, c.cfg.MaxBackoff) + attempt++ + c.log.Warn("disconnected from controller", "err", err, "retry_in", wait.Round(time.Millisecond)) + select { + case <-time.After(wait): + case <-ctx.Done(): + return ctx.Err() + } + } +} + +// backoff is exponential from 500ms with ±25% jitter, capped at max. +func backoff(attempt int, max time.Duration) time.Duration { + d := time.Duration(float64(500*time.Millisecond) * math.Pow(2, float64(min(attempt, 20)))) + d = min(d, max) + return time.Duration(float64(d) * (0.75 + 0.5*mrand.Float64())) +} + +func (c *Client) status() (uint32, []string) { + if c.cfg.Status == nil { + return 0, nil + } + return c.cfg.Status() +} + +func (c *Client) session(ctx context.Context) error { + ctx, cancel := context.WithCancel(ctx) + defer cancel() + stream := c.rpc.Connect(ctx) + + _, running := c.status() + if err := stream.Send(&opificav1.ConnectRequest{Message: &opificav1.ConnectRequest_Hello{Hello: &opificav1.Hello{ + WorkerId: c.cfg.Identity.WorkerID, Info: c.cfg.Info, RunningJobIds: running, + }}}); err != nil { + // The real error, e.g. a rejected certificate, comes from Receive. + _, rerr := stream.Receive() + return errors.Join(err, rerr) + } + + c.mu.Lock() + c.stream = stream + close(c.ready) + c.mu.Unlock() + defer func() { + c.mu.Lock() + c.stream = nil + c.ready = make(chan struct{}) + c.mu.Unlock() + stream.CloseRequest() + stream.CloseResponse() + }() + + go func() { + t := time.NewTicker(c.cfg.HeartbeatInterval) + defer t.Stop() + for { + select { + case <-t.C: + warm, running := c.status() + err := c.Send(&opificav1.ConnectRequest{Message: &opificav1.ConnectRequest_Heartbeat{Heartbeat: &opificav1.Heartbeat{ + WarmVms: warm, RunningJobIds: running, + }}}) + if err != nil { + cancel() + return + } + case <-ctx.Done(): + return + } + } + }() + + c.log.Info("connected to controller") + for { + msg, err := stream.Receive() + if err != nil { + return err + } + if c.cfg.OnMessage != nil { + c.cfg.OnMessage(ctx, msg) + } + } +} From cd54d2f417419477324ca55683f5937435770f0e Mon Sep 17 00:00:00 2001 From: "K.Milchev" <49519924+KristiforMilchev@users.noreply.github.com> Date: Tue, 3 Nov 2026 14:38:39 +0200 Subject: [PATCH 2/2] Scheduler: runs, jobs, leases, retries and cancellation - Runs and jobs in PostgreSQL. A run is planned from the pipeline file and the event; jobs wait for their needs, are skipped when a need fails, and the run closes when nothing can change any more - Jobs go to workers matching pool, OS, architecture, image and labels, up to each worker's capacity. Claiming is one statement with FOR UPDATE SKIP LOCKED, so several controllers never assign a job twice - Leases renewed by heartbeats; expired leases requeue the job, up to 3 attempts. Jobs a disconnected worker never started go back at once without using an attempt. Updates from a worker that no longer holds a job are ignored and answered with a cancel - Cancel from any process: the serving controller stops the jobs on the workers it holds - controller run create|list|show|cancel; serve runs the scheduler - Worker client: stopping now ends the stream at once (Receive ignored context cancellation and waited ~400ms for the controller), and HTTP/2 pings detect a controller that vanished - Integration tests with fake workers built on the real client, checked by breaking the queries on purpose Co-Authored-By: Claude Opus 5.5 --- README.md | 14 +- cmd/controller/main.go | 23 +- cmd/controller/main_test.go | 36 + cmd/controller/runs.go | 235 ++++++ docs/ROADMAP.md | 13 +- internal/controller/scheduler/scheduler.go | 554 ++++++++++++ internal/controller/store/db/models.go | 44 + internal/controller/store/db/runs.sql.go | 799 ++++++++++++++++++ .../controller/store/migrations/0002_runs.sql | 68 ++ internal/controller/store/queries/runs.sql | 163 ++++ internal/controller/workers/connect.go | 32 +- internal/controller/workers/registry.go | 21 +- internal/integration/client_stop_test.go | 28 + internal/integration/connect_test.go | 30 +- internal/integration/scheduler_test.go | 620 ++++++++++++++ internal/server/server.go | 3 + pkg/worker/client.go | 18 +- 17 files changed, 2667 insertions(+), 34 deletions(-) create mode 100644 cmd/controller/runs.go create mode 100644 internal/controller/scheduler/scheduler.go create mode 100644 internal/controller/store/db/runs.sql.go create mode 100644 internal/controller/store/migrations/0002_runs.sql create mode 100644 internal/controller/store/queries/runs.sql create mode 100644 internal/integration/client_stop_test.go create mode 100644 internal/integration/scheduler_test.go diff --git a/README.md b/README.md index 65a6439..e9217a4 100644 --- a/README.md +++ b/README.md @@ -3,9 +3,9 @@ Core of [Opifica](https://github.com/opifica), an open-source CI/CD system built in the EU where every job runs in a fresh virtual machine. -> **Status:** early development. Working and tested: the pipeline language, the worker kit's warm pool, and -> workers joining and staying connected to the controller. Scheduling and the agent come next; see -> [docs/ROADMAP.md](docs/ROADMAP.md). +> **Status:** early development. Working and tested: the pipeline language, the worker kit's warm pool, +> workers joining and staying connected to the controller, and scheduling runs onto them. The job runner and +> agent come next; see [docs/ROADMAP.md](docs/ROADMAP.md). > See the [organisation page](https://github.com/opifica) for the architecture. ## What lives here @@ -19,7 +19,8 @@ where every job runs in a fresh virtual machine. | `pkg/worker` | Shared worker kit: `Hypervisor` interface, warm VM pool, image verification, controller client | Working | | `pkg/worker/workertest` | Fake hypervisor for testing worker code without VMs | Working | | `cmd/controller` | Join tokens, worker joining over TLS, the Connect stream, worker revocation | Working | -| | Scheduler, runs, logs, secrets, Git providers | Not started | +| `internal/controller/scheduler` | Runs and jobs, matching jobs to workers, leases, retries, cancellation | Working | +| | Logs, secrets, Git providers, admin API | Not started | | `cmd/agent` | Runs inside each VM, executes steps, streams logs and results | Skeleton | | `cmd/signer` | Signing service for release artifacts | Not started | @@ -129,6 +130,11 @@ export OPIFICA_DATABASE_URL=postgres://localhost/opifica go run ./cmd/controller token create -pool trusted # prints a join token and the CA fingerprint go run ./cmd/controller serve # TLS on :8443 go run ./cmd/controller worker list + +# Until Git providers exist, start runs by hand: +go run ./cmd/controller run create -repo acme/app -pool trusted -sha $(git rev-parse HEAD) -file .opifica/pipeline.yaml +go run ./cmd/controller run list +go run ./cmd/controller run show ``` ## Stack diff --git a/cmd/controller/main.go b/cmd/controller/main.go index 90ba172..d845eeb 100644 --- a/cmd/controller/main.go +++ b/cmd/controller/main.go @@ -2,8 +2,8 @@ // integration, pipeline engine, scheduler and secrets broker. // // Implemented so far: the database, the built-in certificate authority, join -// tokens, Join and the Connect stream (sessions, heartbeats, revocation). -// Scheduling comes next. +// tokens, Join, the Connect stream, and the scheduler (runs, jobs, leases, +// retries, cancellation). Logs, secrets and Git providers come next. package main import ( @@ -24,6 +24,7 @@ import ( "github.com/opifica/opifica/internal/controller/api" "github.com/opifica/opifica/internal/controller/ca" + "github.com/opifica/opifica/internal/controller/scheduler" "github.com/opifica/opifica/internal/controller/store" "github.com/opifica/opifica/internal/controller/workers" "github.com/opifica/opifica/internal/server" @@ -37,6 +38,10 @@ Commands: token create Create a one-time join token for a new build machine worker list List workers worker revoke Revoke a worker; it is disconnected and can't reconnect + run create Start a run from a pipeline file (until Git providers exist) + run list List recent runs + run show Show a run and its jobs + run cancel Cancel a run version Print the version Common options (or environment variables): @@ -62,6 +67,8 @@ func run(args []string, stdout, stderr io.Writer) int { return workerList(args[2:], stdout, stderr) case args[0] == "worker" && len(args) > 1 && args[1] == "revoke": return workerRevoke(args[2:], stdout, stderr) + case args[0] == "run" && len(args) > 1: + return runCommand(args[1], args[2:], stdout, stderr) case args[0] == "version": fmt.Fprintln(stdout, version.Version) return 0 @@ -132,6 +139,9 @@ func serve(args []string, stderr io.Writer) int { return 1 } log.Info("certificate authority ready", "fingerprint", svc.CA.Fingerprint()) + sched := scheduler.New(svc.Store, svc.Sessions(), log.With("component", "scheduler"), scheduler.Config{}) + svc.Messages = sched + go sched.Run(ctx) if err := server.RunTLS(ctx, *addr, api.NewHandler(svc), tlsCfg, log); err != nil { log.Error("server stopped", "err", err) return 1 @@ -257,16 +267,17 @@ func workerRevoke(args []string, stdout, stderr io.Writer) int { var c common c.register(fs) fs.Usage = func() { fmt.Fprintln(stderr, "Usage: controller worker revoke [options] ") } - if err := fs.Parse(args); err != nil { + pos, err := parseWithArgs(fs, args) + if err != nil { return 2 } - if fs.NArg() != 1 { + if len(pos) != 1 { fs.Usage() return 2 } - id, err := uuid.Parse(fs.Arg(0)) + id, err := uuid.Parse(pos[0]) if err != nil { - fmt.Fprintf(stderr, "controller: invalid worker id %q\n", fs.Arg(0)) + fmt.Fprintf(stderr, "controller: invalid worker id %q\n", pos[0]) return 2 } ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) diff --git a/cmd/controller/main_test.go b/cmd/controller/main_test.go index 3648b9d..4fee3ce 100644 --- a/cmd/controller/main_test.go +++ b/cmd/controller/main_test.go @@ -140,3 +140,39 @@ func TestUsage(t *testing.T) { t.Errorf("token create without -pool: exit %d", code) } } + +func TestRunCommands(t *testing.T) { + dbURL := pgtest.New(t) + dir := t.TempDir() + file := filepath.Join(dir, "pipeline.yaml") + os.WriteFile(file, []byte("on: push\njobs:\n build:\n runs-on: linux\n steps: [{ run: make }]\n deploy:\n needs: build\n runs-on: linux\n steps: [{ run: make deploy }]\n"), 0o644) + common := []string{"-database-url", dbURL, "-data-dir", dir} + + var out, errOut bytes.Buffer + if code := run(append([]string{"run", "create", "-repo", "acme/app", "-pool", "trusted", "-sha", "abc123", "-file", file}, common...), &out, &errOut); code != 0 { + t.Fatalf("run create: exit %d: %s", code, errOut.String()) + } + id := regexp.MustCompile(`Run (\S+) created with 2 job\(s\)`).FindStringSubmatch(out.String()) + if id == nil { + t.Fatalf("run create output: %s", out.String()) + } + + out.Reset() + if code := run(append([]string{"run", "list"}, common...), &out, &errOut); code != 0 || !regexp.MustCompile(`acme/app\s+push\s+main\s+abc123\s+trusted\s+running`).MatchString(out.String()) { + t.Errorf("run list (exit %d):\n%s", code, out.String()) + } + out.Reset() + if code := run(append([]string{"run", "show", id[1]}, common...), &out, &errOut); code != 0 || + !regexp.MustCompile(`build\s+linux\s+queued`).MatchString(out.String()) || !regexp.MustCompile(`deploy\s+linux\s+pending`).MatchString(out.String()) { + t.Errorf("run show (exit %d):\n%s", code, out.String()) + } + out.Reset() + if code := run(append([]string{"run", "cancel", id[1]}, common...), &out, &errOut); code != 0 { + t.Errorf("run cancel: exit %d: %s", code, errOut.String()) + } + out.Reset() + run(append([]string{"run", "show", id[1]}, common...), &out, &errOut) + if !regexp.MustCompile(`: cancelled`).MatchString(out.String()) { + t.Errorf("after cancel:\n%s", out.String()) + } +} diff --git a/cmd/controller/runs.go b/cmd/controller/runs.go new file mode 100644 index 0000000..058a245 --- /dev/null +++ b/cmd/controller/runs.go @@ -0,0 +1,235 @@ +package main + +import ( + "context" + "flag" + "fmt" + "io" + "log/slog" + "os" + "strings" + "text/tabwriter" + "time" + + "github.com/google/uuid" + + "github.com/opifica/opifica/internal/controller/scheduler" + "github.com/opifica/opifica/internal/controller/store/db" + "github.com/opifica/opifica/internal/pipeline" +) + +func runCommand(sub string, args []string, stdout, stderr io.Writer) int { + switch sub { + case "create": + return runCreate(args, stdout, stderr) + case "list": + return runList(args, stdout, stderr) + case "show": + return runShow(args, stdout, stderr) + case "cancel": + return runCancel(args, stdout, stderr) + } + fmt.Fprintf(stderr, "controller: unknown command \"run %s\"\n\n%s", sub, usage) + return 2 +} + +// withScheduler opens the store and builds a scheduler without running its +// loop: the serving controller picks up what this process writes. +func withScheduler(c common, stderr io.Writer, fn func(ctx context.Context, sched *scheduler.Scheduler, q *db.Queries) error) int { + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + svc, closeFn, err := c.open(ctx) + if err != nil { + fmt.Fprintf(stderr, "controller: %v\n", err) + return 1 + } + defer closeFn() + quiet := slog.New(slog.NewTextHandler(io.Discard, nil)) + sched := scheduler.New(svc.Store, svc.Sessions(), quiet, scheduler.Config{}) + if err := fn(ctx, sched, svc.Store.Queries); err != nil { + fmt.Fprintf(stderr, "controller: %v\n", err) + return 1 + } + return 0 +} + +func runCreate(args []string, stdout, stderr io.Writer) int { + fs := flag.NewFlagSet("run create", flag.ContinueOnError) + fs.SetOutput(stderr) + var c common + c.register(fs) + repo := fs.String("repo", "", "repository name, e.g. acme/app (required)") + cloneURL := fs.String("clone-url", "", "URL the worker clones from") + file := fs.String("file", pipeline.DefaultPath, "pipeline file") + pool := fs.String("pool", "", "pool to run in (required); later decided by routing policy") + event := fs.String("event", "push", "push, pull_request or tag") + branch := fs.String("branch", "main", "branch (push) or source branch (pull request)") + base := fs.String("base", "main", "pull request target branch") + tag := fs.String("tag", "", "tag name, for -event tag") + sha := fs.String("sha", "", "commit SHA (required)") + if err := fs.Parse(args); err != nil { + return 2 + } + if *repo == "" || *pool == "" || *sha == "" { + fmt.Fprintln(stderr, "controller: -repo, -pool and -sha are required") + return 2 + } + src, err := os.ReadFile(*file) + if err != nil { + fmt.Fprintf(stderr, "controller: %v\n", err) + return 1 + } + ev := pipeline.Event{Type: pipeline.EventType(*event), SHA: *sha} + switch ev.Type { + case pipeline.EventPush: + ev.Branch, ev.Ref = *branch, "refs/heads/"+*branch + case pipeline.EventPullRequest: + ev.Branch, ev.BaseBranch = *branch, *base + case pipeline.EventTag: + ev.Tag, ev.Ref = *tag, "refs/tags/"+*tag + default: + fmt.Fprintf(stderr, "controller: unknown event %q\n", *event) + return 2 + } + return withScheduler(c, stderr, func(ctx context.Context, sched *scheduler.Scheduler, q *db.Queries) error { + run, err := sched.CreateRun(ctx, scheduler.RunRequest{ + Repo: *repo, CloneURL: *cloneURL, Event: ev, Pool: *pool, Pipeline: src, + }) + if err != nil { + return err + } + jobs, err := q.ListRunJobs(ctx, run.ID) + if err != nil { + return err + } + fmt.Fprintf(stdout, "Run %s created with %d job(s).\n", run.ID, len(jobs)) + return nil + }) +} + +func runList(args []string, stdout, stderr io.Writer) int { + fs := flag.NewFlagSet("run list", flag.ContinueOnError) + fs.SetOutput(stderr) + var c common + c.register(fs) + limit := fs.Int("n", 20, "number of runs to show") + if err := fs.Parse(args); err != nil { + return 2 + } + return withScheduler(c, stderr, func(ctx context.Context, _ *scheduler.Scheduler, q *db.Queries) error { + runs, err := q.ListRuns(ctx, int32(*limit)) + if err != nil { + return err + } + tw := tabwriter.NewWriter(stdout, 0, 4, 2, ' ', 0) + fmt.Fprintln(tw, "ID\tREPO\tEVENT\tREF\tCOMMIT\tPOOL\tSTATE\tCREATED") + for _, r := range runs { + fmt.Fprintf(tw, "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n", r.ID, r.Repo, r.EventType, refName(r), + short(r.CommitSha), r.Pool, r.State, r.CreatedAt.Local().Format(time.DateTime)) + } + return tw.Flush() + }) +} + +func runShow(args []string, stdout, stderr io.Writer) int { + fs := flag.NewFlagSet("run show", flag.ContinueOnError) + fs.SetOutput(stderr) + var c common + c.register(fs) + id, ok := runIDArg(fs, args, stderr, "show") + if !ok { + return 2 + } + return withScheduler(c, stderr, func(ctx context.Context, _ *scheduler.Scheduler, q *db.Queries) error { + r, err := q.GetRun(ctx, id) + if err != nil { + return fmt.Errorf("run %s: %w", id, err) + } + fmt.Fprintf(stdout, "Run %s\n %s %s @ %s in pool %s: %s\n\n", r.ID, r.Repo, refName(r), short(r.CommitSha), r.Pool, r.State) + jobs, err := q.ListRunJobs(ctx, id) + if err != nil { + return err + } + tw := tabwriter.NewWriter(stdout, 0, 4, 2, ' ', 0) + fmt.Fprintln(tw, "JOB\tOS\tSTATE\tATTEMPT\tWORKER\tEXIT\tNOTE") + for _, j := range jobs { + worker, exit := "-", "-" + if j.WorkerID != nil { + worker = j.WorkerID.String()[:8] + } + if j.ExitCode != nil { + exit = fmt.Sprint(*j.ExitCode) + } + fmt.Fprintf(tw, "%s\t%s\t%s\t%d/%d\t%s\t%s\t%s\n", j.Name, j.Os, j.State, j.Attempt, j.MaxAttempts, worker, exit, j.Error) + } + return tw.Flush() + }) +} + +func runCancel(args []string, stdout, stderr io.Writer) int { + fs := flag.NewFlagSet("run cancel", flag.ContinueOnError) + fs.SetOutput(stderr) + var c common + c.register(fs) + id, ok := runIDArg(fs, args, stderr, "cancel") + if !ok { + return 2 + } + return withScheduler(c, stderr, func(ctx context.Context, sched *scheduler.Scheduler, _ *db.Queries) error { + if err := sched.CancelRun(ctx, id); err != nil { + return err + } + fmt.Fprintf(stdout, "Run %s is being cancelled; running jobs are stopped on their workers.\n", id) + return nil + }) +} + +func runIDArg(fs *flag.FlagSet, args []string, stderr io.Writer, cmd string) (uuid.UUID, bool) { + pos, err := parseWithArgs(fs, args) + if err != nil { + return uuid.UUID{}, false + } + if len(pos) != 1 { + fmt.Fprintf(stderr, "Usage: controller run %s [options] \n", cmd) + return uuid.UUID{}, false + } + id, err := uuid.Parse(pos[0]) + if err != nil { + fmt.Fprintf(stderr, "controller: invalid run id %q\n", pos[0]) + return uuid.UUID{}, false + } + return id, true +} + +// parseWithArgs parses flags that may come before or after positional +// arguments, e.g. both "show -data-dir x ID" and "show ID -data-dir x". +func parseWithArgs(fs *flag.FlagSet, args []string) ([]string, error) { + var pos []string + for { + if err := fs.Parse(args); err != nil { + return nil, err + } + if fs.NArg() == 0 { + return pos, nil + } + pos = append(pos, fs.Arg(0)) + args = fs.Args()[1:] + } +} + +func refName(r db.Run) string { + switch { + case r.Tag != "": + return r.Tag + case r.EventType == "pull_request": + return r.Branch + " -> " + r.BaseBranch + } + return strings.TrimPrefix(r.Branch, "refs/heads/") +} + +func short(sha string) string { + if len(sha) > 12 { + return sha[:12] + } + return sha +} diff --git a/docs/ROADMAP.md b/docs/ROADMAP.md index 4fdf977..242f0b0 100644 --- a/docs/ROADMAP.md +++ b/docs/ROADMAP.md @@ -25,22 +25,25 @@ result are stored. Tested in ordinary CI, using the fake hypervisor with a real ### Controller - [x] **Store:** PostgreSQL migrations and sqlc queries, for workers and join tokens so far -- [ ] **Store:** runs and jobs +- [x] **Store:** runs and jobs - [x] **Certificate authority:** built-in CA that issues worker client certificates and the server certificate - [x] **Join:** admins create one-time join tokens (`controller token create`); `Join` checks the token and signs the worker's CSR; first contact is verified against the CA fingerprint printed with the token - [x] **Connect stream:** Hello, heartbeats, worker registry, disconnect and reconnect handling, revocation (`controller worker list`, `controller worker revoke`) -- [ ] **Leases:** job assignment with a lease renewed by heartbeats; expired leases put the job back in the queue -- [ ] **Scheduler:** match planned jobs to workers by OS, architecture, image, labels and pool; respect `needs` -- [ ] **Runs:** create a run from a pipeline file and an event; track job and run state +- [x] **Leases:** job assignment with a lease renewed by heartbeats; expired leases put the job back in the queue + (up to 3 attempts); jobs a disconnected worker never started go back at once +- [x] **Scheduler:** match planned jobs to workers by OS, architecture, image, labels and pool; respect `needs`; + safe with several controllers (`FOR UPDATE SKIP LOCKED`) +- [x] **Runs:** create a run from a pipeline file and an event; track job and run state; cancel - [ ] **Logs:** receive log chunks, drop duplicates after reconnects, store them on disk - [ ] **Secrets:** encrypted store; requests checked against the job token and policy - [ ] **Admin config:** static file for pools, routing (event → pool) and secret policy - [x] **TLS** on the worker API, with worker client certificates accepted - [x] **mTLS required** for `Connect`; the certificate must be the worker's current one - [ ] Certificate renewal before the 30-day expiry -- [ ] **Admin API + CLI:** `opifica worker token create`, `opifica runs list`, `opifica logs ` +- [x] **Admin CLI (direct database access):** `controller token create`, `worker list|revoke`, `run create|list|show|cancel` +- [ ] **Admin API** so the `opifica` CLI and web UI can do the same remotely; `opifica logs ` ### Worker kit and agent diff --git a/internal/controller/scheduler/scheduler.go b/internal/controller/scheduler/scheduler.go new file mode 100644 index 0000000..d31a619 --- /dev/null +++ b/internal/controller/scheduler/scheduler.go @@ -0,0 +1,554 @@ +// Package scheduler turns pipelines into runs and jobs, hands queued jobs to +// matching workers, and follows each job to its end. +// +// All state lives in PostgreSQL. Claiming a job is one statement using +// SELECT ... FOR UPDATE SKIP LOCKED, so several controllers can schedule at +// once without giving a job to two workers. +// +// A worker holds each job it was given under a lease, renewed by its +// heartbeats. If the lease expires (the worker crashed or lost its network), +// the job is queued again, up to its maximum number of attempts. +package scheduler + +import ( + "context" + "crypto/rand" + "crypto/sha256" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "log/slog" + mrand "math/rand/v2" + "time" + + "github.com/google/uuid" + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/types/known/durationpb" + + opificav1 "github.com/opifica/opifica/gen/opifica/v1" + "github.com/opifica/opifica/internal/controller/store" + "github.com/opifica/opifica/internal/controller/store/db" + "github.com/opifica/opifica/internal/controller/workers" + "github.com/opifica/opifica/internal/pipeline" +) + +// Defaults for Config. +const ( + DefaultLease = 60 * time.Second + DefaultInterval = 2 * time.Second +) + +// Config tunes the scheduler. Zero values use the defaults. +type Config struct { + // Lease is how long a worker holds a job without renewing it. Workers + // renew with every heartbeat (every 10s by default). + Lease time.Duration + // Interval is how often the scheduler looks for work and expired leases + // even when nothing woke it. + Interval time.Duration +} + +// Scheduler implements workers.MessageHandler. +type Scheduler struct { + store *store.Store + sessions *workers.Registry + log *slog.Logger + cfg Config + kick chan struct{} +} + +var _ workers.MessageHandler = (*Scheduler)(nil) + +// New creates a scheduler. Call Run to start it. +func New(st *store.Store, sessions *workers.Registry, log *slog.Logger, cfg Config) *Scheduler { + if cfg.Lease <= 0 { + cfg.Lease = DefaultLease + } + if cfg.Interval <= 0 { + cfg.Interval = DefaultInterval + } + if log == nil { + log = slog.Default() + } + return &Scheduler{store: st, sessions: sessions, log: log, cfg: cfg, kick: make(chan struct{}, 1)} +} + +// Kick makes the scheduler look for work now. +func (s *Scheduler) Kick() { + select { + case s.kick <- struct{}{}: + default: + } +} + +// Run schedules until ctx is cancelled. +func (s *Scheduler) Run(ctx context.Context) error { + t := time.NewTicker(s.cfg.Interval) + defer t.Stop() + for { + s.reapLeases(ctx) + s.propagateCancels(ctx) + s.assign(ctx) + select { + case <-ctx.Done(): + return ctx.Err() + case <-t.C: + case <-s.kick: + } + } +} + +// ---------------------------------------------------------------- runs + +// RunRequest asks for a new run. +type RunRequest struct { + Repo string + CloneURL string + Event pipeline.Event + // Pool comes from the controller's routing policy, never from the + // pipeline file. + Pool string + // Pipeline is the pipeline file as it is at Event.SHA. + Pipeline []byte +} + +// ErrNotTriggered means the event doesn't match the pipeline's triggers. +var ErrNotTriggered = errors.New("the event does not trigger this pipeline") + +// ErrNothingToRun means every job was skipped by its condition. +var ErrNothingToRun = errors.New("no job's condition matched; nothing to run") + +// CreateRun plans the pipeline for the event and stores the run and its jobs. +func (s *Scheduler) CreateRun(ctx context.Context, req RunRequest) (db.Run, error) { + if req.Repo == "" || req.Pool == "" || req.Event.SHA == "" { + return db.Run{}, errors.New("repo, pool and commit SHA are required") + } + p, err := pipeline.Parse(req.Pipeline) + if err != nil { + return db.Run{}, err + } + if !p.Matches(req.Event) { + return db.Run{}, ErrNotTriggered + } + plan, err := p.Plan(req.Event) + if err != nil { + return db.Run{}, err + } + if len(plan.Jobs) == 0 { + return db.Run{}, ErrNothingToRun + } + + runID := uuid.New() + ev := req.Event + var run db.Run + err = s.store.Tx(ctx, func(q *db.Queries) error { + var err error + run, err = q.CreateRun(ctx, db.CreateRunParams{ + ID: runID, Repo: req.Repo, CloneUrl: req.CloneURL, EventType: string(ev.Type), Ref: ev.Ref, + CommitSha: ev.SHA, Branch: ev.Branch, BaseBranch: ev.BaseBranch, Tag: ev.Tag, + Pool: req.Pool, PipelineSource: string(req.Pipeline), + }) + if err != nil { + return err + } + for _, pj := range plan.Jobs { + jobID := uuid.New() + spec, err := proto.Marshal(jobSpec(jobID, runID, pj, req.CloneURL, ev)) + if err != nil { + return err + } + labels := pj.RunsOn.Labels + if labels == nil { + labels = map[string]string{} + } + labelsJSON, err := json.Marshal(labels) + if err != nil { + return err + } + state := "queued" + if len(pj.Needs) > 0 { + state = "pending" + } + if err := q.CreateJob(ctx, db.CreateJobParams{ + ID: jobID, RunID: runID, Key: pj.ID, Name: pj.Name, Pool: req.Pool, + Os: pj.RunsOn.OS, Arch: pj.RunsOn.Arch, Image: pj.RunsOn.Image, Labels: labelsJSON, + Needs: nonNil(pj.Needs), Spec: spec, State: state, + }); err != nil { + return err + } + } + return nil + }) + if err != nil { + return db.Run{}, fmt.Errorf("storing run: %w", err) + } + s.log.Info("run created", "run", runID, "repo", req.Repo, "sha", ev.SHA, "pool", req.Pool, "jobs", len(plan.Jobs)) + s.Kick() + return run, nil +} + +func nonNil(s []string) []string { + if s == nil { + return []string{} + } + return s +} + +// jobSpec is what the worker receives, minus the per-attempt job token. +func jobSpec(jobID, runID uuid.UUID, pj *pipeline.PlannedJob, cloneURL string, ev pipeline.Event) *opificav1.Job { + j := &opificav1.Job{ + Id: jobID.String(), + RunId: runID.String(), + Name: pj.Name, + Source: &opificav1.Source{CloneUrl: cloneURL, CommitSha: ev.SHA, Ref: ev.Ref}, + Image: pj.RunsOn.Image, + Env: pj.Env, + AllowedSecrets: pj.Secrets, + Timeout: durationpb.New(pj.Timeout), + } + for _, st := range pj.Steps { + step := &opificav1.Step{ + Name: st.Name, Env: st.Env, Secrets: st.Secrets, ContinueOnError: st.ContinueOnError, + } + if st.Timeout > 0 { + step.Timeout = durationpb.New(st.Timeout) + } + switch { + case st.Run != nil: + step.Action = &opificav1.Step_Run{Run: &opificav1.RunAction{ + Command: st.Run.Command, Shell: st.Run.Shell, WorkingDirectory: st.Run.WorkingDirectory, + }} + case st.Report != nil: + step.Action = &opificav1.Step_Report{Report: &opificav1.ReportAction{ + Format: opificav1.ReportFormat_REPORT_FORMAT_JUNIT, Paths: st.Report.Paths, + }} + case st.Upload != nil: + step.Action = &opificav1.Step_Upload{Upload: &opificav1.UploadAction{Name: st.Upload.Name, Paths: st.Upload.Paths}} + case st.Sign != nil: + step.Action = &opificav1.Step_Sign{Sign: &opificav1.SignAction{Paths: st.Sign.Paths, Profile: st.Sign.Profile}} + } + j.Steps = append(j.Steps, step) + } + return j +} + +// CancelRun stops a run: waiting jobs are cancelled at once, and workers +// running its jobs are told to stop. +func (s *Scheduler) CancelRun(ctx context.Context, runID uuid.UUID) error { + if _, err := s.store.RequestRunCancel(ctx, runID); err != nil { + if store.IsNotFound(err) { + return fmt.Errorf("run %s not found or already finished", runID) + } + return err + } + if err := s.store.CancelWaitingJobs(ctx, runID); err != nil { + return err + } + active, err := s.store.ActiveRunJobs(ctx, runID) + if err != nil { + return err + } + for _, j := range active { + if j.WorkerID != nil { + s.cancelOnWorker(ctx, *j.WorkerID, j.ID, "run cancelled") + } + } + s.advance(ctx, runID) + s.log.Info("run cancel requested", "run", runID, "running_jobs", len(active)) + return nil +} + +// advance moves a run forward after a job finished: jobs that can no longer +// run are skipped, jobs whose needs are met are queued, and the run is +// closed when nothing is left. +func (s *Scheduler) advance(ctx context.Context, runID uuid.UUID) { + for { + skipped, err := s.store.SkipBlockedJobs(ctx, runID) + if err != nil { + s.log.Error("skipping blocked jobs", "run", runID, "err", err) + return + } + if len(skipped) == 0 { + break + } + } + ready, err := s.store.QueueReadyJobs(ctx, runID) + if err != nil { + s.log.Error("queueing ready jobs", "run", runID, "err", err) + return + } + if len(ready) > 0 { + s.Kick() + } + run, err := s.store.FinishRunIfDone(ctx, runID) + switch { + case store.IsNotFound(err): + case err != nil: + s.log.Error("finishing run", "run", runID, "err", err) + default: + s.log.Info("run finished", "run", runID, "state", run.State) + } +} + +// ---------------------------------------------------------------- assignment + +func (s *Scheduler) assign(ctx context.Context) { + sessions := s.sessions.List() + mrand.Shuffle(len(sessions), func(i, j int) { sessions[i], sessions[j] = sessions[j], sessions[i] }) + for _, w := range sessions { + if len(w.Images) == 0 { + continue + } + active, err := s.store.CountWorkerJobs(ctx, &w.WorkerID) + if err != nil { + s.log.Error("counting worker jobs", "worker", w.WorkerID, "err", err) + return + } + for free := w.MaxJobs - int(active); free > 0; free-- { + if !s.assignOne(ctx, w) { + break + } + } + } +} + +// assignOne claims one job for w and sends it. It reports whether a job was sent. +func (s *Scheduler) assignOne(ctx context.Context, w workers.SessionInfo) bool { + labels, err := json.Marshal(nonNilMap(w.Labels)) + if err != nil { + return false + } + token, tokenHash, err := newJobToken() + if err != nil { + return false + } + lease := time.Now().Add(s.cfg.Lease) + job, err := s.store.ClaimJob(ctx, db.ClaimJobParams{ + WorkerID: &w.WorkerID, LeaseExpiresAt: &lease, TokenHash: tokenHash, + Pool: w.Pool, Os: w.OS, Arch: w.Arch, Images: w.Images, WorkerLabels: labels, + }) + if store.IsNotFound(err) { + return false + } + if err != nil { + s.log.Error("claiming job", "worker", w.WorkerID, "err", err) + return false + } + + spec := &opificav1.Job{} + if err := proto.Unmarshal(job.Spec, spec); err != nil { + s.log.Error("stored job spec is invalid", "job", job.ID, "err", err) + s.store.RetryJob(ctx, db.RetryJobParams{ID: job.ID, Error: "stored job spec is invalid"}) + return true + } + spec.JobToken = token + + sendCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + err = s.sessions.Send(sendCtx, w.WorkerID, &opificav1.ConnectResponse{Message: &opificav1.ConnectResponse_AssignJob{ + AssignJob: &opificav1.AssignJob{Job: spec, Lease: durationpb.New(s.cfg.Lease)}, + }}) + if err != nil { + s.log.Warn("could not deliver job; putting it back", "job", job.ID, "worker", w.WorkerID, "err", err) + if err := s.store.UnclaimJob(ctx, db.UnclaimJobParams{ID: job.ID, WorkerID: &w.WorkerID}); err != nil { + s.log.Error("unclaiming job", "job", job.ID, "err", err) + } + return false + } + s.log.Info("job assigned", "job", job.ID, "key", job.Key, "worker", w.WorkerID, "attempt", job.Attempt) + return true +} + +func nonNilMap(m map[string]string) map[string]string { + if m == nil { + return map[string]string{} + } + return m +} + +// newJobToken returns a token for one job attempt and its hash. The token +// will scope the job's secret requests and uploads; only the hash is stored. +func newJobToken() (string, []byte, error) { + b := make([]byte, 32) + if _, err := rand.Read(b); err != nil { + return "", nil, err + } + tok := "opt_" + base64.RawURLEncoding.EncodeToString(b) + h := sha256.Sum256([]byte(tok)) + return tok, h[:], nil +} + +func (s *Scheduler) reapLeases(ctx context.Context) { + expired, err := s.store.ExpiredLeases(ctx) + if err != nil { + s.log.Error("finding expired leases", "err", err) + return + } + for _, e := range expired { + job, err := s.store.RetryJob(ctx, db.RetryJobParams{ + ID: e.ID, WorkerID: e.WorkerID, Error: "the worker stopped responding (lease expired)", + }) + if store.IsNotFound(err) { + continue // finished meanwhile + } + if err != nil { + s.log.Error("retrying job", "job", e.ID, "err", err) + continue + } + s.log.Warn("lease expired", "job", job.ID, "worker", e.WorkerID, "attempt", job.Attempt, "now", job.State) + if e.WorkerID != nil { + // In case the worker is alive but lost track of the job. + s.cancelOnWorker(ctx, *e.WorkerID, job.ID, "lease expired") + } + s.settled(ctx, job) + } +} + +// settled continues after a job changed state through a retry. +func (s *Scheduler) settled(ctx context.Context, job db.Job) { + if job.State == "queued" { + s.Kick() + return + } + s.advance(ctx, job.RunID) +} + +// propagateCancels tells workers connected here to stop jobs of runs that +// were cancelled, possibly by another process. +func (s *Scheduler) propagateCancels(ctx context.Context) { + jobs, err := s.store.JobsToCancel(ctx) + if err != nil { + s.log.Error("finding jobs to cancel", "err", err) + return + } + for _, j := range jobs { + if _, ok := s.sessions.Get(*j.WorkerID); ok { + s.cancelOnWorker(ctx, *j.WorkerID, j.ID, "run cancelled") + } + } +} + +func (s *Scheduler) cancelOnWorker(ctx context.Context, worker, job uuid.UUID, reason string) { + sendCtx, cancel := context.WithTimeout(ctx, 2*time.Second) + defer cancel() + err := s.sessions.Send(sendCtx, worker, &opificav1.ConnectResponse{Message: &opificav1.ConnectResponse_CancelJob{ + CancelJob: &opificav1.CancelJob{JobId: job.String(), Reason: reason}, + }}) + if err != nil && !errors.Is(err, workers.ErrNotConnected) { + s.log.Warn("could not send cancel", "job", job, "worker", worker, "err", err) + } +} + +// ---------------------------------------------------------------- worker messages + +// WorkerConnected offers work to a newly connected worker. +func (s *Scheduler) WorkerConnected(uuid.UUID) { s.Kick() } + +// WorkerDisconnected puts back jobs the worker never reported starting. +func (s *Scheduler) WorkerDisconnected(worker uuid.UUID) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + ids, err := s.store.ReleaseUnstartedJobs(ctx, &worker) + if err != nil { + s.log.Error("releasing unstarted jobs", "worker", worker, "err", err) + return + } + if len(ids) > 0 { + s.log.Info("worker disconnected before starting jobs; queued them again", "worker", worker, "jobs", len(ids)) + s.Kick() + } +} + +// HandleWorkerMessage handles heartbeats and job updates. Database errors are +// logged rather than returned, so a database hiccup doesn't drop workers. +func (s *Scheduler) HandleWorkerMessage(ctx context.Context, worker uuid.UUID, msg *opificav1.ConnectRequest) error { + switch m := msg.GetMessage().(type) { + case *opificav1.ConnectRequest_Heartbeat: + s.renew(ctx, worker, m.Heartbeat.GetRunningJobIds()) + case *opificav1.ConnectRequest_JobUpdate: + s.jobUpdate(ctx, worker, m.JobUpdate) + case *opificav1.ConnectRequest_SecretRequest: + // Until the secrets broker exists, answer so the job doesn't hang. + return s.sessions.Send(ctx, worker, &opificav1.ConnectResponse{Message: &opificav1.ConnectResponse_SecretResponse{ + SecretResponse: &opificav1.SecretResponse{ + RequestId: m.SecretRequest.GetRequestId(), + Result: &opificav1.SecretResponse_DeniedReason{DeniedReason: "secrets are not implemented yet"}, + }, + }}) + case *opificav1.ConnectRequest_LogChunk: + // Log storage comes next; drop for now. + } + return nil +} + +func (s *Scheduler) renew(ctx context.Context, worker uuid.UUID, ids []string) { + if len(ids) == 0 { + return + } + jobIDs := make([]uuid.UUID, 0, len(ids)) + for _, id := range ids { + if u, err := uuid.Parse(id); err == nil { + jobIDs = append(jobIDs, u) + } + } + lease := time.Now().Add(s.cfg.Lease) + if err := s.store.RenewLeases(ctx, db.RenewLeasesParams{LeaseExpiresAt: &lease, WorkerID: &worker, JobIds: jobIDs}); err != nil { + s.log.Error("renewing leases", "worker", worker, "err", err) + } +} + +func (s *Scheduler) jobUpdate(ctx context.Context, worker uuid.UUID, u *opificav1.JobUpdate) { + jobID, err := uuid.Parse(u.GetJobId()) + if err != nil { + s.log.Warn("job update with invalid job id", "worker", worker, "job", u.GetJobId()) + return + } + log := s.log.With("job", jobID, "worker", worker) + + switch u.GetState() { + case opificav1.JobState_JOB_STATE_PREPARING, opificav1.JobState_JOB_STATE_RUNNING, opificav1.JobState_JOB_STATE_UPLOADING: + if _, err := s.store.MarkJobRunning(ctx, db.MarkJobRunningParams{ID: jobID, WorkerID: &worker}); err != nil { + s.stale(ctx, log, worker, jobID, err) + } + + case opificav1.JobState_JOB_STATE_SUCCEEDED, opificav1.JobState_JOB_STATE_FAILED, opificav1.JobState_JOB_STATE_CANCELLED: + state := map[opificav1.JobState]string{ + opificav1.JobState_JOB_STATE_SUCCEEDED: "succeeded", + opificav1.JobState_JOB_STATE_FAILED: "failed", + opificav1.JobState_JOB_STATE_CANCELLED: "cancelled", + }[u.GetState()] + job, err := s.store.FinishJob(ctx, db.FinishJobParams{ + ID: jobID, WorkerID: &worker, State: state, ExitCode: u.ExitCode, Error: u.GetError(), + }) + if err != nil { + s.stale(ctx, log, worker, jobID, err) + return + } + log.Info("job finished", "key", job.Key, "state", state) + s.advance(ctx, job.RunID) + + case opificav1.JobState_JOB_STATE_ERRORED: + job, err := s.store.RetryJob(ctx, db.RetryJobParams{ID: jobID, WorkerID: &worker, Error: u.GetError()}) + if err != nil { + s.stale(ctx, log, worker, jobID, err) + return + } + log.Warn("job errored on the worker", "err", u.GetError(), "attempt", job.Attempt, "now", job.State) + s.settled(ctx, job) + + default: + log.Warn("job update with unexpected state", "state", u.GetState()) + } +} + +// stale handles an update for a job this worker no longer holds, e.g. after +// its lease expired and the job went to another worker: tell it to stop. +func (s *Scheduler) stale(ctx context.Context, log *slog.Logger, worker, jobID uuid.UUID, err error) { + if !store.IsNotFound(err) { + log.Error("recording job update", "err", err) + return + } + log.Warn("ignoring update for a job this worker no longer holds") + s.cancelOnWorker(ctx, worker, jobID, "this job is no longer assigned to this worker") +} diff --git a/internal/controller/store/db/models.go b/internal/controller/store/db/models.go index 524dbe8..468e5a4 100644 --- a/internal/controller/store/db/models.go +++ b/internal/controller/store/db/models.go @@ -11,6 +11,32 @@ import ( "github.com/google/uuid" ) +type Job struct { + ID uuid.UUID + RunID uuid.UUID + Key string + Name string + Pool string + Os string + Arch string + Image string + Labels json.RawMessage + Needs []string + Spec []byte + State string + Attempt int32 + MaxAttempts int32 + WorkerID *uuid.UUID + LeaseExpiresAt *time.Time + TokenHash []byte + ExitCode *int32 + Error string + CreatedAt time.Time + QueuedAt *time.Time + StartedAt *time.Time + FinishedAt *time.Time +} + type JoinToken struct { ID uuid.UUID TokenHash []byte @@ -23,6 +49,24 @@ type JoinToken struct { UsedBy *uuid.UUID } +type Run struct { + ID uuid.UUID + Repo string + CloneUrl string + EventType string + Ref string + CommitSha string + Branch string + BaseBranch string + Tag string + Pool string + PipelineSource string + State string + CancelRequestedAt *time.Time + CreatedAt time.Time + FinishedAt *time.Time +} + type Worker struct { ID uuid.UUID Hostname string diff --git a/internal/controller/store/db/runs.sql.go b/internal/controller/store/db/runs.sql.go new file mode 100644 index 0000000..754cd4f --- /dev/null +++ b/internal/controller/store/db/runs.sql.go @@ -0,0 +1,799 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.1 +// source: runs.sql + +package db + +import ( + "context" + "encoding/json" + "time" + + "github.com/google/uuid" +) + +const activeRunJobs = `-- name: ActiveRunJobs :many +SELECT id, worker_id FROM jobs WHERE run_id = $1 AND state IN ('assigned', 'running') +` + +type ActiveRunJobsRow struct { + ID uuid.UUID + WorkerID *uuid.UUID +} + +func (q *Queries) ActiveRunJobs(ctx context.Context, runID uuid.UUID) ([]ActiveRunJobsRow, error) { + rows, err := q.db.Query(ctx, activeRunJobs, runID) + if err != nil { + return nil, err + } + defer rows.Close() + var items []ActiveRunJobsRow + for rows.Next() { + var i ActiveRunJobsRow + if err := rows.Scan(&i.ID, &i.WorkerID); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const cancelWaitingJobs = `-- name: CancelWaitingJobs :exec +UPDATE jobs SET state = 'cancelled', finished_at = now(), error = 'run cancelled' +WHERE run_id = $1 AND state IN ('pending', 'queued') +` + +func (q *Queries) CancelWaitingJobs(ctx context.Context, runID uuid.UUID) error { + _, err := q.db.Exec(ctx, cancelWaitingJobs, runID) + return err +} + +const claimJob = `-- name: ClaimJob :one +UPDATE jobs +SET state = 'assigned', worker_id = $1, lease_expires_at = $2, + token_hash = $3, attempt = attempt + 1, started_at = NULL +WHERE id = ( + SELECT q.id FROM jobs q + JOIN runs r ON r.id = q.run_id + WHERE q.state = 'queued' + AND r.cancel_requested_at IS NULL + AND q.pool = $4 + AND q.os = $5 + AND (q.arch = '' OR q.arch = $6) + AND q.image = ANY ($7::text[]) + AND q.labels <@ $8::jsonb + ORDER BY q.queued_at, q.id + LIMIT 1 + FOR UPDATE OF q SKIP LOCKED +) + -- Re-checked after the row lock, as a second guard against double claims. + AND state = 'queued' +RETURNING id, run_id, key, name, pool, os, arch, image, labels, needs, spec, state, attempt, max_attempts, worker_id, lease_expires_at, token_hash, exit_code, error, created_at, queued_at, started_at, finished_at +` + +type ClaimJobParams struct { + WorkerID *uuid.UUID + LeaseExpiresAt *time.Time + TokenHash []byte + Pool string + Os string + Arch string + Images []string + WorkerLabels json.RawMessage +} + +// Hands the oldest matching queued job to a worker. SKIP LOCKED makes it +// safe for several controllers to claim at once: a job goes to one worker. +func (q *Queries) ClaimJob(ctx context.Context, arg ClaimJobParams) (Job, error) { + row := q.db.QueryRow(ctx, claimJob, + arg.WorkerID, + arg.LeaseExpiresAt, + arg.TokenHash, + arg.Pool, + arg.Os, + arg.Arch, + arg.Images, + arg.WorkerLabels, + ) + var i Job + err := row.Scan( + &i.ID, + &i.RunID, + &i.Key, + &i.Name, + &i.Pool, + &i.Os, + &i.Arch, + &i.Image, + &i.Labels, + &i.Needs, + &i.Spec, + &i.State, + &i.Attempt, + &i.MaxAttempts, + &i.WorkerID, + &i.LeaseExpiresAt, + &i.TokenHash, + &i.ExitCode, + &i.Error, + &i.CreatedAt, + &i.QueuedAt, + &i.StartedAt, + &i.FinishedAt, + ) + return i, err +} + +const countWorkerJobs = `-- name: CountWorkerJobs :one +SELECT count(*) FROM jobs WHERE worker_id = $1 AND state IN ('assigned', 'running') +` + +func (q *Queries) CountWorkerJobs(ctx context.Context, workerID *uuid.UUID) (int64, error) { + row := q.db.QueryRow(ctx, countWorkerJobs, workerID) + var count int64 + err := row.Scan(&count) + return count, err +} + +const createJob = `-- name: CreateJob :exec +INSERT INTO jobs (id, run_id, key, name, pool, os, arch, image, labels, needs, spec, state, queued_at) +VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, + CASE WHEN $12::text = 'queued' THEN now() END) +` + +type CreateJobParams struct { + ID uuid.UUID + RunID uuid.UUID + Key string + Name string + Pool string + Os string + Arch string + Image string + Labels json.RawMessage + Needs []string + Spec []byte + State string +} + +func (q *Queries) CreateJob(ctx context.Context, arg CreateJobParams) error { + _, err := q.db.Exec(ctx, createJob, + arg.ID, + arg.RunID, + arg.Key, + arg.Name, + arg.Pool, + arg.Os, + arg.Arch, + arg.Image, + arg.Labels, + arg.Needs, + arg.Spec, + arg.State, + ) + return err +} + +const createRun = `-- name: CreateRun :one +INSERT INTO runs (id, repo, clone_url, event_type, ref, commit_sha, branch, base_branch, tag, pool, pipeline_source) +VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) +RETURNING id, repo, clone_url, event_type, ref, commit_sha, branch, base_branch, tag, pool, pipeline_source, state, cancel_requested_at, created_at, finished_at +` + +type CreateRunParams struct { + ID uuid.UUID + Repo string + CloneUrl string + EventType string + Ref string + CommitSha string + Branch string + BaseBranch string + Tag string + Pool string + PipelineSource string +} + +func (q *Queries) CreateRun(ctx context.Context, arg CreateRunParams) (Run, error) { + row := q.db.QueryRow(ctx, createRun, + arg.ID, + arg.Repo, + arg.CloneUrl, + arg.EventType, + arg.Ref, + arg.CommitSha, + arg.Branch, + arg.BaseBranch, + arg.Tag, + arg.Pool, + arg.PipelineSource, + ) + var i Run + err := row.Scan( + &i.ID, + &i.Repo, + &i.CloneUrl, + &i.EventType, + &i.Ref, + &i.CommitSha, + &i.Branch, + &i.BaseBranch, + &i.Tag, + &i.Pool, + &i.PipelineSource, + &i.State, + &i.CancelRequestedAt, + &i.CreatedAt, + &i.FinishedAt, + ) + return i, err +} + +const expiredLeases = `-- name: ExpiredLeases :many +SELECT id, worker_id FROM jobs +WHERE state IN ('assigned', 'running') AND lease_expires_at < now() +ORDER BY lease_expires_at +LIMIT 100 +` + +type ExpiredLeasesRow struct { + ID uuid.UUID + WorkerID *uuid.UUID +} + +func (q *Queries) ExpiredLeases(ctx context.Context) ([]ExpiredLeasesRow, error) { + rows, err := q.db.Query(ctx, expiredLeases) + if err != nil { + return nil, err + } + defer rows.Close() + var items []ExpiredLeasesRow + for rows.Next() { + var i ExpiredLeasesRow + if err := rows.Scan(&i.ID, &i.WorkerID); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const finishJob = `-- name: FinishJob :one +UPDATE jobs +SET state = $1, exit_code = $2, error = $3, + finished_at = now(), lease_expires_at = NULL, token_hash = NULL +WHERE id = $4 AND worker_id = $5 AND state IN ('assigned', 'running') +RETURNING id, run_id, key, name, pool, os, arch, image, labels, needs, spec, state, attempt, max_attempts, worker_id, lease_expires_at, token_hash, exit_code, error, created_at, queued_at, started_at, finished_at +` + +type FinishJobParams struct { + State string + ExitCode *int32 + Error string + ID uuid.UUID + WorkerID *uuid.UUID +} + +// Records a final result reported by the worker holding the job. +func (q *Queries) FinishJob(ctx context.Context, arg FinishJobParams) (Job, error) { + row := q.db.QueryRow(ctx, finishJob, + arg.State, + arg.ExitCode, + arg.Error, + arg.ID, + arg.WorkerID, + ) + var i Job + err := row.Scan( + &i.ID, + &i.RunID, + &i.Key, + &i.Name, + &i.Pool, + &i.Os, + &i.Arch, + &i.Image, + &i.Labels, + &i.Needs, + &i.Spec, + &i.State, + &i.Attempt, + &i.MaxAttempts, + &i.WorkerID, + &i.LeaseExpiresAt, + &i.TokenHash, + &i.ExitCode, + &i.Error, + &i.CreatedAt, + &i.QueuedAt, + &i.StartedAt, + &i.FinishedAt, + ) + return i, err +} + +const finishRunIfDone = `-- name: FinishRunIfDone :one +UPDATE runs r +SET state = CASE + WHEN EXISTS (SELECT 1 FROM jobs j WHERE j.run_id = r.id AND j.state = 'failed') THEN 'failed' + WHEN r.cancel_requested_at IS NOT NULL + OR EXISTS (SELECT 1 FROM jobs j WHERE j.run_id = r.id AND j.state = 'cancelled') THEN 'cancelled' + ELSE 'succeeded' END, + finished_at = now() +WHERE r.id = $1 AND r.state = 'running' + AND NOT EXISTS ( + SELECT 1 FROM jobs j + WHERE j.run_id = r.id AND j.state IN ('pending', 'queued', 'assigned', 'running')) +RETURNING id, repo, clone_url, event_type, ref, commit_sha, branch, base_branch, tag, pool, pipeline_source, state, cancel_requested_at, created_at, finished_at +` + +// Closes the run once no job can change state any more. +func (q *Queries) FinishRunIfDone(ctx context.Context, id uuid.UUID) (Run, error) { + row := q.db.QueryRow(ctx, finishRunIfDone, id) + var i Run + err := row.Scan( + &i.ID, + &i.Repo, + &i.CloneUrl, + &i.EventType, + &i.Ref, + &i.CommitSha, + &i.Branch, + &i.BaseBranch, + &i.Tag, + &i.Pool, + &i.PipelineSource, + &i.State, + &i.CancelRequestedAt, + &i.CreatedAt, + &i.FinishedAt, + ) + return i, err +} + +const getJob = `-- name: GetJob :one +SELECT id, run_id, key, name, pool, os, arch, image, labels, needs, spec, state, attempt, max_attempts, worker_id, lease_expires_at, token_hash, exit_code, error, created_at, queued_at, started_at, finished_at FROM jobs WHERE id = $1 +` + +func (q *Queries) GetJob(ctx context.Context, id uuid.UUID) (Job, error) { + row := q.db.QueryRow(ctx, getJob, id) + var i Job + err := row.Scan( + &i.ID, + &i.RunID, + &i.Key, + &i.Name, + &i.Pool, + &i.Os, + &i.Arch, + &i.Image, + &i.Labels, + &i.Needs, + &i.Spec, + &i.State, + &i.Attempt, + &i.MaxAttempts, + &i.WorkerID, + &i.LeaseExpiresAt, + &i.TokenHash, + &i.ExitCode, + &i.Error, + &i.CreatedAt, + &i.QueuedAt, + &i.StartedAt, + &i.FinishedAt, + ) + return i, err +} + +const getRun = `-- name: GetRun :one +SELECT id, repo, clone_url, event_type, ref, commit_sha, branch, base_branch, tag, pool, pipeline_source, state, cancel_requested_at, created_at, finished_at FROM runs WHERE id = $1 +` + +func (q *Queries) GetRun(ctx context.Context, id uuid.UUID) (Run, error) { + row := q.db.QueryRow(ctx, getRun, id) + var i Run + err := row.Scan( + &i.ID, + &i.Repo, + &i.CloneUrl, + &i.EventType, + &i.Ref, + &i.CommitSha, + &i.Branch, + &i.BaseBranch, + &i.Tag, + &i.Pool, + &i.PipelineSource, + &i.State, + &i.CancelRequestedAt, + &i.CreatedAt, + &i.FinishedAt, + ) + return i, err +} + +const jobsToCancel = `-- name: JobsToCancel :many +SELECT j.id, j.worker_id FROM jobs j +JOIN runs r ON r.id = j.run_id +WHERE r.cancel_requested_at IS NOT NULL AND j.state IN ('assigned', 'running') AND j.worker_id IS NOT NULL +` + +type JobsToCancelRow struct { + ID uuid.UUID + WorkerID *uuid.UUID +} + +// Jobs still running in runs being cancelled. Each controller tells the +// workers it holds sessions for, so cancelling works from any process. +func (q *Queries) JobsToCancel(ctx context.Context) ([]JobsToCancelRow, error) { + rows, err := q.db.Query(ctx, jobsToCancel) + if err != nil { + return nil, err + } + defer rows.Close() + var items []JobsToCancelRow + for rows.Next() { + var i JobsToCancelRow + if err := rows.Scan(&i.ID, &i.WorkerID); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const listRunJobs = `-- name: ListRunJobs :many +SELECT id, run_id, key, name, pool, os, arch, image, labels, needs, spec, state, attempt, max_attempts, worker_id, lease_expires_at, token_hash, exit_code, error, created_at, queued_at, started_at, finished_at FROM jobs WHERE run_id = $1 ORDER BY created_at, key +` + +func (q *Queries) ListRunJobs(ctx context.Context, runID uuid.UUID) ([]Job, error) { + rows, err := q.db.Query(ctx, listRunJobs, runID) + if err != nil { + return nil, err + } + defer rows.Close() + var items []Job + for rows.Next() { + var i Job + if err := rows.Scan( + &i.ID, + &i.RunID, + &i.Key, + &i.Name, + &i.Pool, + &i.Os, + &i.Arch, + &i.Image, + &i.Labels, + &i.Needs, + &i.Spec, + &i.State, + &i.Attempt, + &i.MaxAttempts, + &i.WorkerID, + &i.LeaseExpiresAt, + &i.TokenHash, + &i.ExitCode, + &i.Error, + &i.CreatedAt, + &i.QueuedAt, + &i.StartedAt, + &i.FinishedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const listRuns = `-- name: ListRuns :many +SELECT id, repo, clone_url, event_type, ref, commit_sha, branch, base_branch, tag, pool, pipeline_source, state, cancel_requested_at, created_at, finished_at FROM runs ORDER BY created_at DESC LIMIT $1 +` + +func (q *Queries) ListRuns(ctx context.Context, limit int32) ([]Run, error) { + rows, err := q.db.Query(ctx, listRuns, limit) + if err != nil { + return nil, err + } + defer rows.Close() + var items []Run + for rows.Next() { + var i Run + if err := rows.Scan( + &i.ID, + &i.Repo, + &i.CloneUrl, + &i.EventType, + &i.Ref, + &i.CommitSha, + &i.Branch, + &i.BaseBranch, + &i.Tag, + &i.Pool, + &i.PipelineSource, + &i.State, + &i.CancelRequestedAt, + &i.CreatedAt, + &i.FinishedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const markJobRunning = `-- name: MarkJobRunning :one +UPDATE jobs SET state = 'running', started_at = coalesce(started_at, now()) +WHERE id = $1 AND worker_id = $2 AND state IN ('assigned', 'running') +RETURNING id, run_id, key, name, pool, os, arch, image, labels, needs, spec, state, attempt, max_attempts, worker_id, lease_expires_at, token_hash, exit_code, error, created_at, queued_at, started_at, finished_at +` + +type MarkJobRunningParams struct { + ID uuid.UUID + WorkerID *uuid.UUID +} + +func (q *Queries) MarkJobRunning(ctx context.Context, arg MarkJobRunningParams) (Job, error) { + row := q.db.QueryRow(ctx, markJobRunning, arg.ID, arg.WorkerID) + var i Job + err := row.Scan( + &i.ID, + &i.RunID, + &i.Key, + &i.Name, + &i.Pool, + &i.Os, + &i.Arch, + &i.Image, + &i.Labels, + &i.Needs, + &i.Spec, + &i.State, + &i.Attempt, + &i.MaxAttempts, + &i.WorkerID, + &i.LeaseExpiresAt, + &i.TokenHash, + &i.ExitCode, + &i.Error, + &i.CreatedAt, + &i.QueuedAt, + &i.StartedAt, + &i.FinishedAt, + ) + return i, err +} + +const queueReadyJobs = `-- name: QueueReadyJobs :many +UPDATE jobs j +SET state = 'queued', queued_at = now() +WHERE j.run_id = $1 AND j.state = 'pending' + AND NOT EXISTS ( + SELECT 1 FROM jobs d + WHERE d.run_id = j.run_id AND d.key = ANY (j.needs) AND d.state <> 'succeeded') +RETURNING j.id +` + +// Pending jobs whose needs have all succeeded become queued. +func (q *Queries) QueueReadyJobs(ctx context.Context, runID uuid.UUID) ([]uuid.UUID, error) { + rows, err := q.db.Query(ctx, queueReadyJobs, runID) + if err != nil { + return nil, err + } + defer rows.Close() + var items []uuid.UUID + for rows.Next() { + var id uuid.UUID + if err := rows.Scan(&id); err != nil { + return nil, err + } + items = append(items, id) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const releaseUnstartedJobs = `-- name: ReleaseUnstartedJobs :many +UPDATE jobs +SET state = 'queued', worker_id = NULL, lease_expires_at = NULL, token_hash = NULL, attempt = attempt - 1 +WHERE worker_id = $1 AND state = 'assigned' +RETURNING id +` + +// When a worker's session ends, jobs it was given but never reported +// starting go back to the queue without counting the attempt: the message +// may never have reached it. Started jobs keep their lease, since the worker +// may reconnect and carry on. +func (q *Queries) ReleaseUnstartedJobs(ctx context.Context, workerID *uuid.UUID) ([]uuid.UUID, error) { + rows, err := q.db.Query(ctx, releaseUnstartedJobs, workerID) + if err != nil { + return nil, err + } + defer rows.Close() + var items []uuid.UUID + for rows.Next() { + var id uuid.UUID + if err := rows.Scan(&id); err != nil { + return nil, err + } + items = append(items, id) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const renewLeases = `-- name: RenewLeases :exec +UPDATE jobs SET lease_expires_at = $1 +WHERE worker_id = $2 AND id = ANY ($3::uuid[]) + AND state IN ('assigned', 'running') +` + +type RenewLeasesParams struct { + LeaseExpiresAt *time.Time + WorkerID *uuid.UUID + JobIds []uuid.UUID +} + +func (q *Queries) RenewLeases(ctx context.Context, arg RenewLeasesParams) error { + _, err := q.db.Exec(ctx, renewLeases, arg.LeaseExpiresAt, arg.WorkerID, arg.JobIds) + return err +} + +const requestRunCancel = `-- name: RequestRunCancel :one +UPDATE runs SET cancel_requested_at = now() +WHERE id = $1 AND state = 'running' AND cancel_requested_at IS NULL +RETURNING id, repo, clone_url, event_type, ref, commit_sha, branch, base_branch, tag, pool, pipeline_source, state, cancel_requested_at, created_at, finished_at +` + +func (q *Queries) RequestRunCancel(ctx context.Context, id uuid.UUID) (Run, error) { + row := q.db.QueryRow(ctx, requestRunCancel, id) + var i Run + err := row.Scan( + &i.ID, + &i.Repo, + &i.CloneUrl, + &i.EventType, + &i.Ref, + &i.CommitSha, + &i.Branch, + &i.BaseBranch, + &i.Tag, + &i.Pool, + &i.PipelineSource, + &i.State, + &i.CancelRequestedAt, + &i.CreatedAt, + &i.FinishedAt, + ) + return i, err +} + +const retryJob = `-- name: RetryJob :one +UPDATE jobs j +SET state = CASE + WHEN j.attempt < j.max_attempts AND r.cancel_requested_at IS NULL THEN 'queued' + WHEN r.cancel_requested_at IS NOT NULL THEN 'cancelled' + ELSE 'failed' END, + queued_at = now(), worker_id = NULL, lease_expires_at = NULL, token_hash = NULL, + error = $1, + finished_at = CASE + WHEN j.attempt < j.max_attempts AND r.cancel_requested_at IS NULL THEN NULL + ELSE now() END +FROM runs r +WHERE j.id = $2 AND r.id = j.run_id + AND ($3::uuid IS NULL OR j.worker_id = $3) + AND j.state IN ('assigned', 'running') +RETURNING j.id, j.run_id, j.key, j.name, j.pool, j.os, j.arch, j.image, j.labels, j.needs, j.spec, j.state, j.attempt, j.max_attempts, j.worker_id, j.lease_expires_at, j.token_hash, j.exit_code, j.error, j.created_at, j.queued_at, j.started_at, j.finished_at +` + +type RetryJobParams struct { + Error string + ID uuid.UUID + WorkerID *uuid.UUID +} + +// An attempt failed for infrastructure reasons: queue it again if attempts +// remain and the run isn't being cancelled, otherwise fail it. +func (q *Queries) RetryJob(ctx context.Context, arg RetryJobParams) (Job, error) { + row := q.db.QueryRow(ctx, retryJob, arg.Error, arg.ID, arg.WorkerID) + var i Job + err := row.Scan( + &i.ID, + &i.RunID, + &i.Key, + &i.Name, + &i.Pool, + &i.Os, + &i.Arch, + &i.Image, + &i.Labels, + &i.Needs, + &i.Spec, + &i.State, + &i.Attempt, + &i.MaxAttempts, + &i.WorkerID, + &i.LeaseExpiresAt, + &i.TokenHash, + &i.ExitCode, + &i.Error, + &i.CreatedAt, + &i.QueuedAt, + &i.StartedAt, + &i.FinishedAt, + ) + return i, err +} + +const skipBlockedJobs = `-- name: SkipBlockedJobs :many +UPDATE jobs j +SET state = 'skipped', finished_at = now(), error = 'a job it needs did not succeed' +WHERE j.run_id = $1 AND j.state = 'pending' + AND EXISTS ( + SELECT 1 FROM jobs d + WHERE d.run_id = j.run_id AND d.key = ANY (j.needs) + AND d.state IN ('failed', 'cancelled', 'skipped')) +RETURNING j.id +` + +// Pending jobs that need a job which didn't succeed will never run. +func (q *Queries) SkipBlockedJobs(ctx context.Context, runID uuid.UUID) ([]uuid.UUID, error) { + rows, err := q.db.Query(ctx, skipBlockedJobs, runID) + if err != nil { + return nil, err + } + defer rows.Close() + var items []uuid.UUID + for rows.Next() { + var id uuid.UUID + if err := rows.Scan(&id); err != nil { + return nil, err + } + items = append(items, id) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const unclaimJob = `-- name: UnclaimJob :exec +UPDATE jobs +SET state = 'queued', worker_id = NULL, lease_expires_at = NULL, token_hash = NULL, attempt = attempt - 1 +WHERE id = $1 AND worker_id = $2 AND state = 'assigned' +` + +type UnclaimJobParams struct { + ID uuid.UUID + WorkerID *uuid.UUID +} + +// Puts back a job that couldn't be delivered to the worker that claimed it, +// without counting the attempt. +func (q *Queries) UnclaimJob(ctx context.Context, arg UnclaimJobParams) error { + _, err := q.db.Exec(ctx, unclaimJob, arg.ID, arg.WorkerID) + return err +} diff --git a/internal/controller/store/migrations/0002_runs.sql b/internal/controller/store/migrations/0002_runs.sql new file mode 100644 index 0000000..6849458 --- /dev/null +++ b/internal/controller/store/migrations/0002_runs.sql @@ -0,0 +1,68 @@ +-- A run is one pipeline execution for one event, pinned to one commit. +CREATE TABLE runs ( + id uuid PRIMARY KEY, + repo text NOT NULL, + clone_url text NOT NULL DEFAULT '', + event_type text NOT NULL CHECK (event_type IN ('push', 'pull_request', 'tag')), + ref text NOT NULL DEFAULT '', + commit_sha text NOT NULL, + branch text NOT NULL DEFAULT '', + base_branch text NOT NULL DEFAULT '', + tag text NOT NULL DEFAULT '', + -- Decided by the controller's routing policy, never by the pipeline file. + pool text NOT NULL, + -- The pipeline file exactly as it was at commit_sha, for auditing. + pipeline_source text NOT NULL, + state text NOT NULL DEFAULT 'running' + CHECK (state IN ('running', 'succeeded', 'failed', 'cancelled')), + cancel_requested_at timestamptz, + created_at timestamptz NOT NULL DEFAULT now(), + finished_at timestamptz +); + +CREATE INDEX runs_created ON runs (created_at DESC); + +-- A job is one planned job instance (one matrix combination). +-- +-- pending waiting for the jobs it needs +-- queued ready; waiting for a matching worker +-- assigned sent to a worker, which holds a lease +-- running the worker reported it started +-- succeeded, failed, cancelled, skipped: finished +CREATE TABLE jobs ( + id uuid PRIMARY KEY, + run_id uuid NOT NULL REFERENCES runs (id) ON DELETE CASCADE, + -- Planned ID, unique in the run, e.g. "test[os=linux]". + key text NOT NULL, + name text NOT NULL, + -- What the job needs from a worker. + pool text NOT NULL, + os text NOT NULL, + arch text NOT NULL DEFAULT '', + image text NOT NULL, + labels jsonb NOT NULL DEFAULT '{}', + -- Keys of jobs in the same run that must succeed first. + needs text[] NOT NULL DEFAULT '{}', + -- The opifica.v1.Job sent to the worker, without the per-attempt token. + spec bytea NOT NULL, + state text NOT NULL + CHECK (state IN ('pending', 'queued', 'assigned', 'running', + 'succeeded', 'failed', 'cancelled', 'skipped')), + attempt integer NOT NULL DEFAULT 0, + max_attempts integer NOT NULL DEFAULT 3, + worker_id uuid REFERENCES workers (id), + lease_expires_at timestamptz, + -- SHA-256 of the current attempt's job token. + token_hash bytea, + exit_code integer, + error text NOT NULL DEFAULT '', + created_at timestamptz NOT NULL DEFAULT now(), + queued_at timestamptz, + started_at timestamptz, + finished_at timestamptz, + UNIQUE (run_id, key) +); + +CREATE INDEX jobs_queued ON jobs (pool, os, queued_at) WHERE state = 'queued'; +CREATE INDEX jobs_leased ON jobs (lease_expires_at) WHERE state IN ('assigned', 'running'); +CREATE INDEX jobs_worker ON jobs (worker_id) WHERE state IN ('assigned', 'running'); diff --git a/internal/controller/store/queries/runs.sql b/internal/controller/store/queries/runs.sql new file mode 100644 index 0000000..cd794ba --- /dev/null +++ b/internal/controller/store/queries/runs.sql @@ -0,0 +1,163 @@ +-- name: CreateRun :one +INSERT INTO runs (id, repo, clone_url, event_type, ref, commit_sha, branch, base_branch, tag, pool, pipeline_source) +VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) +RETURNING *; + +-- name: GetRun :one +SELECT * FROM runs WHERE id = $1; + +-- name: ListRuns :many +SELECT * FROM runs ORDER BY created_at DESC LIMIT $1; + +-- name: CreateJob :exec +INSERT INTO jobs (id, run_id, key, name, pool, os, arch, image, labels, needs, spec, state, queued_at) +VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, + CASE WHEN $12::text = 'queued' THEN now() END); + +-- name: GetJob :one +SELECT * FROM jobs WHERE id = $1; + +-- name: ListRunJobs :many +SELECT * FROM jobs WHERE run_id = $1 ORDER BY created_at, key; + +-- name: ClaimJob :one +-- Hands the oldest matching queued job to a worker. SKIP LOCKED makes it +-- safe for several controllers to claim at once: a job goes to one worker. +UPDATE jobs +SET state = 'assigned', worker_id = $1, lease_expires_at = sqlc.arg(lease_expires_at), + token_hash = sqlc.arg(token_hash), attempt = attempt + 1, started_at = NULL +WHERE id = ( + SELECT q.id FROM jobs q + JOIN runs r ON r.id = q.run_id + WHERE q.state = 'queued' + AND r.cancel_requested_at IS NULL + AND q.pool = sqlc.arg(pool) + AND q.os = sqlc.arg(os) + AND (q.arch = '' OR q.arch = sqlc.arg(arch)) + AND q.image = ANY (sqlc.arg(images)::text[]) + AND q.labels <@ sqlc.arg(worker_labels)::jsonb + ORDER BY q.queued_at, q.id + LIMIT 1 + FOR UPDATE OF q SKIP LOCKED +) + -- Re-checked after the row lock, as a second guard against double claims. + AND state = 'queued' +RETURNING *; + +-- name: UnclaimJob :exec +-- Puts back a job that couldn't be delivered to the worker that claimed it, +-- without counting the attempt. +UPDATE jobs +SET state = 'queued', worker_id = NULL, lease_expires_at = NULL, token_hash = NULL, attempt = attempt - 1 +WHERE id = $1 AND worker_id = $2 AND state = 'assigned'; + +-- name: CountWorkerJobs :one +SELECT count(*) FROM jobs WHERE worker_id = $1 AND state IN ('assigned', 'running'); + +-- name: RenewLeases :exec +UPDATE jobs SET lease_expires_at = sqlc.arg(lease_expires_at) +WHERE worker_id = sqlc.arg(worker_id) AND id = ANY (sqlc.arg(job_ids)::uuid[]) + AND state IN ('assigned', 'running'); + +-- name: MarkJobRunning :one +UPDATE jobs SET state = 'running', started_at = coalesce(started_at, now()) +WHERE id = $1 AND worker_id = $2 AND state IN ('assigned', 'running') +RETURNING *; + +-- name: FinishJob :one +-- Records a final result reported by the worker holding the job. +UPDATE jobs +SET state = sqlc.arg(state), exit_code = sqlc.arg(exit_code), error = sqlc.arg(error), + finished_at = now(), lease_expires_at = NULL, token_hash = NULL +WHERE id = sqlc.arg(id) AND worker_id = sqlc.arg(worker_id) AND state IN ('assigned', 'running') +RETURNING *; + +-- name: RetryJob :one +-- An attempt failed for infrastructure reasons: queue it again if attempts +-- remain and the run isn't being cancelled, otherwise fail it. +UPDATE jobs j +SET state = CASE + WHEN j.attempt < j.max_attempts AND r.cancel_requested_at IS NULL THEN 'queued' + WHEN r.cancel_requested_at IS NOT NULL THEN 'cancelled' + ELSE 'failed' END, + queued_at = now(), worker_id = NULL, lease_expires_at = NULL, token_hash = NULL, + error = sqlc.arg(error), + finished_at = CASE + WHEN j.attempt < j.max_attempts AND r.cancel_requested_at IS NULL THEN NULL + ELSE now() END +FROM runs r +WHERE j.id = sqlc.arg(id) AND r.id = j.run_id + AND (sqlc.narg(worker_id)::uuid IS NULL OR j.worker_id = sqlc.narg(worker_id)) + AND j.state IN ('assigned', 'running') +RETURNING j.*; + +-- name: ExpiredLeases :many +SELECT id, worker_id FROM jobs +WHERE state IN ('assigned', 'running') AND lease_expires_at < now() +ORDER BY lease_expires_at +LIMIT 100; + +-- name: SkipBlockedJobs :many +-- Pending jobs that need a job which didn't succeed will never run. +UPDATE jobs j +SET state = 'skipped', finished_at = now(), error = 'a job it needs did not succeed' +WHERE j.run_id = $1 AND j.state = 'pending' + AND EXISTS ( + SELECT 1 FROM jobs d + WHERE d.run_id = j.run_id AND d.key = ANY (j.needs) + AND d.state IN ('failed', 'cancelled', 'skipped')) +RETURNING j.id; + +-- name: QueueReadyJobs :many +-- Pending jobs whose needs have all succeeded become queued. +UPDATE jobs j +SET state = 'queued', queued_at = now() +WHERE j.run_id = $1 AND j.state = 'pending' + AND NOT EXISTS ( + SELECT 1 FROM jobs d + WHERE d.run_id = j.run_id AND d.key = ANY (j.needs) AND d.state <> 'succeeded') +RETURNING j.id; + +-- name: CancelWaitingJobs :exec +UPDATE jobs SET state = 'cancelled', finished_at = now(), error = 'run cancelled' +WHERE run_id = $1 AND state IN ('pending', 'queued'); + +-- name: ActiveRunJobs :many +SELECT id, worker_id FROM jobs WHERE run_id = $1 AND state IN ('assigned', 'running'); + +-- name: RequestRunCancel :one +UPDATE runs SET cancel_requested_at = now() +WHERE id = $1 AND state = 'running' AND cancel_requested_at IS NULL +RETURNING *; + +-- name: FinishRunIfDone :one +-- Closes the run once no job can change state any more. +UPDATE runs r +SET state = CASE + WHEN EXISTS (SELECT 1 FROM jobs j WHERE j.run_id = r.id AND j.state = 'failed') THEN 'failed' + WHEN r.cancel_requested_at IS NOT NULL + OR EXISTS (SELECT 1 FROM jobs j WHERE j.run_id = r.id AND j.state = 'cancelled') THEN 'cancelled' + ELSE 'succeeded' END, + finished_at = now() +WHERE r.id = $1 AND r.state = 'running' + AND NOT EXISTS ( + SELECT 1 FROM jobs j + WHERE j.run_id = r.id AND j.state IN ('pending', 'queued', 'assigned', 'running')) +RETURNING *; + +-- name: JobsToCancel :many +-- Jobs still running in runs being cancelled. Each controller tells the +-- workers it holds sessions for, so cancelling works from any process. +SELECT j.id, j.worker_id FROM jobs j +JOIN runs r ON r.id = j.run_id +WHERE r.cancel_requested_at IS NOT NULL AND j.state IN ('assigned', 'running') AND j.worker_id IS NOT NULL; + +-- name: ReleaseUnstartedJobs :many +-- When a worker's session ends, jobs it was given but never reported +-- starting go back to the queue without counting the attempt: the message +-- may never have reached it. Started jobs keep their lease, since the worker +-- may reconnect and carry on. +UPDATE jobs +SET state = 'queued', worker_id = NULL, lease_expires_at = NULL, token_hash = NULL, attempt = attempt - 1 +WHERE worker_id = $1 AND state = 'assigned' +RETURNING id; diff --git a/internal/controller/workers/connect.go b/internal/controller/workers/connect.go index 99151d1..e7e1e18 100644 --- a/internal/controller/workers/connect.go +++ b/internal/controller/workers/connect.go @@ -58,10 +58,16 @@ type Stream interface { Send(*opificav1.ConnectResponse) error } -// MessageHandler receives what workers send about jobs. It is nil until the -// scheduler exists; until then those messages are logged and dropped. +// MessageHandler receives heartbeats and job messages from workers, after +// the session has been updated. Returning an error ends the session. type MessageHandler interface { HandleWorkerMessage(ctx context.Context, worker uuid.UUID, msg *opificav1.ConnectRequest) error + // WorkerConnected is called when a session starts, e.g. so the + // scheduler can offer it work. + WorkerConnected(worker uuid.UUID) + // WorkerDisconnected is called when a worker's current session ends, + // but not when a newer connection from the same worker replaced it. + WorkerDisconnected(worker uuid.UUID) } type peerCertKey struct{} @@ -171,16 +177,29 @@ func (s *Service) Connect(ctx context.Context, cert *x509.Certificate, stream St if err := json.Unmarshal(w.Labels, &labels); err != nil { return connect.NewError(connect.CodeInternal, errors.New("stored worker labels are invalid")) } + var images []string + for _, img := range info.GetImages() { + if img.GetName() != "" { + images = append(images, img.GetName()) + } + } now := s.now() sess := &Session{ - WorkerID: id, Pool: w.Pool, OS: w.Os, Arch: w.Arch, Labels: labels, MaxJobs: maxJobs, + WorkerID: id, Pool: w.Pool, OS: w.Os, Arch: w.Arch, Labels: labels, Images: images, MaxJobs: maxJobs, ConnectedAt: now, lastSeen: now, running: hello.GetRunningJobIds(), out: make(chan *opificav1.ConnectResponse, 16), ctx: ctx, cancel: cancel, } s.Sessions().register(sess) - defer s.Sessions().unregister(sess) + defer func() { + if s.Sessions().unregister(sess) && s.Messages != nil { + s.Messages.WorkerDisconnected(id) + } + }() log := s.log().With("worker", id, "hostname", info.GetHostname()) - log.Info("worker connected", "pool", w.Pool, "os", w.Os, "arch", w.Arch, "max_jobs", maxJobs) + log.Info("worker connected", "pool", w.Pool, "os", w.Os, "arch", w.Arch, "max_jobs", maxJobs, "images", images) + if s.Messages != nil { + s.Messages.WorkerConnected(id) + } // Writer: the only goroutine that sends on the stream. writer.Add(1) @@ -271,6 +290,9 @@ func (s *Service) handle(ctx context.Context, sess *Session, m *opificav1.Connec switch msg := m.GetMessage().(type) { case *opificav1.ConnectRequest_Heartbeat: sess.heartbeat(msg.Heartbeat, s.now()) + if s.Messages != nil { + return s.Messages.HandleWorkerMessage(ctx, sess.WorkerID, m) + } return nil case *opificav1.ConnectRequest_Hello: return connect.NewError(connect.CodeInvalidArgument, errors.New("hello may only be sent once")) diff --git a/internal/controller/workers/registry.go b/internal/controller/workers/registry.go index f41863c..d43ef5c 100644 --- a/internal/controller/workers/registry.go +++ b/internal/controller/workers/registry.go @@ -18,11 +18,13 @@ var ErrNotConnected = errors.New("worker is not connected") // Session is one live Connect stream from a worker. type Session struct { - WorkerID uuid.UUID - Pool string - OS string - Arch string - Labels map[string]string + WorkerID uuid.UUID + Pool string + OS string + Arch string + Labels map[string]string + // Images are the golden image names the worker can boot. + Images []string MaxJobs int ConnectedAt time.Time @@ -69,6 +71,7 @@ type SessionInfo struct { OS string Arch string Labels map[string]string + Images []string MaxJobs int ConnectedAt time.Time LastSeen time.Time @@ -81,7 +84,7 @@ func (s *Session) info() SessionInfo { defer s.mu.Unlock() return SessionInfo{ WorkerID: s.WorkerID, Pool: s.Pool, OS: s.OS, Arch: s.Arch, - Labels: maps.Clone(s.Labels), MaxJobs: s.MaxJobs, ConnectedAt: s.ConnectedAt, + Labels: maps.Clone(s.Labels), Images: slices.Clone(s.Images), MaxJobs: s.MaxJobs, ConnectedAt: s.ConnectedAt, LastSeen: s.lastSeen, WarmVMs: s.warmVMs, RunningJobs: slices.Clone(s.running), } } @@ -111,12 +114,16 @@ func (r *Registry) register(s *Session) { } } -func (r *Registry) unregister(s *Session) { +// unregister removes s and reports whether it was still the worker's +// current session (false if a newer connection replaced it). +func (r *Registry) unregister(s *Session) bool { r.mu.Lock() defer r.mu.Unlock() if r.sessions[s.WorkerID] == s { delete(r.sessions, s.WorkerID) + return true } + return false } // Get returns the live session for a worker. diff --git a/internal/integration/client_stop_test.go b/internal/integration/client_stop_test.go new file mode 100644 index 0000000..19c1fcd --- /dev/null +++ b/internal/integration/client_stop_test.go @@ -0,0 +1,28 @@ +package integration + +import ( + "testing" + "time" + + opificav1 "github.com/opifica/opifica/gen/opifica/v1" + "github.com/opifica/opifica/internal/pipeline" +) + +// Stopping a worker must end its stream at once, even while it waits for a +// message, so the controller can put its unstarted jobs back right away. +// (Cancelling the context alone left Receive blocked until the controller +// ended the stream.) +func TestClientStopsPromptly(t *testing.T) { + e := newEnv(t) + w := e.fakeWorker("w", workerOpts{behave: func(*opificav1.Job) outcome { return outcome{silent: true} }}) + e.createRun("on: push\njobs:\n build:\n runs-on: linux\n steps: [{ run: make }]\n", pipeline.Event{}, "") + eventually(t, "job", func() bool { return len(w.jobs()) == 1 }) + + start := time.Now() + w.cancel() + err := <-w.done + w.done <- err + if d := time.Since(start); d > 200*time.Millisecond { + t.Errorf("Run took %v to return after cancel", d) + } +} diff --git a/internal/integration/connect_test.go b/internal/integration/connect_test.go index aa8060f..a7d31e2 100644 --- a/internal/integration/connect_test.go +++ b/internal/integration/connect_test.go @@ -21,6 +21,7 @@ import ( "github.com/opifica/opifica/gen/opifica/v1/opificav1connect" "github.com/opifica/opifica/internal/controller/api" "github.com/opifica/opifica/internal/controller/ca" + "github.com/opifica/opifica/internal/controller/scheduler" "github.com/opifica/opifica/internal/controller/store" "github.com/opifica/opifica/internal/controller/workers" "github.com/opifica/opifica/internal/server" @@ -30,12 +31,18 @@ import ( func TestMain(m *testing.M) { os.Exit(pgtest.Run(m)) } -func quiet() *slog.Logger { return slog.New(slog.NewTextHandler(io.Discard, nil)) } +func quiet() *slog.Logger { + if os.Getenv("OPIFICA_TEST_LOG") != "" { + return slog.New(slog.NewTextHandler(os.Stderr, nil)) + } + return slog.New(slog.NewTextHandler(io.Discard, nil)) +} type env struct { - t *testing.T - svc *workers.Service - url string + t *testing.T + svc *workers.Service + sched *scheduler.Scheduler + url string } // newEnv starts a controller with short timeouts on a random port. @@ -58,6 +65,11 @@ func newEnv(t *testing.T) *env { TouchInterval: 50 * time.Millisecond, }, } + sched := scheduler.New(st, svc.Sessions(), quiet(), scheduler.Config{ + Lease: 400 * time.Millisecond, Interval: 20 * time.Millisecond, + }) + svc.Messages = sched + go sched.Run(ctx) tlsCfg, err := api.TLSConfig(authority, []string{"127.0.0.1"}) if err != nil { t.Fatal(err) @@ -76,7 +88,7 @@ func newEnv(t *testing.T) *env { <-done st.Close() }) - return &env{t: t, svc: svc, url: "https://" + ln.Addr().String()} + return &env{t: t, svc: svc, sched: sched, url: "https://" + ln.Addr().String()} } func info() *opificav1.WorkerInfo { @@ -84,11 +96,17 @@ func info() *opificav1.WorkerInfo { Hostname: "build-1", Os: opificav1.OperatingSystem_OPERATING_SYSTEM_LINUX, Arch: opificav1.Architecture_ARCHITECTURE_AMD64, MaxConcurrentJobs: 4, Version: "test", Labels: map[string]string{"kvm": "true"}, + Images: []*opificav1.ImageRef{{Name: "default"}}, } } // join creates a token for pool and joins a worker with it. func (e *env) join(pool string) *worker.Identity { + e.t.Helper() + return e.joinWith(pool, info()) +} + +func (e *env) joinWith(pool string, wi *opificav1.WorkerInfo) *worker.Identity { e.t.Helper() tok, err := e.svc.CreateToken(context.Background(), workers.TokenRequest{Pool: pool}) if err != nil { @@ -96,7 +114,7 @@ func (e *env) join(pool string) *worker.Identity { } id, err := worker.Join(context.Background(), worker.JoinConfig{ ControllerURL: e.url, Token: tok.Secret, CAFingerprint: tok.CAFingerprint, - StateDir: e.t.TempDir(), Info: info(), + StateDir: e.t.TempDir(), Info: wi, }) if err != nil { e.t.Fatalf("Join: %v", err) diff --git a/internal/integration/scheduler_test.go b/internal/integration/scheduler_test.go new file mode 100644 index 0000000..acf4702 --- /dev/null +++ b/internal/integration/scheduler_test.go @@ -0,0 +1,620 @@ +package integration + +import ( + "context" + "slices" + "strings" + "sync" + "testing" + "time" + + "github.com/google/uuid" + + opificav1 "github.com/opifica/opifica/gen/opifica/v1" + "github.com/opifica/opifica/internal/controller/scheduler" + "github.com/opifica/opifica/internal/controller/store/db" + "github.com/opifica/opifica/internal/controller/workers" + "github.com/opifica/opifica/internal/pipeline" + "github.com/opifica/opifica/pkg/worker" +) + +// outcome is what a fake worker does with a job. +type outcome struct { + state opificav1.JobState // final state to report + exit int32 + delay time.Duration + // silent: accept the job, then never mention it again, like a worker + // that crashed right after receiving it. + silent bool + // vanish: report the job running, then never mention it again, like a + // worker that crashed mid-job. + vanish bool + // hang: run until cancelled. + hang bool +} + +func succeed() outcome { return outcome{state: opificav1.JobState_JOB_STATE_SUCCEEDED} } + +// fakeWorker is a real worker client whose "jobs" follow a script. +type fakeWorker struct { + *runningClient + name string + id *worker.Identity + + mu sync.Mutex + running map[string]chan struct{} // job ID -> cancel + received []*opificav1.Job + maxRunning int + behave func(*opificav1.Job) outcome +} + +type workerOpts struct { + pool string + os opificav1.OperatingSystem + images []string + labels map[string]string + maxJobs uint32 + behave func(*opificav1.Job) outcome +} + +func (e *env) fakeWorker(name string, o workerOpts) *fakeWorker { + e.t.Helper() + if o.pool == "" { + o.pool = "trusted" + } + if o.os == 0 { + o.os = opificav1.OperatingSystem_OPERATING_SYSTEM_LINUX + } + if o.images == nil { + o.images = []string{"default"} + } + if o.maxJobs == 0 { + o.maxJobs = 4 + } + if o.behave == nil { + o.behave = func(*opificav1.Job) outcome { return succeed() } + } + wi := info() + wi.Hostname, wi.Os, wi.MaxConcurrentJobs, wi.Labels, wi.Images = name, o.os, o.maxJobs, o.labels, nil + for _, img := range o.images { + wi.Images = append(wi.Images, &opificav1.ImageRef{Name: img}) + } + + id := e.joinWith(o.pool, wi) + fw := &fakeWorker{name: name, id: id, running: map[string]chan struct{}{}, behave: o.behave} + fw.runningClient = e.start(id, func(c *worker.ClientConfig) { + c.Info = wi + c.Status = fw.status + c.OnMessage = fw.onMessage + }) + return fw +} + +func (fw *fakeWorker) status() (uint32, []string) { + fw.mu.Lock() + defer fw.mu.Unlock() + ids := make([]string, 0, len(fw.running)) + for id := range fw.running { + ids = append(ids, id) + } + return 0, ids +} + +func (fw *fakeWorker) onMessage(_ context.Context, m *opificav1.ConnectResponse) { + switch msg := m.GetMessage().(type) { + case *opificav1.ConnectResponse_AssignJob: + job := msg.AssignJob.GetJob() + out := fw.behave(job) + fw.mu.Lock() + fw.received = append(fw.received, job) + if out.silent { + fw.mu.Unlock() + return + } + cancel := make(chan struct{}) + fw.running[job.GetId()] = cancel + fw.maxRunning = max(fw.maxRunning, len(fw.running)) + fw.mu.Unlock() + go fw.run(job, out, cancel) + case *opificav1.ConnectResponse_CancelJob: + fw.mu.Lock() + if c, ok := fw.running[msg.CancelJob.GetJobId()]; ok { + close(c) + } + fw.mu.Unlock() + } +} + +func (fw *fakeWorker) update(job string, state opificav1.JobState, exit *int32) { + fw.Send(&opificav1.ConnectRequest{Message: &opificav1.ConnectRequest_JobUpdate{JobUpdate: &opificav1.JobUpdate{ + JobId: job, State: state, ExitCode: exit, + }}}) +} + +func (fw *fakeWorker) run(job *opificav1.Job, out outcome, cancel chan struct{}) { + fw.update(job.GetId(), opificav1.JobState_JOB_STATE_RUNNING, nil) + if out.vanish { + fw.mu.Lock() + delete(fw.running, job.GetId()) // stops renewing the lease + fw.mu.Unlock() + return + } + final := out.state + var wait <-chan time.Time + if !out.hang { + wait = time.After(out.delay) + } + select { + case <-wait: + case <-cancel: + final = opificav1.JobState_JOB_STATE_CANCELLED + } + fw.mu.Lock() + delete(fw.running, job.GetId()) + fw.mu.Unlock() + exit := out.exit + fw.update(job.GetId(), final, &exit) +} + +func (fw *fakeWorker) jobs() []*opificav1.Job { + fw.mu.Lock() + defer fw.mu.Unlock() + return slices.Clone(fw.received) +} + +func (e *env) createRun(src string, ev pipeline.Event, pool string) db.Run { + e.t.Helper() + if ev.Type == "" { + ev = pipeline.Event{Type: pipeline.EventPush, Branch: "main", Ref: "refs/heads/main"} + } + if ev.SHA == "" { + ev.SHA = "0123456789abcdef0123456789abcdef01234567" + } + if pool == "" { + pool = "trusted" + } + run, err := e.sched.CreateRun(context.Background(), scheduler.RunRequest{ + Repo: "acme/app", CloneURL: "https://forge.example/acme/app.git", Event: ev, Pool: pool, Pipeline: []byte(src), + }) + if err != nil { + e.t.Fatalf("CreateRun: %v", err) + } + return run +} + +// waitRun waits for the run to finish and returns it with its jobs by key. +func (e *env) waitRun(id uuid.UUID) (db.Run, map[string]db.Job) { + e.t.Helper() + ctx := context.Background() + var run db.Run + eventually(e.t, "run to finish", func() bool { + var err error + run, err = e.svc.Store.GetRun(ctx, id) + return err == nil && run.State != "running" + }) + jobs, err := e.svc.Store.ListRunJobs(ctx, id) + if err != nil { + e.t.Fatal(err) + } + byKey := map[string]db.Job{} + for _, j := range jobs { + byKey[j.Key] = j + } + return run, byKey +} + +func TestSingleJobRun(t *testing.T) { + e := newEnv(t) + w := e.fakeWorker("linux-1", workerOpts{}) + run := e.createRun(` +on: push +jobs: + build: + runs-on: linux + steps: + - run: make +`, pipeline.Event{}, "") + + got, jobs := e.waitRun(run.ID) + if got.State != "succeeded" { + t.Errorf("run state = %s", got.State) + } + if j := jobs["build"]; j.State != "succeeded" || j.Attempt != 1 || *j.ExitCode != 0 { + t.Errorf("job = state %s attempt %d", j.State, j.Attempt) + } + + // The worker received the full job spec with a fresh token. + received := w.jobs() + if len(received) != 1 { + t.Fatalf("worker received %d jobs", len(received)) + } + spec := received[0] + if spec.GetSource().GetCommitSha() != run.CommitSha || spec.GetSource().GetCloneUrl() == "" { + t.Errorf("source = %v", spec.GetSource()) + } + if !strings.HasPrefix(spec.GetJobToken(), "opt_") { + t.Errorf("job token = %q", spec.GetJobToken()) + } + if spec.GetSteps()[0].GetRun().GetCommand() != "make" { + t.Errorf("steps = %v", spec.GetSteps()) + } +} + +// Each matrix job lands on a worker of its OS; release waits for all of them. +func TestMatrixAcrossOperatingSystems(t *testing.T) { + e := newEnv(t) + slow := func(*opificav1.Job) outcome { o := succeed(); o.delay = 30 * time.Millisecond; return o } + linux := e.fakeWorker("linux", workerOpts{behave: slow}) + windows := e.fakeWorker("windows", workerOpts{os: opificav1.OperatingSystem_OPERATING_SYSTEM_WINDOWS, behave: slow}) + macos := e.fakeWorker("macos", workerOpts{os: opificav1.OperatingSystem_OPERATING_SYSTEM_MACOS, behave: slow, + images: []string{"default", "xcode-16"}}) + + run := e.createRun(` +on: push +jobs: + test: + matrix: { os: [linux, windows, macos] } + runs-on: + os: ${{ matrix.os }} + steps: [{ run: make test }] + release: + needs: test + runs-on: { os: macos, image: xcode-16 } + steps: [{ run: make release }] +`, pipeline.Event{}, "") + + got, jobs := e.waitRun(run.ID) + if got.State != "succeeded" { + t.Fatalf("run state = %s", got.State) + } + for w, want := range map[*fakeWorker][]string{ + linux: {"test (linux)"}, + windows: {"test (windows)"}, + macos: {"test (macos)", "release"}, + } { + var names []string + for _, j := range w.jobs() { + names = append(names, j.GetName()) + } + slices.Sort(names) + slices.Sort(want) + if !slices.Equal(names, want) { + t.Errorf("%s ran %v, want %v", w.name, names, want) + } + } + rel := jobs["release"] + for _, k := range []string{"test[os=linux]", "test[os=windows]", "test[os=macos]"} { + if !jobs[k].FinishedAt.Before(*rel.StartedAt) && !jobs[k].FinishedAt.Equal(*rel.StartedAt) { + t.Errorf("release started before %s finished", k) + } + } +} + +func TestFailureSkipsDependents(t *testing.T) { + e := newEnv(t) + e.fakeWorker("linux", workerOpts{behave: func(j *opificav1.Job) outcome { + if j.GetName() == "test" { + return outcome{state: opificav1.JobState_JOB_STATE_FAILED, exit: 2} + } + return succeed() + }}) + run := e.createRun(` +on: push +jobs: + test: + runs-on: linux + steps: [{ run: make test }] + lint: + runs-on: linux + steps: [{ run: make lint }] + package: + needs: test + runs-on: linux + steps: [{ run: make package }] + deploy: + needs: package + runs-on: linux + steps: [{ run: make deploy }] +`, pipeline.Event{}, "") + + got, jobs := e.waitRun(run.ID) + if got.State != "failed" { + t.Errorf("run state = %s, want failed", got.State) + } + want := map[string]string{"test": "failed", "lint": "succeeded", "package": "skipped", "deploy": "skipped"} + for k, st := range want { + if jobs[k].State != st { + t.Errorf("%s = %s, want %s", k, jobs[k].State, st) + } + } + if *jobs["test"].ExitCode != 2 { + t.Errorf("exit code = %d", *jobs["test"].ExitCode) + } +} + +// A release job never runs on a worker from a less trusted pool. The pool +// comes from the run (the controller's routing), not from the pipeline. +func TestPoolsAreIsolated(t *testing.T) { + e := newEnv(t) + untrusted := e.fakeWorker("pr-builder", workerOpts{pool: "untrusted"}) + run := e.createRun(` +on: { tag: { pattern: "v*" } } +jobs: + release: + runs-on: linux + steps: [{ run: make release }] +`, pipeline.Event{Type: pipeline.EventTag, Tag: "v1.0.0", Ref: "refs/tags/v1.0.0"}, "release") + + time.Sleep(200 * time.Millisecond) + if n := len(untrusted.jobs()); n != 0 { + t.Fatalf("untrusted worker got %d release jobs", n) + } + release := e.fakeWorker("release-builder", workerOpts{pool: "release"}) + got, _ := e.waitRun(run.ID) + if got.State != "succeeded" || len(release.jobs()) != 1 || len(untrusted.jobs()) != 0 { + t.Errorf("run %s; release worker ran %d, untrusted ran %d", got.State, len(release.jobs()), len(untrusted.jobs())) + } +} + +func TestImagesAndLabelsMatch(t *testing.T) { + e := newEnv(t) + plain := e.fakeWorker("plain", workerOpts{}) + run := e.createRun(` +on: push +jobs: + gpu: + runs-on: + os: linux + image: cuda + labels: { gpu: "true" } + steps: [{ run: nvidia-smi }] +`, pipeline.Event{}, "") + time.Sleep(200 * time.Millisecond) + if n := len(plain.jobs()); n != 0 { + t.Fatalf("worker without the image or label got the job") + } + // Right label, wrong image: still no. + labelOnly := e.fakeWorker("label-only", workerOpts{labels: map[string]string{"gpu": "true"}}) + time.Sleep(200 * time.Millisecond) + if n := len(labelOnly.jobs()); n != 0 { + t.Fatalf("worker without the image got the job") + } + // Right image, missing label: still no. + imageOnly := e.fakeWorker("image-only", workerOpts{images: []string{"cuda"}, labels: map[string]string{"gpu": "false"}}) + time.Sleep(200 * time.Millisecond) + if n := len(imageOnly.jobs()); n != 0 { + t.Fatalf("worker without the gpu=true label got the job") + } + gpu := e.fakeWorker("gpu", workerOpts{images: []string{"cuda"}, labels: map[string]string{"gpu": "true", "vram": "24g"}}) + got, _ := e.waitRun(run.ID) + if got.State != "succeeded" || len(gpu.jobs()) != 1 { + t.Errorf("run = %s, gpu worker ran %d jobs", got.State, len(gpu.jobs())) + } +} + +func TestCapacityIsRespected(t *testing.T) { + e := newEnv(t) + w := e.fakeWorker("mac", workerOpts{maxJobs: 2, behave: func(*opificav1.Job) outcome { + o := succeed() + o.delay = 50 * time.Millisecond + return o + }}) + run := e.createRun(` +on: push +jobs: + t: + matrix: { n: ["1", "2", "3", "4", "5", "6"] } + runs-on: linux + steps: [{ run: make }] +`, pipeline.Event{}, "") + got, _ := e.waitRun(run.ID) + if got.State != "succeeded" { + t.Fatalf("run = %s", got.State) + } + w.mu.Lock() + defer w.mu.Unlock() + if w.maxRunning > 2 { + t.Errorf("worker ran %d jobs at once, capacity 2", w.maxRunning) + } + if len(w.received) != 6 { + t.Errorf("worker ran %d jobs, want 6", len(w.received)) + } +} + +// A worker starts a job and crashes. When the lease expires the job goes to +// another worker. When the crashed worker comes back and reports a result +// for the job, it is ignored. +func TestExpiredLeaseIsRetriedElsewhere(t *testing.T) { + e := newEnv(t) + flaky := e.fakeWorker("flaky", workerOpts{behave: func(*opificav1.Job) outcome { return outcome{vanish: true} }}) + + run := e.createRun("on: push\njobs:\n build:\n runs-on: linux\n steps: [{ run: make }]\n", pipeline.Event{}, "") + eventually(t, "flaky worker to start the job", func() bool { + j, err := e.svc.Store.ListRunJobs(context.Background(), run.ID) + return err == nil && j[0].State == "running" + }) + flaky.cancel() // the crash + flaky.done <- <-flaky.done + + healthy := e.fakeWorker("healthy", workerOpts{}) + got, jobs := e.waitRun(run.ID) + if got.State != "succeeded" { + t.Fatalf("run = %s (job: %s, %q)", got.State, jobs["build"].State, jobs["build"].Error) + } + if len(healthy.jobs()) != 1 || jobs["build"].Attempt != 2 { + t.Errorf("healthy ran %d jobs, attempt %d (want 1 and 2)", len(healthy.jobs()), jobs["build"].Attempt) + } + + // The crashed worker restarts and reports the job it lost track of. + restarted := e.start(flaky.id, nil) + select { + case <-restarted.Connected(): + case <-time.After(5 * time.Second): + t.Fatal("restarted worker did not connect") + } + exit := int32(99) + restarted.Send(&opificav1.ConnectRequest{Message: &opificav1.ConnectRequest_JobUpdate{JobUpdate: &opificav1.JobUpdate{ + JobId: jobs["build"].ID.String(), State: opificav1.JobState_JOB_STATE_FAILED, ExitCode: &exit, + }}}) + // It is told to stop the job it no longer holds. + select { + case m := <-restarted.msgs: + if m.GetCancelJob().GetJobId() != jobs["build"].ID.String() { + t.Errorf("restarted worker got %v, want a cancel for the stale job", m) + } + case <-time.After(5 * time.Second): + t.Error("stale update was not answered with a cancel") + } + j, _ := e.svc.Store.GetJob(context.Background(), jobs["build"].ID) + if j.State != "succeeded" || *j.ExitCode != 0 { + t.Errorf("a stale update changed the job: %s exit %d", j.State, *j.ExitCode) + } +} + +// A worker disconnects after being given a job it never started: the job +// is queued again at once, without waiting for the lease and without +// using up an attempt. +func TestUnstartedJobsReturnOnDisconnect(t *testing.T) { + e := newEnv(t) + flaky := e.fakeWorker("flaky", workerOpts{behave: func(*opificav1.Job) outcome { return outcome{silent: true} }}) + run := e.createRun("on: push\njobs:\n build:\n runs-on: linux\n steps: [{ run: make }]\n", pipeline.Event{}, "") + eventually(t, "flaky worker to receive the job", func() bool { return len(flaky.jobs()) == 1 }) + flaky.cancel() + flaky.done <- <-flaky.done + + e.fakeWorker("healthy", workerOpts{}) + got, jobs := e.waitRun(run.ID) + if got.State != "succeeded" || jobs["build"].Attempt != 1 { + t.Errorf("run %s, attempt %d (want succeeded on attempt 1)", got.State, jobs["build"].Attempt) + } +} + +func TestErroredJobsRetryThenFail(t *testing.T) { + e := newEnv(t) + w := e.fakeWorker("broken", workerOpts{behave: func(*opificav1.Job) outcome { + return outcome{state: opificav1.JobState_JOB_STATE_ERRORED} + }}) + run := e.createRun("on: push\njobs:\n build:\n runs-on: linux\n steps: [{ run: make }]\n", pipeline.Event{}, "") + got, jobs := e.waitRun(run.ID) + if got.State != "failed" || jobs["build"].State != "failed" { + t.Errorf("run = %s, job = %s", got.State, jobs["build"].State) + } + if n := len(w.jobs()); n != 3 || jobs["build"].Attempt != 3 { + t.Errorf("job tried %d times (attempt %d), want 3", n, jobs["build"].Attempt) + } +} + +func TestCancelRun(t *testing.T) { + e := newEnv(t) + w := e.fakeWorker("linux", workerOpts{behave: func(*opificav1.Job) outcome { return outcome{hang: true} }}) + run := e.createRun(` +on: push +jobs: + long: + runs-on: linux + steps: [{ run: sleep 3600 }] + after: + needs: long + runs-on: linux + steps: [{ run: echo done }] +`, pipeline.Event{}, "") + eventually(t, "the long job to start", func() bool { + j, err := e.svc.Store.ListRunJobs(context.Background(), run.ID) + return err == nil && slices.ContainsFunc(j, func(j db.Job) bool { return j.State == "running" }) + }) + + if err := e.sched.CancelRun(context.Background(), run.ID); err != nil { + t.Fatal(err) + } + got, jobs := e.waitRun(run.ID) + if got.State != "cancelled" || jobs["long"].State != "cancelled" || jobs["after"].State != "cancelled" { + t.Errorf("run %s, long %s, after %s", got.State, jobs["long"].State, jobs["after"].State) + } + if len(w.jobs()) != 1 { + t.Errorf("worker got %d jobs after cancel", len(w.jobs())) + } + if err := e.sched.CancelRun(context.Background(), run.ID); err == nil { + t.Error("cancelling a finished run succeeded") + } +} + +// Two controllers share one database: every job runs exactly once. +func TestTwoControllersNeverDoubleAssign(t *testing.T) { + e := newEnv(t) + second := scheduler.New(e.svc.Store, e.svc.Sessions(), quiet(), scheduler.Config{Lease: 400 * time.Millisecond, Interval: 5 * time.Millisecond}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go second.Run(ctx) + + var workers []*fakeWorker + for i := range 4 { + workers = append(workers, e.fakeWorker("w"+string(rune('a'+i)), workerOpts{maxJobs: 3})) + } + run := e.createRun(` +on: push +jobs: + t: + matrix: { a: ["1","2","3","4","5"], b: ["1","2","3","4"] } + runs-on: linux + steps: [{ run: make }] +`, pipeline.Event{}, "") + got, jobs := e.waitRun(run.ID) + if got.State != "succeeded" { + t.Fatalf("run = %s", got.State) + } + seen := map[string]int{} + for _, w := range workers { + for _, j := range w.jobs() { + seen[j.GetId()]++ + } + } + for k, j := range jobs { + if seen[j.ID.String()] != 1 { + t.Errorf("%s ran %d times", k, seen[j.ID.String()]) + } + } +} + +func TestCreateRunRejects(t *testing.T) { + e := newEnv(t) + ctx := context.Background() + ev := pipeline.Event{Type: pipeline.EventPush, Branch: "develop", SHA: "abc"} + for name, tc := range map[string]struct { + src string + want error + }{ + "not triggered": {"on: { push: { branches: [main] } }\njobs:\n a:\n runs-on: linux\n steps: [{run: x}]\n", scheduler.ErrNotTriggered}, + "all skipped": {"on: push\njobs:\n a:\n if: event.tag != \"\"\n runs-on: linux\n steps: [{run: x}]\n", scheduler.ErrNothingToRun}, + "invalid": {"on: push\njobs: {}\n", nil}, + } { + _, err := e.sched.CreateRun(ctx, scheduler.RunRequest{Repo: "r", Pool: "trusted", Event: ev, Pipeline: []byte(tc.src)}) + if err == nil || tc.want != nil && err != tc.want { + t.Errorf("%s: err = %v, want %v", name, err, tc.want) + } + } + runs, _ := e.svc.Store.ListRuns(ctx, 10) + if len(runs) != 0 { + t.Errorf("%d runs stored from rejected requests", len(runs)) + } +} + +// `controller run cancel` runs in its own process with no worker sessions. +// The serving controller must still stop the job on its worker. +func TestCancelFromAnotherProcess(t *testing.T) { + e := newEnv(t) + e.fakeWorker("linux", workerOpts{behave: func(*opificav1.Job) outcome { return outcome{hang: true} }}) + run := e.createRun("on: push\njobs:\n long:\n runs-on: linux\n steps: [{ run: sleep 3600 }]\n", pipeline.Event{}, "") + eventually(t, "the job to start", func() bool { + j, err := e.svc.Store.ListRunJobs(context.Background(), run.ID) + return err == nil && len(j) == 1 && j[0].State == "running" + }) + + cli := scheduler.New(e.svc.Store, workers.NewRegistry(), quiet(), scheduler.Config{}) + if err := cli.CancelRun(context.Background(), run.ID); err != nil { + t.Fatal(err) + } + got, jobs := e.waitRun(run.ID) + if got.State != "cancelled" || jobs["long"].State != "cancelled" { + t.Errorf("run %s, job %s", got.State, jobs["long"].State) + } +} diff --git a/internal/server/server.go b/internal/server/server.go index 9a79dba..b3297fc 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -61,6 +61,9 @@ func serve(ctx context.Context, ln net.Listener, h http.Handler, tlsCfg *tls.Con Protocols: protocols, TLSConfig: tlsCfg, ReadHeaderTimeout: 10 * time.Second, + // TLS handshake failures from bad clients are routine; keep them out + // of the default logger. + ErrorLog: slog.NewLogLogger(log.Handler(), slog.LevelDebug), } errc := make(chan error, 1) diff --git a/pkg/worker/client.go b/pkg/worker/client.go index c3c09b8..bcd40fd 100644 --- a/pkg/worker/client.go +++ b/pkg/worker/client.go @@ -151,6 +151,10 @@ func httpClient(cfg *tls.Config) *http.Client { TLSClientConfig: cfg, ForceAttemptHTTP2: true, IdleConnTimeout: 90 * time.Second, + // Pings detect a controller that vanished without closing the + // connection (network cut, crash), which would otherwise leave the + // stream waiting forever. + HTTP2: &http.HTTP2Config{SendPingTimeout: 15 * time.Second, PingTimeout: 10 * time.Second}, }} } @@ -306,10 +310,22 @@ func (c *Client) session(ctx context.Context) error { c.stream = nil c.ready = make(chan struct{}) c.mu.Unlock() - stream.CloseRequest() stream.CloseResponse() }() + // Cancelling the context doesn't interrupt a Receive that is waiting + // for data. Closing our side of the stream does: the controller sees + // the end of the stream and ends its side too. + go func() { + <-ctx.Done() + c.mu.Lock() + defer c.mu.Unlock() + if c.stream == stream { + c.stream = nil + } + stream.CloseRequest() + }() + go func() { t := time.NewTicker(c.cfg.HeartbeatInterval) defer t.Stop()