Lab 4.3 — Build It: WavingVertexManager

Lab type: Build It — a VertexManagerPlugin with a full JUnit + Mockito test suite Estimated time: 120–150 min Maven module: book/projects/level-4-waving-manager Key class: org.apache.tez.learning.l4.WavingVertexManager

Background

In Lab 4.2 you read the VertexManagerPlugin contract and watched ShuffleVertexManager schedule reduce tasks in a slow-start ramp. This lab has you build a plugin of your own — a WavingVertexManager that schedules a vertex's tasks in fixed-size waves:

• Wave 0: tasks 0 .. waveSize-1
• Wave 1: tasks waveSize .. 2*waveSize-1
• Wave N: launched only after every task in wave N-1 has succeeded

The wave size is read from UserPayload as "waveSize=N"; the default is WavingVertexManager.DEFAULT_WAVE_SIZE = 2. This is a small but complete plugin — the same architectural shape as ImmediateStartVertexManager (schedule tasks via the context) plus the completion-gated advancement pattern that every real slow-start manager uses. The companion project book/projects/level-4-waving-manager already contains the finished code and a 11-test suite; your job is to understand it deeply enough to have written it, then extend it.

Why waves? When downstream tasks each open a scarce resource — a database connection, a rate- limited external API — launching all of them at once overwhelms it. Waves cap concurrency at waveSize while still letting Tez manage retries and locality.

Deep-dive companions: vertex-lifecycle.md, event-routing.md, and the Level 4 overview.

Why This Lab Matters for Contributors

Every non-trivial Tez integration — Hive, Pig, Cascading — ships a custom VertexManagerPlugin. Writing one, and unit-testing it against a mocked VertexManagerPluginContext, is the single most transferable skill in Level 4: it is exactly how the Tez project tests ShuffleVertexManager and ImmediateStartVertexManager, and exactly the diff shape a scheduling-behavior PR takes. Building WavingVertexManager from scratch teaches you the constructor/initialize split, the completion bookkeeping (BitSets), the no-double-schedule guard, and the Mockito ArgumentCaptor pattern for asserting which tasks were scheduled — all of which you will use when you fix a real one in Lab 4.4.

Prerequisites

• Completed Lab 4.2; you know scheduleTasks(List<ScheduleTaskRequest>) is the current API and how initialize / onVertexStarted fire.
• The level-4-waving-manager module builds. Its pom.xml inherits tez-api, tez-common, tez-runtime-library, JUnit 4.13.2, and Mockito from the learning parent pom.
• These files open side by side:
  • book/projects/level-4-waving-manager/src/main/java/org/apache/tez/learning/l4/WavingVertexManager.java
  • book/projects/level-4-waving-manager/src/test/java/org/apache/tez/learning/l4/TestWavingVertexManager.java
  • tez-dag/.../impl/ImmediateStartVertexManager.java (the closest real analog)

Step 1 — Understand the contract, then compile the project

Re-read the three Tez sources so the project code has context:

find ~/src/oss-repos/tez -name "VertexManagerPlugin.java" | grep -v target
find ~/src/oss-repos/tez -name "VertexManagerPluginContext.java" | grep -v target
find ~/src/oss-repos/tez -name "ImmediateStartVertexManager.java" | grep -v target

Then build and run the companion project's tests:

cd book/projects
mvn -pl level-4-waving-manager test

Expected (11 tests, all green):

Tests run: 11, Failures: 0, Errors: 0, Skipped: 0

Answer before reading further:

1. List the lifecycle methods a VertexManagerPlugin must implement. Which are abstract (onVertexManagerEventReceived, onRootVertexInitialized) and which have a default body?
2. When does the AM call initialize()? Can you call scheduleTasks(...) from inside it? (You can, but the project deliberately does not — see Step 3.)
3. What does getContext().scheduleTasks(List<ScheduleTaskRequest>) do to the execution engine? (It enqueues those task indices for the scheduler; unscheduled tasks simply never run.)

Step 2 — Read WavingVertexManager field by field

Open the source. The class extends VertexManagerPlugin and keeps its scheduling state in two BitSets and a cursor (quoted from module level-4-waving-manager, class WavingVertexManager):

public static final String PAYLOAD_KEY  = "waveSize";
public static final int    DEFAULT_WAVE_SIZE = 2;

