refactor(pageserver): propagate RequestContext to layer downloads (#11001)

For some reason the layer download API never fully got
`RequestContext`-infected.

This PR fixes that as a precursor to
- https://github.com/neondatabase/neon/issues/6107
This commit is contained in:
Christian Schwarz
2025-02-27 10:26:25 +01:00
committed by GitHub
parent c92a36740b
commit f09843ef17
11 changed files with 84 additions and 53 deletions

View File

@@ -927,11 +927,10 @@ async fn get_lsn_by_timestamp_handler(
let with_lease = parse_query_param(&request, "with_lease")?.unwrap_or(false);
let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
let timeline =
active_timeline_of_active_tenant(&state.tenant_manager, tenant_shard_id, timeline_id)
.await?;
let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
let result = timeline
.find_lsn_for_timestamp(timestamp_pg, &cancel, &ctx)
.await?;
@@ -1000,10 +999,10 @@ async fn get_timestamp_of_lsn_handler(
.with_context(|| format!("Invalid LSN: {lsn_str:?}"))
.map_err(ApiError::BadRequest)?;
let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
let timeline =
active_timeline_of_active_tenant(&state.tenant_manager, tenant_shard_id, timeline_id)
.await?;
let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
let result = timeline.get_timestamp_for_lsn(lsn, &ctx).await?;
match result {
@@ -1368,7 +1367,7 @@ async fn timeline_layer_scan_disposable_keys(
};
let resident_layer = layer
.download_and_keep_resident()
.download_and_keep_resident(&ctx)
.await
.map_err(|err| match err {
tenant::storage_layer::layer::DownloadError::TimelineShutdown
@@ -1443,6 +1442,7 @@ async fn timeline_download_heatmap_layers_handler(
let timeline =
active_timeline_of_active_tenant(&state.tenant_manager, tenant_shard_id, timeline_id)
.await?;
let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
let max_concurrency = get_config(&request)
.remote_storage_config
@@ -1451,7 +1451,9 @@ async fn timeline_download_heatmap_layers_handler(
.unwrap_or(DEFAULT_MAX_CONCURRENCY);
let concurrency = std::cmp::min(max_concurrency, desired_concurrency);
timeline.start_heatmap_layers_download(concurrency).await?;
timeline
.start_heatmap_layers_download(concurrency, &ctx)
.await?;
json_response(StatusCode::ACCEPTED, ())
}
@@ -1490,8 +1492,9 @@ async fn layer_download_handler(
let timeline =
active_timeline_of_active_tenant(&state.tenant_manager, tenant_shard_id, timeline_id)
.await?;
let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
let downloaded = timeline
.download_layer(&layer_name)
.download_layer(&layer_name, &ctx)
.await
.map_err(|e| match e {
tenant::storage_layer::layer::DownloadError::TimelineShutdown
@@ -2389,7 +2392,8 @@ async fn timeline_download_remote_layers_handler_post(
let timeline =
active_timeline_of_active_tenant(&state.tenant_manager, tenant_shard_id, timeline_id)
.await?;
match timeline.spawn_download_all_remote_layers(body).await {
let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
match timeline.spawn_download_all_remote_layers(body, &ctx).await {
Ok(st) => json_response(StatusCode::ACCEPTED, st),
Err(st) => json_response(StatusCode::CONFLICT, st),
}