task-orchestrator - A Container Orchestrator in Go

Overview

task-orchestrator is a container orchestrator written in Go, in the spirit of Kubernetes, Nomad, and Mesos. A manager accepts tasks over HTTP and schedules each one onto a pool of worker nodes; every worker runs its tasks as Docker containers and reports health and resource stats back; and the manager runs background loops that reconcile the cluster toward the desired state. A single orchestrator binary runs the manager, runs a worker, or drives tasks from the command line.

It began as a Cube-style ("Build an Orchestrator in Go") system but was re-architected around a few opinions: push all decisions into pure, testable functions; make illegal states unrepresentable in the type system; and write down an explicit concurrency model instead of hoping -race catches everything. The interesting parts aren't the happy path — they're what happens under concurrency and partial failure.

Tech stack: Go 1.26, the Docker Engine API (moby/moby/client), spf13/cobra for the CLI, go-chi/chi for the manager and worker HTTP APIs, go.etcd.io/bbolt for persistent task storage, gopsutil for worker stats, and Go generics throughout the store and queue.

Architecture

One binary, three roles

The main package is a three-line front door — all wiring lives in internal/cli, where cobra assembles the command tree. worker and manager are long-running foreground servers; task groups the client verbs, with run/stop/status promoted to top-level shortcuts.

cmd.AddCommand(
	workerCmd(),
	managerCmd(),
	taskCmd(),
	newStartCommand("run"),
	newStopCommand("stop"),
	newStatusCommand("status"),
	statsCmd(),
)
return cmd

Signal handling is centralized in a single shutdown channel that every long-running command selects on. The manager command launches its three background loops and the HTTP API against that one done channel, so a Ctrl-C fans out to all of them at once:

done := waitForShutdown(logger)
go m.UpdateTasks(done)
go m.DoHealthChecks(done)
go m.ProcessTasks(done)

api := manager.NewAPI(host, port, m, mLogger)
// Blocks until done is closed, then shuts the HTTP server down.
api.Start(done)

The manager and its reconcile loop

The Manager holds all mutable routing state — the two worker/task maps, the pending queue, the task and event stores, the worker list, and the scheduler — behind a single RWMutex. The loop intervals are struct fields so tests can drive the loops in milliseconds, and it carries its own http.Client purely to enforce a timeout on outbound worker calls (http.DefaultClient has none).

type Manager struct {
	mu            sync.RWMutex
	Pending       *queue.Queue[task.TaskEvent]
	TaskDb        store.Store[task.Task]
	EventDb       store.Store[task.TaskEvent]
	WorkerTaskMap map[string][]uuid.UUID
	TaskWorkerMap map[uuid.UUID]string
	WorkerNodes   []node.Node
	Scheduler     scheduler.Scheduler
	logger        *slog.Logger

	httpClient          *http.Client
	updateInterval      time.Duration
	processInterval     time.Duration
	healthCheckInterval time.Duration
}

The heart of the system is the "send work" path. It dequeues a pending event, selects a worker, and then does something subtle: it commits the assignment before sending the HTTP request, and rolls it back with a compensating unassignTask if the worker turns out to be unreachable. The comment explains why the ordering matters.

selected, err := m.SelectWorker(taskEvent.Task)
if err != nil {
	m.logger.Error("selecting worker failed", "err", err, "taskID", t.ID)
	return
}

workerName := selected.Name
m.assignTask(workerName, taskEvent)

newTask, err := m.postTaskEvent("manager.SendWork", workerName, selected.API, taskEvent)
if err != nil {
	if errors.Is(err, errWorkerUnreachable) {
		// The assignment was committed before the send, so it has to be
		// rolled back: a requeued task that is still routed to a worker
		// looks like a duplicate on the next pass and gets dropped.
		m.unassignTask(workerName, t.ID)
		m.logger.Warn("worker unreachable, requeueing task", "err", err, "taskID", t.ID)
		m.enqueueEvent(taskEvent)
		return
	}
	m.logger.Error("sending task to worker failed", "err", err, "taskID", t.ID)
	return
}

Pushing decisions into pure functions

Rather than branch inline, the health-check loop dispatches on a pure decision helper. The loop itself only does I/O; the policy is a function with no receiver and no side effects, which makes it trivially testable without a live worker.

func decideHealthAction(t task.Task) healthAction {
	if t.State != task.Running && t.State != task.Failed {
		return healthSkipNotEligible
	}
	if t.RestartCount >= TaskRestartMax {
		return healthSkipRestartMax
	}
	if t.State == task.Running {
		return healthActionCheck
	}
	return healthActionRestart
}