private VertexManagerPluginContext context;
private int waveSize;
private int totalTasks;

/** Tracks which task indices have already been scheduled (not scheduled twice). */
private BitSet scheduled;
/** Tracks which task indices in the current wave have finished. */
private BitSet waveFinished;
/** Index of the first task in the wave that has not yet been scheduled. */
private int nextTaskToSchedule;

Note the constructor — it just stores the context, doing no real work, exactly as the contract requires:

public WavingVertexManager(VertexManagerPluginContext context) {
    super(context);
    this.context = context;
}

Step 3 — Trace the four methods that do the work

initialize()

It parses the payload and sizes the BitSets (quoted):

totalTasks      = context.getVertexNumTasks(context.getVertexName());
scheduled       = new BitSet(totalTasks);
waveFinished    = new BitSet(totalTasks);
nextTaskToSchedule = 0;

The payload parsing reads UserPayload as UTF-8, splits on ;, and for the waveSize token parses a positive int — falling back to DEFAULT_WAVE_SIZE on absence, a non-integer, or a non-positive value. (This is exactly what testDefaultWaveSizeWhenNoPayload, testPayloadParsed, testNegativePayloadFallsBackToDefault, and testGarbagePayloadFallsBackToDefault assert.)

onVertexStarted(Map<String, List<Integer>> completions)

The whole body is one line:

@Override
public void onVertexStarted(Map<String, List<Integer>> completions) {
    scheduleNextWave();
}

Why is scheduleNextWave() here and not in initialize()? Because onVertexStarted is the AM's signal that inputs are ready — scheduling from initialize() would launch tasks before the vertex is truly startable. (This project overrides the deprecated Map<String,List<Integer>> overload of onVertexStarted; a new plugin would prefer the List<TaskAttemptIdentifier> overload from Lab 4.2. Both are safe — the base class adapts one to the other — but the modern overload gives you full attempt identifiers instead of bare task indices. Migrating the project to the modern overload, and adjusting the tests that pass Collections.emptyMap(), is a worthwhile refactor and a good first "improve an internal API usage" exercise.)

onTaskAttemptCompleted(int taskIndex, boolean successful)

This is the wave gate:

public void onTaskAttemptCompleted(int taskIndex, boolean successful) {
    if (!successful) {
        return; // let the AM handle retries
    }
    waveFinished.set(taskIndex);
    checkAndScheduleNextWave();
}

Failed attempts are silently ignored so the AM's retry policy can re-run them; a failed task must not count toward wave completion. (This is what testFailedAttemptDoesNotAdvanceWave pins.)

scheduleNextWave() and checkAndScheduleNextWave()

private void scheduleNextWave() {
    List<ScheduleTaskRequest> toSchedule = new ArrayList<>();
    int count = 0;
    while (nextTaskToSchedule < totalTasks && count < waveSize) {
        int idx = nextTaskToSchedule++;
        if (!scheduled.get(idx)) {
            scheduled.set(idx);
            toSchedule.add(ScheduleTaskRequest.create(idx, null));
            count++;
        }
    }
    if (!toSchedule.isEmpty()) {
        context.scheduleTasks(toSchedule);
    }
}

private void checkAndScheduleNextWave() {
    BitSet scheduledCopy = (BitSet) scheduled.clone();
    scheduledCopy.andNot(waveFinished);
    if (scheduledCopy.isEmpty() && nextTaskToSchedule < totalTasks) {
        scheduleNextWave();
    }
}

Study three things:

1. The loop terminates on count < waveSize for full waves and on nextTaskToSchedule < totalTasks for the last (short) wave — which is why waveSize=1000, totalTasks=5 is safe.
2. ScheduleTaskRequest.create(idx, null) passes null for the location hint — the plugin expresses no locality preference. Note it uses the current scheduleTasks API, matching Lab 4.2.
3. checkAndScheduleNextWave clones scheduled before andNot(waveFinished). Without the clone, andNot would mutate scheduled, corrupting the double-schedule guard. (Step 5 makes you break this deliberately.)

Trace it by hand. For totalTasks=4, waveSize=2, tabulate scheduled, waveFinished, and nextTaskToSchedule after each callback. Check your work against this table:

