libs: make backoff::retry() take a cancellation token (#5065)

## Problem

Currently, anything that uses backoff::retry will delay the join of its
task by however long its backoff sleep is, multiplied by its max
retries.

Whenever we call a function that sleeps, we should be passing in a
CancellationToken.

## Summary of changes

- Add a `Cancel` type to backoff::retry that wraps a CancellationToken
and an error `Fn` to generate an error if the cancellation token fires.
- In call sites that already run in a `task_mgr` task, use
`shutdown_token()` to provide the token. In other locations, use a dead
`CancellationToken` to satisfy the interface, and leave a TODO to fix it
up when we broaden the use of explicit cancellation tokens.

pageserver/src/tenant/delete.rs

@@ -7,6 +7,7 @@ use anyhow::Context;
 use pageserver_api::models::TenantState;
 use remote_storage::{DownloadError, GenericRemoteStorage, RemotePath};
 use tokio::sync::OwnedMutexGuard;
+use tokio_util::sync::CancellationToken;
 use tracing::{error, info, instrument, warn, Instrument, Span};
 
 use utils::{
@@ -82,6 +83,8 @@ async fn create_remote_delete_mark(
         FAILED_UPLOAD_WARN_THRESHOLD,
         FAILED_REMOTE_OP_RETRIES,
         "mark_upload",
+        // TODO: use a cancellation token (https://github.com/neondatabase/neon/issues/5066)
+        backoff::Cancel::new(CancellationToken::new(), || unreachable!()),
     )
     .await
     .context("mark_upload")?;
@@ -171,6 +174,8 @@ async fn remove_tenant_remote_delete_mark(
             FAILED_UPLOAD_WARN_THRESHOLD,
             FAILED_REMOTE_OP_RETRIES,
             "remove_tenant_remote_delete_mark",
+            // TODO: use a cancellation token (https://github.com/neondatabase/neon/issues/5066)
+            backoff::Cancel::new(CancellationToken::new(), || unreachable!()),
         )
         .await
         .context("remove_tenant_remote_delete_mark")?;
@@ -252,6 +257,8 @@ pub(crate) async fn remote_delete_mark_exists(
         SHOULD_RESUME_DELETION_FETCH_MARK_ATTEMPTS,
         SHOULD_RESUME_DELETION_FETCH_MARK_ATTEMPTS,
         "fetch_tenant_deletion_mark",
+        // TODO: use a cancellation token (https://github.com/neondatabase/neon/issues/5066)
+        backoff::Cancel::new(CancellationToken::new(), || unreachable!()),
     )
     .await;
 

pageserver/src/tenant/remote_timeline_client.rs

@@ -211,6 +211,7 @@ use chrono::{NaiveDateTime, Utc};
 // re-export these
 pub use download::{is_temp_download_file, list_remote_timelines};
 use scopeguard::ScopeGuard;
+use tokio_util::sync::CancellationToken;
 use utils::backoff::{
     self, exponential_backoff, DEFAULT_BASE_BACKOFF_SECONDS, DEFAULT_MAX_BACKOFF_SECONDS,
 };
@@ -231,6 +232,7 @@ use crate::metrics::{
     RemoteTimelineClientMetricsCallTrackSize, REMOTE_ONDEMAND_DOWNLOADED_BYTES,
     REMOTE_ONDEMAND_DOWNLOADED_LAYERS,
 };
+use crate::task_mgr::shutdown_token;
 use crate::tenant::debug_assert_current_span_has_tenant_and_timeline_id;
 use crate::tenant::remote_timeline_client::index::LayerFileMetadata;
 use crate::tenant::upload_queue::Delete;
@@ -754,7 +756,7 @@ impl RemoteTimelineClient {
         pausable_failpoint!("persist_deleted_index_part");
 
         backoff::retry(
-            || async {
+            || {
                 upload::upload_index_part(
                     self.conf,
                     &self.storage_impl,
@@ -762,7 +764,6 @@ impl RemoteTimelineClient {
                     &self.timeline_id,
                     &index_part_with_deleted_at,
                 )
-                .await
             },
             |_e| false,
             1,
@@ -771,6 +772,8 @@ impl RemoteTimelineClient {
             // when executed as part of tenant deletion this happens in the background
             2,
             "persist_index_part_with_deleted_flag",
+            // TODO: use a cancellation token (https://github.com/neondatabase/neon/issues/5066)
+            backoff::Cancel::new(CancellationToken::new(), || unreachable!()),
         )
         .await?;
 
@@ -857,6 +860,7 @@ impl RemoteTimelineClient {
             FAILED_DOWNLOAD_WARN_THRESHOLD,
             FAILED_REMOTE_OP_RETRIES,
             "list_prefixes",
+            backoff::Cancel::new(shutdown_token(), || anyhow::anyhow!("Cancelled!")),
         )
         .await
         .context("list prefixes")?;
@@ -880,6 +884,7 @@ impl RemoteTimelineClient {
                 FAILED_UPLOAD_WARN_THRESHOLD,
                 FAILED_REMOTE_OP_RETRIES,
                 "delete_objects",
+                backoff::Cancel::new(shutdown_token(), || anyhow::anyhow!("Cancelled!")),
             )
             .await
             .context("delete_objects")?;
@@ -901,6 +906,7 @@ impl RemoteTimelineClient {
             FAILED_UPLOAD_WARN_THRESHOLD,
             FAILED_REMOTE_OP_RETRIES,
             "delete_index",
+            backoff::Cancel::new(shutdown_token(), || anyhow::anyhow!("Cancelled")),
         )
         .await
         .context("delete_index")?;
@@ -1134,14 +1140,13 @@ impl RemoteTimelineClient {
                     }
 
                     // sleep until it's time to retry, or we're cancelled
-                    tokio::select! {
-                        _ = task_mgr::shutdown_watcher() => { },
-                        _ = exponential_backoff(
-                            retries,
-                            DEFAULT_BASE_BACKOFF_SECONDS,
-                            DEFAULT_MAX_BACKOFF_SECONDS,
-                        ) => { },
-                    };
+                    exponential_backoff(
+                        retries,
+                        DEFAULT_BASE_BACKOFF_SECONDS,
+                        DEFAULT_MAX_BACKOFF_SECONDS,
+                        &shutdown_token(),
+                    )
+                    .await;
                 }
             }
         }

pageserver/src/tenant/remote_timeline_client/download.rs

@@ -11,6 +11,7 @@ use std::time::Duration;
 use anyhow::{anyhow, Context};
 use tokio::fs;
 use tokio::io::AsyncWriteExt;
+use tokio_util::sync::CancellationToken;
 use utils::{backoff, crashsafe};
 
 use crate::config::PageServerConf;
@@ -280,6 +281,10 @@ where
         FAILED_DOWNLOAD_WARN_THRESHOLD,
         FAILED_REMOTE_OP_RETRIES,
         description,
+        // TODO: use a cancellation token (https://github.com/neondatabase/neon/issues/5066)
+        backoff::Cancel::new(CancellationToken::new(), || -> DownloadError {
+            unreachable!()
+        }),
     )
     .await
 }

pageserver/src/tenant/timeline/walreceiver/connection_manager.rs

@@ -17,7 +17,7 @@ use crate::metrics::{
     WALRECEIVER_ACTIVE_MANAGERS, WALRECEIVER_BROKER_UPDATES, WALRECEIVER_CANDIDATES_ADDED,
     WALRECEIVER_CANDIDATES_REMOVED, WALRECEIVER_SWITCHES,
 };
-use crate::task_mgr::TaskKind;
+use crate::task_mgr::{shutdown_token, TaskKind};
 use crate::tenant::{debug_assert_current_span_has_tenant_and_timeline_id, Timeline};
 use anyhow::Context;
 use chrono::{NaiveDateTime, Utc};
@@ -211,11 +211,14 @@ async fn subscribe_for_timeline_updates(
     id: TenantTimelineId,
 ) -> Streaming<SafekeeperTimelineInfo> {
     let mut attempt = 0;
+    let cancel = shutdown_token();
+
     loop {
         exponential_backoff(
             attempt,
             DEFAULT_BASE_BACKOFF_SECONDS,
             DEFAULT_MAX_BACKOFF_SECONDS,
+            &cancel,
         )
         .await;
         attempt += 1;