cloud/bmaas/bmdb: init
This adds the initial Bare Metal Database structure. This change focuses
on a session/work mechanism which is the foundation on which we will
build worker components. It allows lease-like mechanics on machines,
letting us not have to use 'standard' work queues in the BMaaS project.
Change-Id: I42c3f4384c64fd90dbeab8ff9652a6f611be81d4
Reviewed-on: https://review.monogon.dev/c/monogon/+/953
Tested-by: Jenkins CI
Reviewed-by: Lorenz Brun <lorenz@monogon.tech>
diff --git a/cloud/bmaas/bmdb/BUILD.bazel b/cloud/bmaas/bmdb/BUILD.bazel
new file mode 100644
index 0000000..957a884
--- /dev/null
+++ b/cloud/bmaas/bmdb/BUILD.bazel
@@ -0,0 +1,32 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+ name = "bmdb",
+ srcs = [
+ "bmdb.go",
+ "sessions.go",
+ ],
+ importpath = "source.monogon.dev/cloud/bmaas/bmdb",
+ visibility = ["//visibility:public"],
+ deps = [
+ "//cloud/bmaas/bmdb/model",
+ "//cloud/lib/component",
+ "@com_github_cockroachdb_cockroach_go_v2//crdb",
+ "@com_github_google_uuid//:uuid",
+ "@com_github_lib_pq//:pq",
+ "@io_k8s_klog_v2//:klog",
+ ],
+)
+
+go_test(
+ name = "bmdb_test",
+ srcs = ["sessions_test.go"],
+ data = [
+ "@cockroach",
+ ],
+ embed = [":bmdb"],
+ deps = [
+ "//cloud/bmaas/bmdb/model",
+ "//cloud/lib/component",
+ ],
+)
diff --git a/cloud/bmaas/bmdb/bmdb.go b/cloud/bmaas/bmdb/bmdb.go
new file mode 100644
index 0000000..d56e083
--- /dev/null
+++ b/cloud/bmaas/bmdb/bmdb.go
@@ -0,0 +1,8 @@
+// Package bmdb implements a connector to the Bare Metal Database, which is the
+// main data store backing information about bare metal machines.
+//
+// All components of the BMaaS project connect directly to the underlying
+// CockroachDB database storing this data via this library. In the future, this
+// library might turn into a shim which instead connects to a coordinator
+// service over gRPC.
+package bmdb
diff --git a/cloud/bmaas/bmdb/model/BUILD.bazel b/cloud/bmaas/bmdb/model/BUILD.bazel
new file mode 100644
index 0000000..7a0f122
--- /dev/null
+++ b/cloud/bmaas/bmdb/model/BUILD.bazel
@@ -0,0 +1,28 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("//build/sqlc:sqlc.bzl", "sqlc_go_library")
+
+sqlc_go_library(
+ name = "sqlc_model",
+ dialect = "cockroachdb",
+ importpath = "source.monogon.dev/cloud/bmaas/bmdb/model",
+ migrations = glob(["migrations/*sql"]),
+ queries = [
+ "queries.sql",
+ ],
+)
+
+go_library(
+ name = "model",
+ srcs = ["migrations.go"],
+ embed = [
+ ":sqlc_model", # keep
+ ],
+ embedsrcs = glob(["migrations/*sql"]),
+ importpath = "source.monogon.dev/cloud/bmaas/bmdb/model",
+ visibility = ["//visibility:public"],
+ deps = [
+ "@com_github_golang_migrate_migrate_v4//source",
+ "@com_github_golang_migrate_migrate_v4//source/iofs",
+ "@com_github_google_uuid//:uuid", # keep
+ ],
+)
diff --git a/cloud/bmaas/bmdb/model/migrations.go b/cloud/bmaas/bmdb/model/migrations.go
new file mode 100644
index 0000000..2c07768
--- /dev/null
+++ b/cloud/bmaas/bmdb/model/migrations.go
@@ -0,0 +1,15 @@
+package model
+
+import (
+ "embed"
+
+ "github.com/golang-migrate/migrate/v4/source"
+ "github.com/golang-migrate/migrate/v4/source/iofs"
+)
+
+//go:embed migrations/*.sql
+var migrationData embed.FS
+
+func MigrationsSource() (source.Driver, error) {
+ return iofs.New(migrationData, "migrations")
+}
diff --git a/cloud/bmaas/bmdb/model/migrations/1662136250_initial.down.sql b/cloud/bmaas/bmdb/model/migrations/1662136250_initial.down.sql
new file mode 100644
index 0000000..6a85991
--- /dev/null
+++ b/cloud/bmaas/bmdb/model/migrations/1662136250_initial.down.sql
@@ -0,0 +1,6 @@
+DROP TABLE machine_agent_report;
+DROP TABLE machine_agent_installed;
+DROP TABLE machine_provided;
+DROP TABLE work;
+DROP TABLE sessions;
+DROP TABLE machines;
diff --git a/cloud/bmaas/bmdb/model/migrations/1662136250_initial.up.sql b/cloud/bmaas/bmdb/model/migrations/1662136250_initial.up.sql
new file mode 100644
index 0000000..b31324f
--- /dev/null
+++ b/cloud/bmaas/bmdb/model/migrations/1662136250_initial.up.sql
@@ -0,0 +1,73 @@
+CREATE TABLE machines (
+ machine_id UUID NOT NULL DEFAULT gen_random_uuid() PRIMARY KEY,
+ machine_created_at TIMESTAMPTZ NOT NULL
+);
+
+
+-- Sessions are maintained by components as they work on the rest of the machine
+-- database. Once a session is created, it must be maintained by its owning
+-- component by repeatedly 'poking' it, ie. updating the heartbeat_deadline
+-- value to be some point in the future.
+--
+-- TODO: garbage collect old sessions.
+CREATE TABLE sessions (
+ session_id UUID NOT NULL DEFAULT gen_random_uuid() PRIMARY KEY,
+ -- Name of component which created this session. Human-readable.
+ session_component_name STRING NOT NULL,
+ -- Node name, hostname:port, pod name, whatever. Something to tell where
+ -- a particular component is running. Human-readable.
+ session_runtime_info STRING NOT NULL,
+ -- Time at which this session was created.
+ session_created_at TIMESTAMPTZ NOT NULL,
+ -- Number of seconds by which session_deadline (counting from now())
+ -- is bumped up every time the session is poked.
+ session_interval_seconds INT NOT NULL,
+ -- Deadline after which this session should not be considered valid anymore.
+ session_deadline TIMESTAMPTZ NOT NULL
+);
+
+CREATE TYPE process AS ENUM (
+ -- Reserved for unit tests.
+ 'UnitTest1',
+ 'UnitTest2'
+);
+
+-- Work items map a session to work performed on a machine. Multiple work items
+-- can exist per session, and thus, a session can back multiple items of work
+-- acting on multiple machines. These are optionally created by components to
+-- indicate some long-running process being performed on a machine, and will
+-- lock out the same process from being run simultaneously, eg. in a
+-- concurrently running instance of the same component.
+CREATE TABLE work (
+ -- Machine that this work is being performed on. Prevent deleting machines
+ -- that have active work on them.
+ machine_id UUID NOT NULL REFERENCES machines(machine_id) ON DELETE RESTRICT,
+ -- Session that this work item is tied to. If the session expires, so does
+ -- the work item.
+ session_id UUID NOT NULL REFERENCES sessions(session_id) ON DELETE CASCADE,
+ -- Human-readable process name.
+ process process NOT NULL,
+ UNIQUE (machine_id, process),
+ CONSTRAINT "primary" PRIMARY KEY (machine_id, session_id, process)
+);
+
+-- The following three tables are for illustrative purposes only.
+
+CREATE TABLE machine_provided (
+ machine_id UUID NOT NULL REFERENCES machines(machine_id) ON DELETE CASCADE,
+ provider STRING NOT NULL,
+ provider_id STRING NOT NULL,
+ CONSTRAINT "primary" PRIMARY KEY (machine_id)
+);
+
+CREATE TABLE machine_agent_installed (
+ machine_id UUID NOT NULL REFERENCES machines(machine_id) ON DELETE CASCADE,
+ CONSTRAINT "primary" PRIMARY KEY (machine_id)
+);
+
+CREATE TABLE machine_agent_report (
+ machine_id UUID NOT NULL REFERENCES machines(machine_id) ON DELETE CASCADE,
+ shape_cpu_count INT NOT NULL,
+ shape_memory_megabytes INT NOT NULL,
+ CONSTRAINT "primary" PRIMARY KEY (machine_id)
+);
diff --git a/cloud/bmaas/bmdb/model/queries.sql b/cloud/bmaas/bmdb/model/queries.sql
new file mode 100644
index 0000000..ee3f618
--- /dev/null
+++ b/cloud/bmaas/bmdb/model/queries.sql
@@ -0,0 +1,64 @@
+-- name: NewMachine :one
+INSERT INTO machines (
+ machine_created_at
+) VALUES (
+ now()
+)
+RETURNING *;
+
+-- name: NewSession :one
+INSERT INTO sessions (
+ session_component_name, session_runtime_info, session_created_at, session_interval_seconds, session_deadline
+) VALUES (
+ $1, $2, now(), $3, (now() + $3 * interval '1 second')
+)
+RETURNING *;
+
+-- name: SessionPoke :exec
+-- Update a given session with a new deadline. Must be called in the same
+-- transaction as SessionCheck to ensure the session is still alive.
+UPDATE sessions
+SET session_deadline = now() + session_interval_seconds * interval '1 second'
+WHERE session_id = $1;
+
+-- name: SessionCheck :many
+-- SessionCheck returns a session by ID if that session is still valid (ie. its
+-- deadline hasn't expired).
+SELECT *
+FROM sessions
+WHERE session_id = $1
+AND session_deadline > now();
+
+-- name: StartWork :exec
+INSERT INTO work (
+ machine_id, session_id, process
+) VALUES (
+ $1, $2, $3
+);
+
+-- name: FinishWork :exec
+DELETE FROM work
+WHERE machine_id = $1
+ AND session_id = $2
+ AND process = $3;
+
+-- Example tag processing queries follow.
+
+-- name: MachineAddProvided :exec
+INSERT INTO machine_provided (
+ machine_id, provider, provider_id
+) VALUES (
+ $1, $2, $3
+);
+
+-- name: GetMachinesNeedingInstall :many
+SELECT
+ machine_provided.*
+FROM machines
+INNER JOIN machine_provided ON machines.machine_id = machine_provided.machine_id
+LEFT JOIN work ON machines.machine_id = work.machine_id AND work.process = 'NecromancerInstall'
+LEFT JOIN machine_agent_installed ON machines.machine_id = machine_agent_installed.machine_id
+WHERE machine_agent_installed.machine_id IS NULL
+ AND work.machine_id IS NULL
+LIMIT $1;
+
diff --git a/cloud/bmaas/bmdb/sessions.go b/cloud/bmaas/bmdb/sessions.go
new file mode 100644
index 0000000..6b19171
--- /dev/null
+++ b/cloud/bmaas/bmdb/sessions.go
@@ -0,0 +1,265 @@
+package bmdb
+
+import (
+ "context"
+ "database/sql"
+ "errors"
+ "fmt"
+ "time"
+
+ "github.com/cockroachdb/cockroach-go/v2/crdb"
+ "github.com/google/uuid"
+ "github.com/lib/pq"
+ "k8s.io/klog/v2"
+
+ "source.monogon.dev/cloud/bmaas/bmdb/model"
+ "source.monogon.dev/cloud/lib/component"
+)
+
+// BMDB is the Bare Metal Database, a common schema to store information about
+// bare metal machines in CockroachDB. This struct is supposed to be
+// embedded/contained by different components that interact with the BMDB, and
+// provides a common interface to BMDB operations to these components.
+//
+// The BMDB provides two mechanisms facilitating a 'reactive work system' being
+// implemented on the bare metal machine data:
+//
+// - Sessions, which are maintained by heartbeats by components and signal the
+// liveness of said components to other components operating on the BMDB. These
+// effectively extend CockroachDB's transactions to be visible as row data. Any
+// session that is not actively being updated by a component can be expired by a
+// component responsible for lease garbage collection.
+// - Work locking, which bases on Sessions and allows long-standing
+// multi-transaction work to be performed on given machines, preventing
+// conflicting work from being performed by other components. As both Work
+// locking and Sessions are plain row data, other components can use SQL queries
+// to exclude machines to act on by constraining SELECT queries to not return
+// machines with some active work being performed on them.
+type BMDB struct {
+ Config
+}
+
+// Config is the configuration of the BMDB connector.
+type Config struct {
+ Database component.CockroachConfig
+
+ // ComponentName is a human-readable name of the component connecting to the
+ // BMDB, and is stored in any Sessions managed by this component's connector.
+ ComponentName string
+ // RuntimeInfo is a human-readable 'runtime information' (eg. software version,
+ // host machine/job information, IP address, etc.) stored alongside the
+ // ComponentName in active Sessions.
+ RuntimeInfo string
+}
+
+// Open creates a new Connection to the BMDB for the calling component. Multiple
+// connections can be opened (although there is no advantage to doing so, as
+// Connections manage an underlying CockroachDB connection pool, which performs
+// required reconnects and connection pooling automatically).
+func (b *BMDB) Open(migrate bool) (*Connection, error) {
+ if migrate {
+ if b.Config.Database.Migrations == nil {
+ klog.Infof("Using default migrations source.")
+ m, err := model.MigrationsSource()
+ if err != nil {
+ klog.Exitf("failed to prepare migrations source: %w", err)
+ }
+ b.Config.Database.Migrations = m
+ }
+ if err := b.Database.MigrateUp(); err != nil {
+ return nil, fmt.Errorf("migration failed: %w", err)
+ }
+ }
+ db, err := b.Database.Connect()
+ if err != nil {
+ return nil, err
+ }
+ return &Connection{
+ bmdb: b,
+ db: db,
+ }, nil
+}
+
+// Connection to the BMDB. Internally, this contains a sql.DB connection pool,
+// so components can (and likely should) reuse Connections as much as possible
+// internally.
+type Connection struct {
+ bmdb *BMDB
+ db *sql.DB
+}
+
+// StartSession creates a new BMDB session which will be maintained in a
+// background goroutine as long as the given context is valid. Each Session is
+// represented by an entry in a sessions table within the BMDB, and subsequent
+// Transact calls emit SQL transactions which depend on that entry still being
+// present and up to date. A garbage collection system (to be implemented) will
+// remove expired sessions from the BMDB, but this mechanism is not necessary
+// for the session expiry mechanism to work.
+//
+// When the session becomes invalid (for example due to network partition),
+// subsequent attempts to call Transact will fail with ErrSessionExpired. This
+// means that the caller within the component is responsible for recreating a
+// new Session if a previously used one expires.
+func (c *Connection) StartSession(ctx context.Context) (*Session, error) {
+ intervalSeconds := 5
+
+ res, err := model.New(c.db).NewSession(ctx, model.NewSessionParams{
+ SessionComponentName: c.bmdb.ComponentName,
+ SessionRuntimeInfo: c.bmdb.RuntimeInfo,
+ SessionIntervalSeconds: int64(intervalSeconds),
+ })
+ if err != nil {
+ return nil, fmt.Errorf("creating session failed: %w", err)
+ }
+
+ klog.Infof("Started session %s", res.SessionID)
+
+ ctx2, ctxC := context.WithCancel(ctx)
+
+ s := &Session{
+ connection: c,
+ interval: time.Duration(intervalSeconds) * time.Second,
+
+ UUID: res.SessionID,
+
+ ctx: ctx2,
+ ctxC: ctxC,
+ }
+ go s.maintainHeartbeat(ctx2)
+ return s, nil
+}
+
+// Session is a session (identified by UUID) that has been started in the BMDB.
+// Its liveness is maintained by a background goroutine, and as long as that
+// session is alive, it can perform transactions and work on the BMDB.
+type Session struct {
+ connection *Connection
+ interval time.Duration
+
+ UUID uuid.UUID
+
+ ctx context.Context
+ ctxC context.CancelFunc
+}
+
+var (
+ // ErrSessionExpired is returned when attempting to Transact or Work on a
+ // Session that has expired or been canceled. Once a Session starts returning
+ // these errors, it must be re-created by another StartSession call, as no other
+ // calls will succeed.
+ ErrSessionExpired = errors.New("session expired")
+ // ErrWorkConflict is returned when attempting to Work on a Session with a
+ // process name that's already performing some work, concurrently, on the
+ // requested machine.
+ ErrWorkConflict = errors.New("conflicting work on machine")
+)
+
+// maintainHeartbeat will attempt to repeatedly poke the session at a frequency
+// twice of that of the minimum frequency mandated by the configured 5-second
+// interval. It will exit if it detects that the session cannot be maintained
+// anymore, canceling the session's internal context and causing future
+// Transact/Work calls to fail.
+func (s *Session) maintainHeartbeat(ctx context.Context) {
+ // Internal deadline, used to check whether we haven't dropped the ball on
+ // performing the updates due to a lot of transient errors.
+ deadline := time.Now().Add(s.interval)
+ for {
+ if ctx.Err() != nil {
+ klog.Infof("Session %s: context over, exiting: %v", s.UUID, ctx.Err())
+ return
+ }
+
+ err := s.Transact(ctx, func(q *model.Queries) error {
+ sessions, err := q.SessionCheck(ctx, s.UUID)
+ if err != nil {
+ return fmt.Errorf("when retrieving session: %w", err)
+ }
+ if len(sessions) < 1 {
+ return ErrSessionExpired
+ }
+ err = q.SessionPoke(ctx, s.UUID)
+ if err != nil {
+ return fmt.Errorf("when poking session: %w", err)
+ }
+ return nil
+ })
+ if err != nil {
+ klog.Errorf("Session %s: update failed: %v", s.UUID, err)
+ if errors.Is(err, ErrSessionExpired) || deadline.After(time.Now()) {
+ // No way to recover.
+ klog.Errorf("Session %s: exiting", s.UUID)
+ s.ctxC()
+ return
+ }
+ // Just retry in a bit. One second seems about right for a 5 second interval.
+ //
+ // TODO(q3k): calculate this based on the configured interval.
+ time.Sleep(time.Second)
+ }
+ // Success. Keep going.
+ deadline = time.Now().Add(s.interval)
+ time.Sleep(s.interval / 2)
+ }
+}
+
+// Transact runs a given function in the context of both a CockroachDB and BMDB
+// transaction, retrying as necessary.
+//
+// Most pure (meaning without side effects outside the database itself) BMDB
+// transactions should be run this way.
+func (s *Session) Transact(ctx context.Context, fn func(q *model.Queries) error) error {
+ return crdb.ExecuteTx(ctx, s.connection.db, nil, func(tx *sql.Tx) error {
+ qtx := model.New(tx)
+ sessions, err := qtx.SessionCheck(ctx, s.UUID)
+ if err != nil {
+ return fmt.Errorf("when retrieving session: %w", err)
+ }
+ if len(sessions) < 1 {
+ return ErrSessionExpired
+ }
+
+ if err := fn(qtx); err != nil {
+ return err
+ }
+
+ return nil
+ })
+}
+
+// Work runs a given function as a work item with a given process name against
+// some identified machine. Not more than one process of a given name can run
+// against a machine concurrently.
+//
+// Most impure (meaning with side effects outside the database itself) BMDB
+// transactions should be run this way.
+func (s *Session) Work(ctx context.Context, machine uuid.UUID, process model.Process, fn func() error) error {
+ err := model.New(s.connection.db).StartWork(ctx, model.StartWorkParams{
+ MachineID: machine,
+ SessionID: s.UUID,
+ Process: process,
+ })
+ if err != nil {
+ var perr *pq.Error
+ if errors.As(err, &perr) && perr.Code == "23505" {
+ return ErrWorkConflict
+ }
+ return fmt.Errorf("could not start work: %w", err)
+ }
+ klog.Infof("Started work: %q on machine %s, session %s", process, machine, s.UUID)
+
+ defer func() {
+ err := model.New(s.connection.db).FinishWork(s.ctx, model.FinishWorkParams{
+ MachineID: machine,
+ SessionID: s.UUID,
+ Process: process,
+ })
+ klog.Errorf("Finished work: %q on machine %s, session %s", process, machine, s.UUID)
+ if err != nil && !errors.Is(err, s.ctx.Err()) {
+ klog.Errorf("Failed to finish work: %v", err)
+ klog.Errorf("Closing session out of an abundance of caution")
+ s.ctxC()
+ }
+ }()
+
+ return fn()
+}
diff --git a/cloud/bmaas/bmdb/sessions_test.go b/cloud/bmaas/bmdb/sessions_test.go
new file mode 100644
index 0000000..0018109
--- /dev/null
+++ b/cloud/bmaas/bmdb/sessions_test.go
@@ -0,0 +1,172 @@
+package bmdb
+
+import (
+ "context"
+ "errors"
+ "testing"
+ "time"
+
+ "source.monogon.dev/cloud/bmaas/bmdb/model"
+ "source.monogon.dev/cloud/lib/component"
+)
+
+func dut() *BMDB {
+ return &BMDB{
+ Config: Config{
+ Database: component.CockroachConfig{
+ InMemory: true,
+ },
+ },
+ }
+}
+
+// TestSessionExpiry exercises the session heartbeat logic, making sure that if
+// a session stops being maintained subsequent Transact calls will fail.
+func TestSessionExpiry(t *testing.T) {
+ b := dut()
+ conn, err := b.Open(true)
+ if err != nil {
+ t.Fatalf("Open failed: %v", err)
+ }
+
+ ctx, ctxC := context.WithCancel(context.Background())
+ defer ctxC()
+
+ session, err := conn.StartSession(ctx)
+ if err != nil {
+ t.Fatalf("Starting session failed: %v", err)
+ }
+
+ // A transaction in a brand-new session should work.
+ var machine model.Machine
+ err = session.Transact(ctx, func(q *model.Queries) error {
+ machine, err = q.NewMachine(ctx)
+ return err
+ })
+ if err != nil {
+ t.Fatalf("First transaction failed: %v", err)
+ }
+
+ time.Sleep(6 * time.Second)
+
+ // A transaction after the 5-second session interval should continue to work.
+ err = session.Transact(ctx, func(q *model.Queries) error {
+ _, err = q.NewMachine(ctx)
+ return err
+ })
+ if err != nil {
+ t.Fatalf("Second transaction failed: %v", err)
+ }
+
+ // A transaction after the 5-second session interval should fail if we don't
+ // maintain its heartbeat.
+ session.ctxC()
+ time.Sleep(6 * time.Second)
+
+ err = session.Transact(ctx, func(q *model.Queries) error {
+ return q.MachineAddProvided(ctx, model.MachineAddProvidedParams{
+ MachineID: machine.MachineID,
+ Provider: "foo",
+ ProviderID: "bar",
+ })
+ })
+ if !errors.Is(err, ErrSessionExpired) {
+ t.Fatalf("Second transaction should've failed due to expired session, got %v", err)
+ }
+
+}
+
+// TestWork exercises the per-{process,machine} mutual exclusion mechanism of
+// Work items.
+func TestWork(t *testing.T) {
+ b := dut()
+ conn, err := b.Open(true)
+ if err != nil {
+ t.Fatalf("Open failed: %v", err)
+ }
+
+ ctx, ctxC := context.WithCancel(context.Background())
+ defer ctxC()
+
+ // Start two session for testing.
+ session1, err := conn.StartSession(ctx)
+ if err != nil {
+ t.Fatalf("Starting session failed: %v", err)
+ }
+ session2, err := conn.StartSession(ctx)
+ if err != nil {
+ t.Fatalf("Starting session failed: %v", err)
+ }
+
+ var machine model.Machine
+ err = session1.Transact(ctx, func(q *model.Queries) error {
+ machine, err = q.NewMachine(ctx)
+ return err
+ })
+ if err != nil {
+ t.Fatalf("Creating machine failed: %v", err)
+ }
+
+ // Create a subcontext for a long-term work request. We'll cancel it later as
+ // part of the test.
+ ctxB, ctxBC := context.WithCancel(ctx)
+ defer ctxBC()
+ // Start work which will block forever. We have to go rendezvous through a
+ // channel to make sure the work actually starts.
+ started := make(chan error)
+ done := make(chan error, 1)
+ go func() {
+ err := session1.Work(ctxB, machine.MachineID, model.ProcessUnitTest1, func() error {
+ started <- nil
+ <-ctxB.Done()
+ return ctxB.Err()
+ })
+ done <- err
+ if err != nil {
+ started <- err
+ }
+ }()
+ err = <-started
+ if err != nil {
+ t.Fatalf("Starting first work failed: %v", err)
+ }
+
+ // Starting more work on the same machine but a different process should still
+ // be allowed.
+ for _, session := range []*Session{session1, session2} {
+ err = session.Work(ctxB, machine.MachineID, model.ProcessUnitTest2, func() error {
+ return nil
+ })
+ if err != nil {
+ t.Errorf("Could not run concurrent process on machine: %v", err)
+ }
+ }
+
+ // However, starting work with the same process on the same machine should
+ // fail.
+ for _, session := range []*Session{session1, session2} {
+ err = session.Work(ctxB, machine.MachineID, model.ProcessUnitTest1, func() error {
+ return nil
+ })
+ if !errors.Is(err, ErrWorkConflict) {
+ t.Errorf("Concurrent work with same process should've been forbidden, got %v", err)
+ }
+ }
+
+ // Now, cancel the first long-running request and wait for it to return.
+ ctxBC()
+ err = <-done
+ if !errors.Is(err, ctxB.Err()) {
+ t.Fatalf("First work item should've failed with %v, got %v", ctxB.Err(), err)
+ }
+
+ // We should now be able to perform 'test1' work again against this machine.
+ for _, session := range []*Session{session1, session2} {
+ err = session.Work(ctx, machine.MachineID, model.ProcessUnitTest1, func() error {
+ return nil
+ })
+ if err != nil {
+ t.Errorf("Could not run work against machine: %v", err)
+ }
+ }
+}