fix(zebra): isolate dispatching of jobs per os_image (#625)

## 📝 Description
Isolate dispatching of cloud jobs per `os_image` (on top of
`machine_type`)

## ✅ Checklist
- [x] I have tested this change
- [ ] This change requires documentation update

Author: hamir-suspect

zebra/AGENTS.md

@@ -0,0 +1,19 @@
+# Repository Guidelines
+
+## Project Structure & Module Organization
+Semaphore’s monorepo mixes Elixir services (`auth/`, `guard/`, `projecthub/`), Go tooling (`bootstrapper/`, `repohub/`), and the Phoenix/React UI in `front/`. Elixir services keep source in `lib/` and ExUnit suites in `test/`. Go utilities follow the `cmd/` entrypoint with reusable code under `pkg/`. Frontend assets, including React components and bundles, live in `front/assets/`. Shared documentation sits in `docs/` and `rfcs/`, while enterprise-specific code is isolated under `ee/`.
+
+## Build, Test, and Development Commands
+Run commands from the relevant service directory. Use `make build` to produce Docker images for CI parity. Execute `make test.ex` (or `make test.ex TEST_FILE=test/<path>.exs`) to run ExUnit suites. For Go modules, run `make test` (`go test ./...`) and add `-race` when debugging concurrency. Frontend checks run through `make test.js`. Start the Phoenix UI locally with `make dev.server`, and consult `LOCAL-DEVELOPMENT.md` for Minikube or dev-container workflows.
+
+## Coding Style & Naming Conventions
+Elixir modules are PascalCase with files in snake_case; tests end in `_test.exs`. Go packages stay lowercase, with exported identifiers PascalCase. React components use PascalCase filenames inside `front/assets/`. Before committing, run `make format.ex`, `make lint` for Go format and static checks, and `make lint.js` for ESLint/Prettier to keep the tree consistent.
+
+## Testing Guidelines
+Favor focused ExUnit `describe` blocks and use tags such as `--only integration` for longer suites. Keep Go tests table-driven in `_test.go` files. Frontend changes require Jest coverage via `make test.js`. Add regression tests alongside fixes and align Phoenix endpoint updates with matching ExUnit cases.
+
+## Commit & Pull Request Guidelines
+Commits follow Conventional Commits, e.g., `feat(auth): add token audit trail`. PRs should explain rationale, outline risk, and link issues when available. Include screenshots or logs for UI or automation changes, and confirm formatters, linters, and relevant `make test*` targets have passed before requesting review.
+
+## Security & Configuration Tips
+Run `make check.ex.deps`, `make check.go.deps`, and `make check.docker` regularly to catch dependency and image issues. Keep secrets in local `.env` files referenced by docker-compose, never checked into the repository, and export `DOCKER_BUILDKIT=1` for reproducible Docker builds. For deeper architectural context, review `DOCUMENTATION.md`.

zebra/DOCUMENTATION.md

@@ -0,0 +1,30 @@
+# Architecture Handbook
+
+## Overview
+Zebra is an Elixir OTP application responsible for job lifecycle orchestration inside Semaphore. It exposes several gRPC services on port `50051`, coordinates queueing and dispatching through background workers, and delegates persistence to a legacy Postgres schema via Ecto.
+
+## Code Layout
+- `lib/zebra/application.ex` boots the supervision tree, wiring caches, the feature flag provider, and optional gRPC servers (public/internal job & task APIs, health checks). Services are toggled with `START_*` env vars.
+- `lib/zebra/apis/` implements gRPC endpoints. Concrete service modules live under `public_job_api/`, `internal_job_api/`, and `internal_task_api/`, with shared helpers in `utils.ex`.
+- `lib/protos/` houses protobuf contracts. Regenerate stubs with `mix grpc.gen` if schemas change upstream.
+- `lib/zebra/models/` defines Ecto schemas and state helpers (`Job`, `Task`, `Project`, etc.). Database access goes through `Zebra.LegacyRepo`, whose migrations live in `priv/legacy_repo/migrations/`.
+- `lib/zebra/workers/` contains GenServer pipelines (`Dispatcher`, `Scheduler`, `TaskFinisher`, callback handlers) driven by `Workers.active/1`. Workers rely on `DbWorker` to poll the legacy database and on Watchman metrics for observability.
+- Support modules include `cache.ex` (Cachex caches), `monitor.ex` (Watchman instrumentation), `machines/` (brownout scheduling), and `feature_hub_provider.ex`.
+
+## Data & External Dependencies
+Endpoints for upstream services (artifacthub, cachehub, projecthub, RBAC, etc.) are configured in `config/*.exs` and default to `localhost:50051` for tests. Workers interact with AMQP (`tackle`), GRPC agents (`HostedAgent`, `SelfHostedAgent`), and feature flags (`FeatureProvider`). Secrets stay out of the repo—load them via `.env` for docker-compose.
+
+## Running & Development
+- `mix deps.get` to install dependencies; `make build` creates Docker images matching CI.
+- `make dev.server` starts the Phoenix UI shelling into this service; alternatively `iex -S mix phx.server` if you need an interactive node.
+- Feature flags require `FEATURE_YAML_PATH` or remote provider credentials before boot.
+
+## Testing & Quality
+- `make test.ex` wraps `mix test`, seeding/migrating the legacy repo automatically. Use `TEST_FILE=test/zebra/...` to scope cases.
+- Workers and gRPC services have fake counterparts under `test/support/fake_servers/` to keep suites hermetic. Extend these fakes when adding new integrations.
+- Linting/formatting: `make format.ex` (mix format + Credo), `make lint` for Go stubs (if any), and `make lint.js` for the Phoenix assets.
+
+## Operations & Troubleshooting
+- Set `WATCHMAN_HOST` to push metrics; logs are structured via Logger with Sentry backend (`Sentry.LoggerBackend`).
+- Scheduler cadence and worker batch sizes live in `config/*.exs` (`Zebra.Workers.Scheduler`, `Dispatcher` timeouts). Tune via env overrides before scaling changes.
+- For local brownout testing, edit `lib/zebra/machines/brownout_schedule.ex` and restart the node; the scheduler pulls updates on boot.

zebra/lib/zebra/workers/db_worker.ex

@@ -9,6 +9,7 @@ defmodule Zebra.Workers.DbWorker do
     :state_field,
     :state_value,
     :machine_type_field,
+    :machine_os_image_field,
     :machine_type_environment,
     :metric_name,
     :order_by,
@@ -45,16 +46,16 @@ defmodule Zebra.Workers.DbWorker do
     Watchman.benchmark("#{worker.metric_name}.tick.duration", fn ->
       if isolate_machine_types do
         query_machine_types(worker)
-        |> Enum.each(fn machine_type -> tick_(worker, machine_type) end)
+        |> Enum.each(fn machine_type_tuple -> tick_(worker, machine_type_tuple) end)
       else
         tick_(worker)
       end
     end)
   end
 
-  def tick_(worker, machine_type \\ nil) do
-    rows = query_jobs(worker, machine_type)
-    submit_batch_size(worker.metric_name, length(rows), machine_type)
+  def tick_(worker, machine_type_tuple \\ nil) do
+    rows = query_jobs(worker, machine_type_tuple)
+    submit_batch_size(worker.metric_name, length(rows), machine_type_tuple)
 
     parallelism = worker.parallelism || 10
 
@@ -65,8 +66,8 @@ defmodule Zebra.Workers.DbWorker do
 
   defp submit_batch_size(name, v, nil), do: Watchman.submit("#{name}.batch_size", v)
 
-  defp submit_batch_size(name, v, machine_type),
-    do: Watchman.submit({"#{name}.batch_size", [machine_type]}, v)
+  defp submit_batch_size(name, v, {machine_type, machine_os_image}),
+    do: Watchman.submit({"#{name}.batch_size", ["#{machine_type}-#{machine_os_image}"]}, v)
 
   def process(worker, id) do
     Watchman.benchmark("#{worker.metric_name}.process.duration", fn ->
@@ -98,14 +99,15 @@ defmodule Zebra.Workers.DbWorker do
 
   defp query_machine_types(worker) do
     machine_type_environment = worker.machine_type_environment || :all
+    machine_os_image_field = worker.machine_os_image_field
 
     cond do
       machine_type_environment == :all ->
         Repo.all(
           from(r in worker.schema,
             where: field(r, ^worker.state_field) == ^worker.state_value,
-            distinct: r.machine_type,
-            select: r.machine_type
+            distinct: [field(r, ^worker.machine_type_field), field(r, ^machine_os_image_field)],
+            select: {field(r, ^worker.machine_type_field), field(r, ^machine_os_image_field)}
           )
         )
 
@@ -114,8 +116,8 @@ defmodule Zebra.Workers.DbWorker do
           from(r in worker.schema,
             where: field(r, ^worker.state_field) == ^worker.state_value,
             where: like(field(r, ^worker.machine_type_field), @self_hosted_prefix),
-            distinct: r.machine_type,
-            select: r.machine_type
+            distinct: [field(r, ^worker.machine_type_field), field(r, ^machine_os_image_field)],
+            select: {field(r, ^worker.machine_type_field), field(r, ^machine_os_image_field)}
           )
         )
 
@@ -124,8 +126,8 @@ defmodule Zebra.Workers.DbWorker do
           from(r in worker.schema,
             where: field(r, ^worker.state_field) == ^worker.state_value,
             where: not like(field(r, ^worker.machine_type_field), @self_hosted_prefix),
-            distinct: r.machine_type,
-            select: r.machine_type
+            distinct: [field(r, ^worker.machine_type_field), field(r, ^machine_os_image_field)],
+            select: {field(r, ^worker.machine_type_field), field(r, ^machine_os_image_field)}
           )
         )
     end