The subtle row is the third: when task 1 finishes, checkAndScheduleNextWave computes scheduled.clone().andNot(waveFinished) = {0,1}.andNot({0,1}) = {} (empty), and nextTaskToSchedule (2) < totalTasks (4), so wave 1 ([2,3]) is scheduled. Note that task 0 was recorded in waveFinished on its own callback, so by the time task 1 completes, both wave-0 tasks are marked finished and the copy comes out empty. Confirm scheduleTasks fires exactly twice.

Step 4 — Read the test suite and its Mockito patterns

Open TestWavingVertexManager.java. Every test builds the manager against a mocked context:

mockContext = mock(VertexManagerPluginContext.class);
when(mockContext.getVertexName()).thenReturn("TestVertex");

The Mockito vocabulary you must recognize:

The two most instructive tests:

• testThreeWavesForSixTasks — the full lifecycle: 6 tasks, wave size 2, drives all six completions and asserts verify(mockContext, times(3)).scheduleTasks(anyList()) plus getScheduled().cardinality() == 6. This is the integration test.
• testWave0ContainsFirstTwoTasks — uses an ArgumentCaptor<List<ScheduleTaskRequest>> to capture wave 0 and asserts wave0.get(0).getTaskIndex() == 0 and .get(1).getTaskIndex() == 1. This proves the identity of scheduled tasks, not just the count.

Run one test in isolation to see the pattern live:

mvn -pl level-4-waving-manager test -Dtest=TestWavingVertexManager#testThreeWavesForSixTasks

Step 5 — Break it: three experiments

Make each change, run the suite, observe the failure, then revert. (Work on a scratch copy — the committed project must stay green; do not leave these edits in place.)

Experiment A — remove the failure guard

Delete if (!successful) return; in onTaskAttemptCompleted, then run:

mvn -pl level-4-waving-manager test -Dtest=TestWavingVertexManager#testFailedAttemptDoesNotAdvanceWave

Which test fails, and what is the actual vs expected scheduleTasks count? Why does treating a failure as a success advance the wave prematurely?

Experiment B — remove the BitSet.clone()

Change checkAndScheduleNextWave to scheduled.andNot(waveFinished) (mutating scheduled directly). Run the full suite. Which tests fail? Trace testThreeWavesForSixTasks by hand: after wave 0, scheduled is corrupted to empty, so the double-schedule guard no longer protects wave 1.

Experiment C — off-by-one in the wave loop

Change count < waveSize to count <= waveSize. How many tasks does wave 0 now schedule, and which test catches it (testWave0ScheduledOnVertexStarted)?

Step 5b — Observe wave scheduling (add logging)

The committed plugin is quiet. To watch the waves, add a temporary log line at the top of scheduleNextWave (on your scratch copy):

LOG.info("WavingVertexManager: scheduling wave starting at {} (count up to {})",
    nextTaskToSchedule, waveSize);

You will need a logger field (private static final Logger LOG = LoggerFactory.getLogger(WavingVertexManager.class);). Re-run testThreeWavesForSixTasks with -Dtest=... and read the surefire output under level-4-waving-manager/target/surefire-reports/; you should see three "scheduling wave" lines at offsets 0, 2, 4. In a real DAG the same lines appear in the AM (syslog_dag_*) container log, interleaved with the VertexImpl "Task Completion:" lines from checkTasksForCompletion you read in Lab 4.1 — that interleaving is how you confirm, on a live run, that wave N+1 starts only after wave N's completions land. Revert the log line before committing.

Step 6 — Extend it (choose at least one)

Keep the committed project green; add these as new methods/tests on your own branch.

1. Consume producer statistics. Give the plugin a real onVertexManagerEventReceived(VertexManagerEvent) body: if a producer's payload says a task's output is empty, mark it done so it does not block wave advancement. Model the decode on ShuffleVertexManagerBase.handleVertexManagerEvent from Lab 4.2. Add a test testSkipEventReleasesWave that sends such an event and verifies the wave advances.
2. Percentage-based waves. Add a "wavePercent=P" payload key so wave size is ceil(P% * totalTasks). Add tests for P=50 on 6 tasks (waves of 3) and P=100 (one wave = immediate start, equivalent to ImmediateStartVertexManager).
3. Handle upstream failures. Override onSourceTaskCompleted(TaskAttemptIdentifier) (the modern overload) so the plugin only starts wave 0 after all source tasks have completed, making the manager safe on a vertex with real input edges. Test with a stubbed source vertex.

