fix: send the first shard assignments right after the data server handshake

[merlimat]

Motivation

On a fresh cluster, and for any data server that has not bound the coordinator instance id yet, data servers get their first shard assignments 5-15s late. Until then, clients connected to them keep retrying with Failed receiving shard assignments, retrying later ... error="oxia: server not initialized yet".

The data server rejects every coordinator RPC except Handshake and the health checks with ErrNotInitialized until the handshake binds it (NewGrpcInsIDVerifyInterceptors). The data server controller starts the assignments dispatcher together with the health watch and ping goroutines that perform the handshake, so the dispatcher usually opens PushShardAssignments first. The stream is rejected: with gRPC, the rejection reaches the drain goroutine from #1216, and the dispatch loop sees the stream end as context.Canceled. The dispatcher then waits for dispatchAssignmentsBackoff before retrying: 10s initial, ±50% jitter.

This is the assignments half of the startup handshake race. #1196 proposed gating on Running || Draining, but it returned an error from the retry function, so it still waited for the backoff before the first retry.

Changes

In oxiad/coordinator/runtime/controller/dataserver/dataserver_controller.go:

• openAssignmentsStream opens the stream only once the data server is Running or Draining. It waits on a statusChanged channel instead of polling, so the first assignments go out right after the handshake and nothing is sent before it.
• advanceStatusEpochLocked replaces the three statusEpoch++ sites and closes and replaces statusChanged, so every status transition wakes the waiters.
• Failures on a Running data server still go through the regular 10s backoff.

Behavior change: a data server that failed its health checks is NotRunning too. The dispatcher now waits for it to come back and be handshaken again before reopening the stream, instead of retrying it on every backoff round. An unreachable node could not take the stream anyway. After a long outage, the next push now follows the recovery right away, instead of waiting for a backoff that has grown to its 60s cap (30-90s with jitter). A Draining node is never re-handshaken, so it is not gated.

Test support: the mock RPC provider counts PushShardAssignments attempts (PushShardAssignmentsCount).

In tests/control: TestControlRequestFeatureEnabled and TestOrderedWrites read the shard leader from the metadata right after creating the client. That only worked because client creation blocked on the late first assignments. They now wait for the shard's steady state first.

This relies on #1325. The first snapshot the reconciler publishes has no leaders yet, because the first elections are still running, and main only hid that behind the delay: NewShardManager blocks until the first snapshot arrives. Without #1325, a client created right after startup would get that snapshot and fail its first requests to the empty leader with a non-retryable Unavailable. Before #1325, TestOIDCWithPerIssuerConfig, TestOIDCWithStaticKeyFile and TestControlRequestRecordChecksum failed every run with this change. On top of #1325 they pass, and much faster than before. Go clients released before #1325 can still fail requests sent in the sub-second window before the first election on a fresh cluster. Main already exposes empty leaders the same way for namespaces created at runtime (per the code, not tested).

Testing

• New TestDataServerController_SendsAssignmentsRightAfterHandshake: the mock rejects the stream with ErrNotInitialized while the handshake is held in flight, with a 1-minute retry backoff. It fails on main (no assignments within 10s, because the retry is 30-90s away). With the fix it passes in ~0.2s, with exactly one stream opened, after the handshake. 50× repeat of the four dispatcher tests with -race: 200/200 pass.
• TestCoordinator_ShrinkCluster with -race: 12.6-22.9s on main (5 runs), 0.33-0.56s with the fix (10 runs). Client server not initialized yet retries in the logs: 54 on main, 0 with the fix.
• Rebased on fix: retry transport Unavailable errors on shard requests in the Go client #1325: go test -race ./oxiad/coordinator/... and every tests/ package pass.
• On top of fix: retry transport Unavailable errors on shard requests in the Go client #1325, with -race, 3 runs each: both OIDC tests pass in 0.5-0.7s (6-14s before), and the three tests/control tests pass in 0.5-5.4s (10-20s before). The test commit also passes on the old main, 5/5 per test.
• golangci-lint v2.13.2: 0 issues on the oxiad and tests modules.

[Copilot]

Copilot review overview

🔵 Needs a closer look

Recovery after a stream ends while the server is not running can still be delayed by the backoff interval.

Review effort: Lite
Findings: None

What changed in this PR

This PR prevents delayed initial shard assignments by waiting for data-server handshake readiness.

Changes:

• Gates assignment streams on server readiness.
• Adds status-transition signaling and regression coverage.
• Tracks stream attempts in mocks.
• Updates control tests to wait for shard steady state.

File	Summary
tests/​control/​ordered_writes_test.go	Waits for shard steady state.
tests/​control/​control_request_test.go	Waits for shard steady state.
oxiad/​coordinator/​runtime/​controller/​mockutils/​mock.go	Tracks assignment stream attempts.
oxiad/​coordinator/​runtime/​controller/​dataserver/​dataserver_controller.go	Gates streams on data-server readiness.
oxiad/​coordinator/​runtime/​controller/​dataserver/​dataserver_controller_test.go	Tests immediate post-handshake assignment delivery.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.