@@ -175,19 +177,35 @@ defmodule Zebra.Workers.DbWorker do
     end
   end
 
-  defp query_jobs(worker, machine_type) do
+  defp query_jobs(worker, {machine_type, machine_os_image}) do
     order_by = worker.order_by || :id
     order_dir = worker.order_direction || :asc
     records_per_tick = worker.records_per_tick || 100
+    machine_os_image_field = worker.machine_os_image_field
 
-    Repo.all(
+    base_query =
       from(r in worker.schema,
         where: field(r, ^worker.state_field) == ^worker.state_value,
-        where: field(r, ^worker.machine_type_field) == ^machine_type,
+        where: field(r, ^worker.machine_type_field) == ^machine_type
+      )
+
+    filtered_query =
+      maybe_filter_machine_os_image(base_query, machine_os_image_field, machine_os_image)
+
+    Repo.all(
+      from(r in filtered_query,
         order_by: [{^order_dir, ^order_by}],
         select: r.id,
         limit: ^records_per_tick
       )
     )
   end
+
+  defp maybe_filter_machine_os_image(query, field, nil) do
+    from(r in query, where: is_nil(field(r, ^field)))
+  end
+
+  defp maybe_filter_machine_os_image(query, field, value) do
+    from(r in query, where: field(r, ^field) == ^value)
+  end
 end

zebra/lib/zebra/workers/dispatcher.ex

@@ -14,6 +14,7 @@ defmodule Zebra.Workers.Dispatcher do
       state_field: :aasm_state,
       state_value: Zebra.Models.Job.state_scheduled(),
       machine_type_field: :machine_type,
+      machine_os_image_field: :machine_os_image,
       machine_type_environment: machine_type_environment(),
       order_by: :scheduled_at,
       order_direction: :asc,

zebra/test/zebra/workers/dispatcher_test.exs

@@ -231,6 +231,136 @@ defmodule Zebra.Workers.DispatcherTest do
       # a NOT_FOUND response from chmura, and stopped trying to occupy agents.
       assert Zebra.Workers.DispatcherTest.Counter.value() == 10
     end
+
+    test "isolates dispatching by both machine_type and os_image" do
+      System.put_env("DISPATCH_SELF_HOSTED_ONLY", "false")
+      System.put_env("DISPATCH_CLOUD_ONLY", "false")
+
+      # we need to have at least 20 for each os_image to ensure that we
+      # don't stop batching when we receive a NOT_FOUND response from chmura
+      # and stop trying to occupy agents.
+
+      ubuntu2404_jobs =
+        Enum.map(1..20, fn _ ->
+          {:ok, job} =
+            Support.Factories.Job.create(:scheduled, %{
+              machine_type: "e1-standard-2",
+              machine_os_image: "ubuntu2404"
+            })
+
+          job
+        end)
+
+      # Create jobs with same machine_type but different os_images
+      ubuntu1804_jobs =
+        Enum.map(1..20, fn _ ->
+          {:ok, job} =
+            Support.Factories.Job.create(:scheduled, %{
+              machine_type: "e1-standard-2",
+              machine_os_image: "ubuntu1804"
+            })
+
+          job
+        end)
+
+      ubuntu2004_jobs =
+        Enum.map(1..20, fn _ ->
+          {:ok, job} =
+            Support.Factories.Job.create(:scheduled, %{
+              machine_type: "e1-standard-2",
+              machine_os_image: "ubuntu2004"
+            })
+
+          job
+        end)
+
+      # Track which os_images were requested
+      agent_requests = Agent.start_link(fn -> [] end)
+
+      GrpcMock.stub(Support.FakeServers.ChmuraApi, :occupy_agent, fn req, _ ->
+        Agent.update(elem(agent_requests, 1), fn list ->
+          [req.machine.os_image | list]
+        end)
+
+        if req.machine.os_image == "ubuntu2404" do
+          raise GRPC.RPCError, status: GRPC.Status.not_found(), message: "No suitable agent found"
+        else
+          %InternalApi.Chmura.OccupyAgentResponse{
+            agent: %InternalApi.Chmura.Agent{
+              id: Ecto.UUID.generate(),
+              ip_address: "1.2.3.4",
+              ssh_port: 80,
+              ctrl_port: 80,
+              auth_token: "asdas"
+            }
+          }
+        end
+      end)
+
+      with_stubbed_http_calls(fn ->
+        Worker.init() |> Zebra.Workers.DbWorker.tick()
+      end)
+
+      # ubuntu1804 and ubuntu2004 jobs should be started
+      (ubuntu1804_jobs ++ ubuntu2004_jobs)
+      |> Enum.each(fn job ->
+        job = Job.reload(job)
+        assert Job.started?(job) == true
+      end)
+
+      # ubuntu2404 jobs should remain scheduled (no agents available)
+      ubuntu2404_jobs
+      |> Enum.each(fn job ->
+        job = Job.reload(job)
+        assert Job.scheduled?(job) == true
+      end)
+
+      # Verify that requests were made with the correct os_images
+      requested_os_images = Agent.get(elem(agent_requests, 1), & &1)
+      assert length(requested_os_images) == 50
+      assert Enum.count(requested_os_images, &(&1 == "ubuntu1804")) == 20
+      assert Enum.count(requested_os_images, &(&1 == "ubuntu2004")) == 20
+      # only one batch requested
+      assert Enum.count(requested_os_images, &(&1 == "ubuntu2404")) == 10
+    end
+
+    test "dispatches self-hosted jobs when os_image is blank or nil" do
+      System.put_env("DISPATCH_SELF_HOSTED_ONLY", "false")
+      System.put_env("DISPATCH_CLOUD_ONLY", "false")
+
+      {:ok, blank_image_job} =
+        Support.Factories.Job.create(:scheduled, %{
+          machine_type: "s1-local-testing",
+          machine_os_image: ""
+        })
+
+      {:ok, nil_image_job} =
+        Support.Factories.Job.create(:scheduled, %{
+          machine_type: "s1-local-testing",
+          machine_os_image: nil
+        })
+
+      response = %InternalApi.SelfHosted.OccupyAgentResponse{
+        agent_id: @agent_id,
+        agent_name: "self-hosted-agent"
+      }
+
+      GrpcMock.stub(Support.FakeServers.SelfHosted, :occupy_agent, fn _, _ -> response end)
+
+      with_stubbed_http_calls(fn ->
+        Worker.init() |> Zebra.Workers.DbWorker.tick()
+      end)
+
+      blank_image_job = Job.reload(blank_image_job)
+      nil_image_job = Job.reload(nil_image_job)
+
+      assert Job.started?(blank_image_job)
+      assert Job.started?(nil_image_job)
+      assert blank_image_job.agent_id == @agent_id
+      assert nil_image_job.agent_id == @agent_id
+      assert blank_image_job.machine_os_image == ""
+      assert nil_image_job.machine_os_image in [nil, ""]
+    end
   end
 
   describe ".process" do