The same file holds mergeTaskUpdate — where the manager owns RestartCount, the worker owns the runtime fields, and a report that walks an invalid state edge is rejected outright — and healthCheckURL, which composes the worker's scheme and host with the container's published port to avoid producing a doubled http://localhost:3001:32768/health.

The scheduler interface

The Scheduler interface deliberately pairs a node with its score. This replaced an earlier map[string]float64 where an absent key read as 0.0 — and because costs are lowest-wins, an unreachable, unscored node would beat every healthy one. Making "a candidate with no score" unrepresentable is a type-driven correctness fix.

type ScoredNode struct {
	Node  node.Node
	Score float64
}

type Scheduler interface {
	SelectCandidateNodes(t task.Task, nodes []node.Node) []node.Node
	// Score returns one ScoredNode per node it could score. A scheduler may
	// drop nodes it could not reach, so the result is not necessarily the
	// same length as its input.
	Score(t task.Task, nodes []node.Node) ([]ScoredNode, error)
	Pick(scored []ScoredNode) node.Node
}

Round-robin is the simple strategy: it carries a single cursor guarded by its own mutex — which is precisely what lets the manager schedule without holding m.mu across network I/O.

The marginal-cost scheduler

The interesting strategy is an E-PVM-derived marginal-cost scheduler. It fans out one goroutine per candidate to fetch live stats, tolerates unreachable nodes, and only fails when nothing could be scored. Crucially, each goroutine writes into its own slice index — no lock — because distinct slice slots are disjoint memory.

func (e *MarginalCostScheduler) Score(t task.Task, candidates []node.Node) ([]ScoredNode, error) {
	op := "scheduler.MarginalCostScheduler.Score"
	// results[i] is meaningful only when scored[i] is true.
	results := make([]ScoredNode, len(candidates))
	scored := make([]bool, len(candidates))

	var wg sync.WaitGroup
	errsStream := make(chan error, len(candidates))
	for i, n := range candidates {
		wg.Go(func() {
			stats, err := n.GetStats()
			if err != nil {
				errsStream <- Wrap(op, "scoring node "+n.Name, err)
				return
			}
			load := nodeLoad{
				MemUsed:   float64(stats.MemUsedKb()) + float64(n.MemoryAllocated),
				MemTotal:  float64(stats.MemTotalKb()),
				CPULoad:   e.cpuLoad(n, stats),
				TaskCount: float64(stats.TaskCount),
			}
			score, err := nodeScore(load, float64(t.Memory))
			if err != nil {
				errsStream <- E(op, "scoring node "+n.Name, err)
				return
			}
			results[i] = ScoredNode{Node: n, Score: score}
			scored[i] = true
		})
	}
	wg.Wait()
	// ...collect scored[i] entries; error only if none scored.
}

