pageserver: add Gate as a partner to CancellationToken for safe shutdown of Tenant & Timeline

[jcsp]

Problem

When shutting down a Tenant, it isn't just important to cause any background tasks to stop. It's also important to wait until they have stopped before declaring shutdown complete, in cases where we may re-use the tenant's local storage for something else, such as running in secondary mode, or creating a new tenant with the same ID.

Summary of changes

A Gate class is added, inspired by seastar::gate. For types that have an important lifetime that corresponds to some physical resource, use of a Gate as well as a CancellationToken provides a robust pattern for async requests & shutdown:

• Requests must always acquire the gate as long as they are using the object
• Shutdown must set the cancellation token, and then close() the gate to wait for requests in progress before returning.

This is not for memory safety: it's for expressing the difference between "Arc exists", and "This tenant's files on disk are eligible to be read/written".

• Both Tenant and Timeline get a Gate & CancellationToken.
• The Timeline gate is held during eviction of layers, and during page_service requests.
• Existing cancellation support in page_service is refined to use the timeline-scope cancellation token instead of a process-scope cancellation token. This replaces the use of task_mgr::associate_with: tasks no longer change their tenant/timelineidentity after being spawned.

The Tenant's Gate is not yet used, but will be important for Tenant-scoped operations in secondary mode, where we must ensure that our secondary-mode downloads for a tenant are gated wrt the activity of an attached Tenant.

This is part of a broader move away from using the global-state driven task_mgr shutdown tokens:

• less global state where we rely on implicit knowledge of what task a given function is running in, and more explicit references to the cancellation token that a particular function/type will respect, making shutdown easier to reason about.
• eventually avoid the big global TASKS mutex.

[github-actions]

2358 tests run: 2243 passed, 0 failed, 115 skipped (full report)

---

Flaky tests (2)

Postgres 16

• test_crafted_wal_end[last_wal_record_crossing_segment]: release

Postgres 14

• test_crafted_wal_end[last_wal_record_xlog_switch_ends_on_page_boundary]: debug

Code coverage (full report)

• functions: 54.6% (8858 of 16232 functions)
• lines: 81.6% (50930 of 62421 lines)

---

_{The comment gets automatically updated with the latest test results
0eb1060 at 2023-11-06T12:17:44.227Z :recycle:}

[koivunej]

I dislike how this seems to duplicate what utils::completion aims to do (which tokio now has as TaskTracker), but I don't have more time to review right now.

[jcsp]

I dislike how this seems to duplicate what utils::completion aims to do (which tokio now has as TaskTracker), but I don't have more time to review right now.

I wouldn't say it duplicates it: neither one is a concurrency primitive (both are wrappers), and they have different goals:

• Completion is a "you must have at least one clone alive until you're done". As long as you want to keep a barrier blocked, you have to keep at least one Completion around.
• Gate is a "you can have zero guards, and acquire more later, but you'll be refused if the gate is closed". But you don't have to keep a guard around all the time to clone it.

One could implement Gate as a funky wrapper on Barrier, but the way I look at it they're both ergonomic wrappers around tokio concurrency types, so it makes sense to have two different types to implement two different usage styles.

[koivunej]

I wouldn't say it duplicates it: neither one is a concurrency primitive (both are wrappers), and they have different goals:

• Completion is a "you must have at least one clone alive until you're done". As long as you want to keep a barrier blocked, you have to keep at least one Completion around.
• Gate is a "you can have zero guards, and acquire more later, but you'll be refused if the gate is closed". But you don't have to keep a guard around all the time to clone it.

It introduces more state than we need: my idea was to add a completions to state (task local) from which we can start new work. If we are not in that state, we cannot obtain a guard thus we cannot start new work. If we transitioning from the state, then we must wait out all outstanding work.