gen_stage_test.exs

defmodule Lockstep.GenStageTest do
  @moduledoc """
  End-to-end demonstration: real upstream `gen_stage` source
  (cloned from `elixir-lang/gen_stage`, ~5,200 LOC of pure Elixir),
  processed through `Lockstep.MixCompiler`, driving a producer →
  consumer pipeline under Lockstep's controlled scheduling.

  GenStage is the foundation under Broadway and Flow. If GenStage
  works under Lockstep, that's a meaningful chunk of the
  data-pipeline ecosystem reachable for race testing.

  Producer / consumer modules are defined at runtime via
  `Module.create/3` after GenStage is loaded -- otherwise their
  `use GenStage` would fail at test-file compile time.

  Skipped if upstream isn't on disk at `/tmp/gen_stage_src`. Clone:

      gh repo clone elixir-lang/gen_stage /tmp/gen_stage_src
  """

  use ExUnit.Case, async: false

  @gen_stage_src "/tmp/gen_stage_src"
  @rewritten_dir "/tmp/gen_stage_lockstep"

  setup_all do
    if not File.exists?(@gen_stage_src) do
      IO.puts(:stderr, "\n[gen_stage] skipped: clone elixir-lang/gen_stage to #{@gen_stage_src}")
      {:ok, available: false}
    else
      File.rm_rf!(@rewritten_dir)

      paths = [Path.join(@gen_stage_src, "lib/**/*.ex")]
      {:ok, rewritten} = Lockstep.MixCompiler.compile(%{paths: paths, output: @rewritten_dir})

      ordered =
        Enum.sort_by(rewritten, fn p ->
          cond do
            String.ends_with?(p, "utils.ex") -> 0
            String.ends_with?(p, "buffer.ex") -> 1
            String.contains?(p, "/dispatchers/") -> 2
            String.ends_with?(p, "dispatcher.ex") -> 3
            String.ends_with?(p, "gen_stage.ex") -> 4
            true -> 5
          end
        end)

      for path <- ordered do
        try do
          Code.compile_file(path)
        rescue
          _e -> :ok
        end
      end

      if Code.ensure_loaded?(GenStage) do
        define_test_modules()
        {:ok, available: true}
      else
        {:ok, available: false}
      end
    end
  end

  defp define_test_modules do
    producer_ast =
      quote do
        use GenStage

        def start_link(initial), do: GenStage.start_link(__MODULE__, initial)

        @impl true
        def init(initial), do: {:producer, initial}

        @impl true
        def handle_demand(demand, state) when demand > 0 do
          events = Enum.to_list(state..(state + demand - 1))
          {:noreply, events, state + demand}
        end
      end

    consumer_ast =
      quote do
        use GenStage

        def start_link({producer, parent}) do
          GenStage.start_link(__MODULE__, {producer, parent})
        end

        @impl true
        def init({producer, parent}) do
          {:consumer, parent, subscribe_to: [{producer, max_demand: 4, min_demand: 0}]}
        end

        @impl true
        def handle_events(events, _from, parent) do
          Lockstep.send(parent, {:events, events})
          {:noreply, [], parent}
        end
      end

    # Producer using BroadcastDispatcher: events go to ALL consumers,
    # constrained by minimum-demand-among-all-subscribers.
    # `demand: :accumulate` -- per GenStage docs, this pauses demand
    # at startup so all consumers subscribe before the first batch
    # ships, otherwise late subscribers miss the first events
    # (documented behaviour, not a bug).
    broadcast_producer_ast =
      quote do
        use GenStage

        def start_link(initial),
          do: GenStage.start_link(__MODULE__, initial)

        @impl true
        def init(initial) do
          {:producer, initial, dispatcher: GenStage.BroadcastDispatcher, demand: :accumulate}
        end

        @impl true
        def handle_demand(demand, state) when demand > 0 do
          events = Enum.to_list(state..(state + demand - 1))
          {:noreply, events, state + demand}
        end
      end

    # Consumer that records (label, events) pairs so we can verify
    # broadcast invariant: every consumer should see the SAME events
    # in the SAME order.
    labeled_consumer_ast =
      quote do
        use GenStage

        def start_link({producer, label, parent}) do
          GenStage.start_link(__MODULE__, {producer, label, parent})
        end

        @impl true
        def init({producer, label, parent}) do
          {:consumer, {label, parent}, subscribe_to: [{producer, max_demand: 3, min_demand: 0}]}
        end

        @impl true
        def handle_events(events, _from, {label, parent} = state) do
          Lockstep.send(parent, {:bcast, label, events})
          {:noreply, [], state}
        end
      end

    Module.create(
      Lockstep.GenStageTest.CountingProducer,
      producer_ast,
      Macro.Env.location(__ENV__)
    )

    Module.create(
      Lockstep.GenStageTest.CollectingConsumer,
      consumer_ast,
      Macro.Env.location(__ENV__)
    )

    Module.create(
      Lockstep.GenStageTest.BroadcastProducer,
      broadcast_producer_ast,
      Macro.Env.location(__ENV__)
    )

    Module.create(
      Lockstep.GenStageTest.LabeledConsumer,
      labeled_consumer_ast,
      Macro.Env.location(__ENV__)
    )
  end

  test "producer-consumer pipeline drives real GenStage code under Lockstep", ctx do
    if not ctx.available do
      assert true
    else
      :ok =
        Lockstep.Runner.run(
          fn ->
            {:ok, prod} = Lockstep.GenStageTest.CountingProducer.start_link(0)
            parent = self()
            {:ok, _cons} = Lockstep.GenStageTest.CollectingConsumer.start_link({prod, parent})

            events =
              for _ <- 1..3 do
                Lockstep.recv_first(fn
                  {:events, _} -> true
                  _ -> false
                end)
              end
              |> Enum.flat_map(fn {:events, list} -> list end)

            sorted = Enum.sort(events)

            if sorted != events do
              raise "producer emitted out of order: #{inspect(events)}"
            end

            if length(events) < 3 do
              raise "expected at least 3 events, got #{inspect(events)}"
            end
          end,
          iterations: 5,
          strategy: :pct,
          max_steps: 5_000,
          seed: 17,
          iter_timeout: 30_000,
          suite: "gen_stage_basic"
        )
    end
  end

  # NOTE: Cancellation under Lockstep requires the consumer's death
  # to flow through the controller (so :DOWN is delivered via
  # Lockstep's monitor system, not the BEAM mailbox). External
  # `Process.exit(pid, :kill)` from the test body bypasses the
  # controller, which is a Lockstep modeling boundary, not a real
  # GenStage bug. So for consumer-crash testing we use a consumer
  # that crashes itself via a raised exception in its handler --
  # cleanly observable through the controller's link/monitor graph.
  test "consumer that raises in handle_events doesn't stall the producer", ctx do
    if not ctx.available do
      assert true
    else
      # Define a self-crashing consumer that raises after the first
      # batch arrives. Because it was started under the controller,
      # the death + monitor :DOWN flow through Lockstep.
      crashing_ast =
        quote do
          use GenStage

          def start_link({producer, parent}) do
            GenStage.start_link(__MODULE__, {producer, parent})
          end

          @impl true
          def init({producer, parent}) do
            {:consumer, parent, subscribe_to: [{producer, max_demand: 2, min_demand: 0}]}
          end

          @impl true
          def handle_events([_ | _], _from, parent) do
            Lockstep.send(parent, :crasher_got_batch)
            raise "crasher self-destruct"
          end
        end

      unless Code.ensure_loaded?(Lockstep.GenStageTest.CrashingConsumer) do
        Module.create(
          Lockstep.GenStageTest.CrashingConsumer,
          crashing_ast,
          Macro.Env.location(__ENV__)
        )
      end

      :ok =
        Lockstep.Runner.run(
          fn ->
            {:ok, prod} = Lockstep.GenStageTest.CountingProducer.start_link(0)
            parent = self()

            {:ok, _stable} =
              Lockstep.GenStageTest.CollectingConsumer.start_link({prod, parent})

            {:ok, _crasher} =
              Lockstep.GenStageTest.CrashingConsumer.start_link({prod, parent})

            # Wait for crasher to receive a batch, then crash.
            Lockstep.recv_first(fn
              :crasher_got_batch -> true
              _ -> false
            end)

            # Stable consumer must continue to receive new batches.
            # The producer's monitor on crasher fires :DOWN, the
            # subscription gets cancelled, and demand redistributes.
            for _ <- 1..2 do
              Lockstep.recv_first(fn
                {:events, _} -> true
                _ -> false
              end)
            end
          end,
          iterations: 3,
          strategy: :pct,
          max_steps: 10_000,
          seed: 0xCAFE_F00D,
          iter_timeout: 30_000,
          suite: "gen_stage_consumer_raises"
        )
    end
  end

  test "BroadcastDispatcher: every consumer sees every event in the same order", ctx do
    if not ctx.available do
      assert true
    else
      :ok =
        Lockstep.Runner.run(
          fn ->
            {:ok, prod} = Lockstep.GenStageTest.BroadcastProducer.start_link(0)
            parent = self()

            {:ok, _a} =
              Lockstep.GenStageTest.LabeledConsumer.start_link({prod, :a, parent})

            {:ok, _b} =
              Lockstep.GenStageTest.LabeledConsumer.start_link({prod, :b, parent})

            {:ok, _c} =
              Lockstep.GenStageTest.LabeledConsumer.start_link({prod, :c, parent})

            # Now that all consumers are subscribed, resume demand so
            # everyone gets the same first batch.
            :ok = GenStage.demand(prod, :forward)

            # Collect 6 batches total -- 2 from each consumer.
            collected =
              for _ <- 1..6 do
                Lockstep.recv_first(fn
                  {:bcast, _, _} -> true
                  _ -> false
                end)
              end

            by_label =
              collected
              |> Enum.group_by(
                fn {:bcast, label, _} -> label end,
                fn {:bcast, _, events} -> events end
              )

            # Each consumer should see at least 2 batches.
            for {label, batches} <- by_label do
              if length(batches) < 1 do
                raise "consumer #{label} got no batches"
              end
            end

            # Compute the merged event sequence per consumer.
            sequences =
              for {label, batches} <- by_label, into: %{} do
                {label, batches |> Enum.flat_map(& &1)}
              end

            # All consumers should have a common prefix of events --
            # broadcast invariant says everyone gets the same sequence.
            min_len = sequences |> Map.values() |> Enum.map(&length/1) |> Enum.min()

            if min_len < 1 do
              raise "no consumer received any events: #{inspect(sequences)}"
            end

            prefixes =
              sequences
              |> Enum.map(fn {label, evs} -> {label, Enum.take(evs, min_len)} end)
              |> Enum.into(%{})

            unique = prefixes |> Map.values() |> Enum.uniq()

            if length(unique) > 1 do
              raise "BROADCAST INVARIANT VIOLATED: consumers saw divergent prefixes:\n#{inspect(prefixes, pretty: true)}"
            end
          end,
          iterations: 20,
          strategy: :pct,
          max_steps: 20_000,
          seed: 0xB1A_5CA5,
          iter_timeout: 60_000,
          suite: "gen_stage_broadcast_dispatcher"
        )
    end
  end

  test "two consumers fan-out: each event goes to exactly one consumer (DemandDispatcher)", ctx do
    if not ctx.available do
      assert true
    else
      :ok =
        Lockstep.Runner.run(
          fn ->
            {:ok, prod} = Lockstep.GenStageTest.CountingProducer.start_link(100)
            parent = self()

            {:ok, _c1} = Lockstep.GenStageTest.CollectingConsumer.start_link({prod, parent})
            {:ok, _c2} = Lockstep.GenStageTest.CollectingConsumer.start_link({prod, parent})

            collected =
              for _ <- 1..4 do
                Lockstep.recv_first(fn
                  {:events, _} -> true
                  _ -> false
                end)
              end
              |> Enum.flat_map(fn {:events, list} -> list end)

            uniq = Enum.uniq(collected)

            if uniq != collected do
              raise "DemandDispatcher delivered duplicates: #{inspect(collected)}"
            end
          end,
          iterations: 5,
          strategy: :pct,
          max_steps: 10_000,
          seed: 31,
          iter_timeout: 60_000,
          suite: "gen_stage_fanout"
        )
    end
  end
end