The cost formula itself is a pure function. Memory and job-count are true marginal costs — cost-after minus cost-before — while CPU enters only as a load penalty, because a task.Task carries no CPU request. The exponent base LIEB (Lieb's square-ice constant, ≈1.5396) makes the cost curve convex, so the same load increment hurts more on an already-busy node.

func powerCost(load float64) float64 { return math.Pow(LIEB, load) }

func nodeScore(l nodeLoad, taskMemory float64) (float64, error) {
	if l.MemTotal <= 0 {
		// Without the guard the node scores NaN, which compares false against
		// every other score and so neither wins nor loses.
		return 0, fmt.Errorf("total memory must be positive, got %v", l.MemTotal)
	}
	memNow := calculateLoad(l.MemUsed, l.MemTotal)
	memNew := calculateLoad(l.MemUsed+taskMemory, l.MemTotal)

	memCost := powerCost(memNew) - powerCost(memNow)
	jobCost := powerCost((l.TaskCount+1)/maxJobsPerNode) - powerCost(l.TaskCount/maxJobsPerNode)
	cpuCost := powerCost(l.CPULoad) - 1
	return memCost + jobCost + cpuCost, nil
}

The worker and the Docker runtime

The worker runs assigned tasks as Docker containers through the Docker Engine API, using the newer split moby/moby/client package (not the classic docker/docker). StartTask pulls-and-runs, inspects the container to capture its published host ports, and then commits the new state through an upsertTask callback that touches only the fields this operation owns.

w.upsertTask(t, func(p *task.Task) {
	p.StartTime = startTime
	p.ContainerID = result.ContainerId
	p.HostPorts = hostPorts
	p.State = task.Running
})

The worker's HTTP API is intentionally asynchronous. A DELETE /tasks doesn't stop the container inline — it enqueues a Completed copy of the task and returns 204 No Content, letting the worker's own run loop pick up the stop. Receipt is decoupled from execution.

taskCopy := taskToStop
taskCopy.State = task.Completed
a.Worker.AddTask(taskCopy)
a.logger.Info("task added to stop container", ...)
// 204 means no content — writing a body here would violate the status.
w.WriteHeader(http.StatusNoContent)

The task state machine

Task state is a hand-rolled machine: an enum plus an explicit adjacency-map transition table. Completed is terminal — an empty transition set — and that terminality is load-bearing across the whole system.

func NewStateMachine() StateMachine {
	return StateMachine{
		machine: map[State][]State{
			Pending:   {Scheduled},
			Scheduled: {Scheduled, Running, Failed},
			Running:   {Running, Completed, Failed, Scheduled},
			Completed: {},
			Failed:    {Scheduled},
		},
	}
}

func (s *StateMachine) IsValidTransition(cur, next State) bool {
	if validStates, exists := s.machine[cur]; exists && slices.Contains(validStates, next) {
		return true
	}
	return false
}

Because Completed accepts no transitions, once the manager marks a stopped task Completed, a stale Running report arriving late from a worker can never walk it back to life. The self-transitions (Running → Running, Scheduled → Scheduled) model the steady state.

Generics to kill a class of panics

The store and the queue are both generic, and the reason is stated plainly in the code: the type parameter fixes the element type at construction, so callers can't write a wrong type assertion — a class of runtime panic the codebase had actually hit.

type Store[T any] interface {
	Put(key string, value T) error
	Get(key string) (T, error)
	List() ([]T, error)
	Count() (int, error)
	Close() error
}

There are two implementations behind that interface: a plain in-memory map (deliberately not internally synchronized, because callers already lock around multi-step invariants) and a persistent store backed by go.etcd.io/bbolt — the embedded key/value engine that powers etcd — which JSON-encodes values into a named bucket inside bbolt Update/View transactions. Swapping between the two is a single IN_MEMORY vs PERSISTENT flag, so the same manager runs disposably in a test or durably against a bbolt file on disk. The generic queue similarly wraps an interface{}-based collection so that exactly one type assertion exists in the entire codebase:

func (q *Queue[T]) Dequeue() (T, bool) {
	v := q.inner.Dequeue()
	if v == nil {
		var zero T
		return zero, false
	}
	return v.(T), true // the ONLY type assertion in the codebase
}

Difficult Parts

Data races are the easy half; lost updates are the hard half

The project's sharpest idea is that per-access locking stops data races but not lost updates. Consider two goroutines racing over one task, where every single access is correctly locked:

updateTasks reads task X          {State: Running, ContainerID: abc}
updateTasks calls Inspect         ... two seconds of I/O ...
RunTask stops X, writes           {State: Completed, FinishTime: now}
updateTasks writes its copy back  {State: Running, FinishTime: zero}

go test -race reports nothing — yet task X is Running again with no container behind it. The fix is upsertTask: re-read the record inside the lock and apply a callback that mutates only the caller-owned fields, so a slow read can never clobber a concurrent write. The whole concurrency model is written down in docs/concurrency-and-state.md, which classifies every state helper into one of four categories (pure read, whole-value write, read-modify-write without I/O, and read-modify-write with I/O) and enforces two rules: never hold the mutex across I/O, and nothing outside state.go touches guarded state directly.

Slices fan out lock-free; maps crash

The marginal-cost scheduler writes goroutine results into a slice by index with no lock — and that's safe. The tempting refactor to a map[string]ScoredNode keyed by node name is not: a Go map shares its count, buckets, and growth state across all keys, so concurrent writes to different keys still trip the uncatchable fatal error: concurrent map writes. The doc's one-liner: "unique keys buy nothing; unique indices buy everything."

Rolling back a half-finished assignment

Scheduling straddles a network boundary, and the manager has to pick an order: assign-then-send, or send-then-assign. It chooses assign-first so the routing tables never lag reality — but that means a send failure leaves a phantom assignment behind. The unassignTask rollback exists solely to close that window, because a requeued task that still looks assigned reads as a duplicate on the next reconcile pass and gets silently dropped.

A restart budget that has to survive a crash

The health loop restarts failed tasks, but a crash-looping task must eventually be given up on. That budget only works if the restart count is incremented in a single critical section and made durable before the restart is dispatched — otherwise a lost write means the counter never advances and the task restarts forever. beginRestart does the read, the increment, and the durable Put under one lock before returning.

Testing what -race can't see

The test strategy leans on go test -race, then adds named tests for exactly the failure modes the race detector is blind to: TestScoreDoesNotDeadlockWhenEveryNodeFails, TestUpsertTaskDoesNotResurrectAStoppedTask, and TestPickNeverChoosesAnUnreachableNode. Each one encodes a bug that was real at some point — the deadlock when every node fails, the lost update that resurrects a stopped task, and the zero-score node that wins by default.