refactor(walreceiver): eliminate task_mgr usage

[problame]

We want to move the code base away from task_mgr.

This PR refactors the walreceiver code such that it doesn't use task_mgr anymore.

Background

As a reminder, there are three tasks in a Timeline that's ingesting WAL.
WalReceiverManager, WalReceiverConnectionHandler, and WalReceiverConnectionPoller.
See the documentation in task_mgr.rs for how they interact.

Before this PR, cancellation was requested through task_mgr::shutdown_token() and TaskHandle::shutdown.

Wait-for-task-finish was implemented using a mixture of task_mgr::shutdown_tasks and TaskHandle::shutdown.

This drawing might help:

Changes

For cancellation, the entire WalReceiver task tree now has a child_token() of Timeline::cancel. The TaskHandle no longer is a cancellation root.
This means that Timeline::cancel.cancel() is propagated.

For wait-for-task-finish, all three tasks in the task tree hold the Timeline::gate open until they exit.

The downside of using the Timeline::gate is that we can no longer wait for just the walreceiver to shut down, which is particularly relevant for Timeline::flush_and_shutdown.
Effectively, it means that we might ingest more WAL while the freeze_and_flush() call is ongoing.

Also, drive-by-fix the assertiosn around task kinds in wait_lsn. The check for WalReceiverConnectionHandler was ineffective because that never was a task_mgr task, but a TaskHandle task. Refine the assertion to check whether we would wait, and only fail in that case.

Alternatives

I contemplated (ab-)using the Gate by having a separate Gate for struct WalReceiver.
All the child tasks would use that gate instead of Timeline::gate.
And struct WalReceiver itself would hold an Option<GateGuard> of the Timeline::gate.
Then we could have a WalReceiver::stop function that closes the WalReceiver's gate, then drops the WalReceiver::Option<GateGuard>.

However, such design would mean sharing the WalReceiver's Gate in an Arc, which seems awkward.
A proper abstraction would be to make gates hierarchical, analogous to CancellationToken.

In the end, @jcsp and I talked it over and we determined that it's not worth the effort at this time.

Refs

part of #7062

[github-actions]

2730 tests run: 2592 passed, 0 failed, 138 skipped (full report)

---

Code coverage* (full report)

• functions: 28.2% (6311 of 22369 functions)
• lines: 47.0% (44295 of 94335 lines)

* collected from Rust tests only

---

_{The comment gets automatically updated with the latest test results
c0d5cea at 2024-03-28T12:09:13.984Z :recycle:}

[problame]

What I read between the lines of your review here is that you would like to preserve the "wait-for-task-finish" semantics that we had before this PR, correct?

If so, I think keeping track of the spawned tasks using JoinHandle / JoinSet is a nice explicit way to do it (somewhere in the history of this PR, I had it half-done that way).

The less explicit alternative of achieving the same thing is to not keep track of the spawned tasks explicitly, but use the approach pointed out in # Alternatives in the PR description.

But, as also pointed out in the PR description:

In the end, @jcsp and I talked it over and we determined that it's not worth the effort at this time.

Elaborating more on that: early walreceiver shutdown is the only use case where we can't just Timeline::cancel.cancel() + Timeline::gate.close().await.
And, early walreceiver shutdown is only relevant for the ShutdownMode::FreezeAndFlush use case. In the ::Hard mode, even before this PR, we'd cancel + gate.close() immediately anyway.

So, we concluded that we shouldn't spend the complexity budget on walreceiver if possible.

EDIT: I realized that ShutdownMode::* is a construct that is added in the PR that builds on top of this one: refactor(Timeline::shutdown): rely more on Timeline::cancel; use it from deletion code path

[problame]

One sensible argument in favor of continuing to have early walreceiver shutdown with proper "wait-for-task-finish" semantics is that, if the walreceiver tasks continue to write data past the remote_client.shutdown() call inside the try_freeze_and_flush=true branch (link), then the flush loop will encounter errors from the remote_client when it tries to queue new layers for upload, but we don't have a bound / back-pressure on the number of frozen layers, so, there's OOM risk.

However, that's a big if and frankly, it doesn't matter practically, because right after remote_client.shutdown().await, we proceed with Timeline::cancel and teardown of the remaining tasks + Timeline::gate.close().await.

[koivunej]

We had a call on this. I agree that waiting for walreceiver after cancelation is not required; I cannot see how this could fail or what would the worst-case outcome be, probably nothing. I would still have preferred to keep the waiting as more of a refactor and as a design that would be future-proof.

There are many ways of achieving a scalable design for starting and stopping tasks without task_mgr. My preferred would have been with joinsets, cancellationtokens, and plain async fn without all of the unnecessary structs.