Step 7 — Map the project onto the Tez source

Fill this table from your reading (paths under your Tez checkout):

Step 8 — Wire it onto a real vertex (reading exercise)

You will not run a full DAG here, but you must be able to read the wiring. On the Vertex API:

grep -n "setVertexManagerPlugin" ~/src/oss-repos/tez/tez-api/src/main/java/org/apache/tez/dag/api/Vertex.java
grep -n "public static VertexManagerPluginDescriptor create\|setUserPayload" \
  ~/src/oss-repos/tez/tez-api/src/main/java/org/apache/tez/dag/api/VertexManagerPluginDescriptor.java

A DAG author attaches the plugin like this (verify each call exists in your source):

VertexManagerPluginDescriptor vmDesc =
    VertexManagerPluginDescriptor.create(WavingVertexManager.class.getName());
vmDesc.setUserPayload(
    UserPayload.create(ByteBuffer.wrap("waveSize=3".getBytes(StandardCharsets.UTF_8))));
reduceVertex.setVertexManagerPlugin(vmDesc);

Answer: which class carries the payload from the DAG definition to initialize()? (VertexManagerPluginDescriptor → the AM reflectively constructs the plugin and exposes the payload via getContext().getUserPayload().)

Deliverables

• mvn -pl level-4-waving-manager test green: Tests run: 11, Failures: 0.
• A by-hand trace table of scheduled / waveFinished / nextTaskToSchedule for 4 tasks, wave size 2.
• Written answers to the field/method questions in Steps 2–4.
• Experiments A, B, C run and reverted, each with the failing test named and the failure explained.
• At least one Step 6 extension implemented with a passing test.
• The Step 7 source-mapping table completed and the Step 8 wiring understood.

Troubleshooting

Stretch Goals

1. Compare to ShuffleVertexManager. Read its onVertexStarted/schedulePendingTasks. It does not schedule immediately — it waits on the slow-start fraction. List the five methods that hold its core scheduling logic (onVertexStarted, onSourceTaskCompleted, getNumOfTasksToSchedule, computeRouting, schedulePendingTasks) and contrast its double-schedule guard (pendingTasks list) with your scheduled BitSet.
2. JIRA research. Search the Tez tracker for a resolved VertexManagerPlugin scheduling bug (git log --oneline -- '*VertexManager*'). Read one commit: was it a race, a double-schedule, or a wrong wave/parallelism boundary? What did the added test mock?
3. Concurrency. onVertexManagerEventReceived may run concurrently with onVertexStarted (see the VertexManagerPlugin Javadoc). If you implemented Step 6.1, add the synchronization the real ShuffleVertexManagerBase uses (synchronized on the event handler) and justify it.

Validation / Self-check

Answer in your own words, citing the class:

1. Why does WavingVertexManager do nothing in its constructor and defer all setup to initialize()?

2. What are scheduled and waveFinished, and why must checkAndScheduleNextWave clone before andNot?

3. Why does onTaskAttemptCompleted ignore failed attempts, and what would break if it did not?

4. Which loop condition guarantees a short final wave is scheduled correctly when totalTasks is not a multiple of waveSize?

5. In the test suite, how does ArgumentCaptor let testWave0ContainsFirstTwoTasks assert which tasks were scheduled, not merely how many?

6. Which API attaches this plugin to a vertex, and how does the "waveSize=N" payload reach initialize()?

7. How does WavingVertexManager's scheduling strategy differ from ImmediateStartVertexManager and from ShuffleVertexManager?

8. Trace the 4-task, wave-size-2 example: at the moment the second wave-0 task completes, what does scheduled.clone().andNot(waveFinished) evaluate to, and why does that trigger wave 1?

Etiquette reminder for when you upstream a plugin. A VertexManagerPlugin PR to Apache Tez needs a unit test built exactly like TestWavingVertexManager (mock the context, drive the callbacks, verify the scheduling), a TEZ-xxxx JIRA reference, and a DCO Signed-off-by (git commit -s). The muscle you built here — asserting scheduling behavior against a mocked context — is the same muscle every reviewer will look for.

When your extension compiles, your test passes, and you can explain every field's role, you have completed Lab 4.3. Continue to Lab 4.4: Fix It — ShuffleVertexManager Zero-Partition Guard.