pageserver: add and stabilize import chaos test (neondatabase#11982)

## Problem

Test coverage of timeline imports is lacking.

## Summary of changes

This PR adds a chaos import test. It runs an import while injecting
various chaos events
in the environment. All the commits that follow the test fix various
issues that were surfaced by it.

Closes neondatabase#10191

Author: VladLazar

pageserver/src/http/routes.rs

@@ -370,6 +370,18 @@ impl From<crate::tenant::secondary::SecondaryTenantError> for ApiError {
     }
 }
 
+impl From<crate::tenant::FinalizeTimelineImportError> for ApiError {
+    fn from(err: crate::tenant::FinalizeTimelineImportError) -> ApiError {
+        use crate::tenant::FinalizeTimelineImportError::*;
+        match err {
+            ImportTaskStillRunning => {
+                ApiError::ResourceUnavailable("Import task still running".into())
+            }
+            ShuttingDown => ApiError::ShuttingDown,
+        }
+    }
+}
+
 // Helper function to construct a TimelineInfo struct for a timeline
 async fn build_timeline_info(
     timeline: &Arc<Timeline>,
@@ -3533,10 +3545,7 @@ async fn activate_post_import_handler(
 
         tenant.wait_to_become_active(ACTIVE_TENANT_TIMEOUT).await?;
 
-        tenant
-            .finalize_importing_timeline(timeline_id)
-            .await
-            .map_err(ApiError::InternalServerError)?;
+        tenant.finalize_importing_timeline(timeline_id).await?;
 
         match tenant.get_timeline(timeline_id, false) {
             Ok(_timeline) => {

pageserver/src/tenant.rs

@@ -864,6 +864,14 @@ impl Debug for SetStoppingError {
     }
 }
 
+#[derive(thiserror::Error, Debug)]
+pub(crate) enum FinalizeTimelineImportError {
+    #[error("Import task not done yet")]
+    ImportTaskStillRunning,
+    #[error("Shutting down")]
+    ShuttingDown,
+}
+
 /// Arguments to [`TenantShard::create_timeline`].
 ///
 /// Not usable as an idempotency key for timeline creation because if [`CreateTimelineParamsBranch::ancestor_start_lsn`]
@@ -1150,10 +1158,20 @@ impl TenantShard {
             ctx,
         )?;
         let disk_consistent_lsn = timeline.get_disk_consistent_lsn();
-        anyhow::ensure!(
-            disk_consistent_lsn.is_valid(),
-            "Timeline {tenant_id}/{timeline_id} has invalid disk_consistent_lsn"
-        );
+
+        if !disk_consistent_lsn.is_valid() {
+            // As opposed to normal timelines which get initialised with a disk consitent LSN
+            // via initdb, imported timelines start from 0. If the import task stops before
+            // it advances disk consitent LSN, allow it to resume.
+            let in_progress_import = import_pgdata
+                .as_ref()
+                .map(|import| !import.is_done())
+                .unwrap_or(false);
+            if !in_progress_import {
+                anyhow::bail!("Timeline {tenant_id}/{timeline_id} has invalid disk_consistent_lsn");
+            }
+        }
+
         assert_eq!(
             disk_consistent_lsn,
             metadata.disk_consistent_lsn(),
@@ -1247,20 +1265,25 @@ impl TenantShard {
                     }
                 }
 
-                // Sanity check: a timeline should have some content.
-                anyhow::ensure!(
-                    ancestor.is_some()
-                        || timeline
-                            .layers
-                            .read()
-                            .await
-                            .layer_map()
-                            .expect("currently loading, layer manager cannot be shutdown already")
-                            .iter_historic_layers()
-                            .next()
-                            .is_some(),
-                    "Timeline has no ancestor and no layer files"
-                );
+                if disk_consistent_lsn.is_valid() {
+                    // Sanity check: a timeline should have some content.
+                    // Exception: importing timelines might not yet have any
+                    anyhow::ensure!(
+                        ancestor.is_some()
+                            || timeline
+                                .layers
+                                .read()
+                                .await
+                                .layer_map()
+                                .expect(
+                                    "currently loading, layer manager cannot be shutdown already"
+                                )
+                                .iter_historic_layers()
+                                .next()
+                                .is_some(),
+                        "Timeline has no ancestor and no layer files"
+                    );
+                }
 
                 Ok(TimelineInitAndSyncResult::ReadyToActivate)
             }
@@ -2860,13 +2883,13 @@ impl TenantShard {
     pub(crate) async fn finalize_importing_timeline(
         &self,
         timeline_id: TimelineId,
-    ) -> anyhow::Result<()> {
+    ) -> Result<(), FinalizeTimelineImportError> {
         let timeline = {
             let locked = self.timelines_importing.lock().unwrap();
             match locked.get(&timeline_id) {
                 Some(importing_timeline) => {
                     if !importing_timeline.import_task_handle.is_finished() {
-                        return Err(anyhow::anyhow!("Import task not done yet"));
+                        return Err(FinalizeTimelineImportError::ImportTaskStillRunning);
                     }
 
                     importing_timeline.timeline.clone()
@@ -2879,8 +2902,13 @@ impl TenantShard {
 
         timeline
             .remote_client
-            .schedule_index_upload_for_import_pgdata_finalize()?;
-        timeline.remote_client.wait_completion().await?;
+            .schedule_index_upload_for_import_pgdata_finalize()
+            .map_err(|_err| FinalizeTimelineImportError::ShuttingDown)?;
+        timeline
+            .remote_client
+            .wait_completion()
+            .await
+            .map_err(|_err| FinalizeTimelineImportError::ShuttingDown)?;
 
         self.timelines_importing
             .lock()
@@ -3484,8 +3512,9 @@ impl TenantShard {
             let mut timelines_importing = self.timelines_importing.lock().unwrap();
             timelines_importing
                 .drain()
-                .for_each(|(_timeline_id, importing_timeline)| {
-                    importing_timeline.shutdown();
+                .for_each(|(timeline_id, importing_timeline)| {
+                    let span = tracing::info_span!("importing_timeline_shutdown", %timeline_id);
+                    js.spawn(async move { importing_timeline.shutdown().instrument(span).await });
                 });
         }
         // test_long_timeline_create_then_tenant_delete is leaning on this message

pageserver/src/tenant/timeline/import_pgdata.rs

@@ -25,8 +25,11 @@ pub(crate) struct ImportingTimeline {
 }
 
 impl ImportingTimeline {
-    pub(crate) fn shutdown(self) {
+    pub(crate) async fn shutdown(self) {
         self.import_task_handle.abort();
+        let _ = self.import_task_handle.await;
+
+        self.timeline.remote_client.shutdown().await;
     }
 }
 
@@ -93,6 +96,11 @@ pub async fn doit(
                 );
             }
 
+            timeline
+                .remote_client
+                .schedule_index_upload_for_file_changes()?;
+            timeline.remote_client.wait_completion().await?;
+
             // Communicate that shard is done.
             // Ensure at-least-once delivery of the upcall to storage controller
             // before we mark the task as done and never come here again.

pageserver/src/tenant/timeline/import_pgdata/flow.rs

@@ -113,14 +113,14 @@ async fn run_v1(
     let plan_hash = hasher.finish();
 
     if let Some(progress) = &import_progress {
-        if plan_hash != progress.import_plan_hash {
-            anyhow::bail!("Import plan does not match storcon metadata");
-        }
-
         // Handle collisions on jobs of unequal length
         if progress.jobs != plan.jobs.len() {
             anyhow::bail!("Import plan job length does not match storcon metadata")
         }
+
+        if plan_hash != progress.import_plan_hash {
+            anyhow::bail!("Import plan does not match storcon metadata");
+        }
     }
 
     pausable_failpoint!("import-timeline-pre-execute-pausable");
@@ -218,6 +218,19 @@ impl Planner {
                 checkpoint_buf,
             )));
 
+        // Sort the tasks by the key ranges they handle.
+        // The plan being generated here needs to be stable across invocations
+        // of this method.
+        self.tasks.sort_by_key(|task| match task {
+            AnyImportTask::SingleKey(key) => (key.key, key.key.next()),
+            AnyImportTask::RelBlocks(rel_blocks) => {
+                (rel_blocks.key_range.start, rel_blocks.key_range.end)
+            }
+            AnyImportTask::SlruBlocks(slru_blocks) => {
+                (slru_blocks.key_range.start, slru_blocks.key_range.end)
+            }
+        });
+
         // Assigns parts of key space to later parallel jobs
         let mut last_end_key = Key::MIN;
         let mut current_chunk = Vec::new();
@@ -426,6 +439,8 @@ impl Plan {
                     }));
                 },
                 maybe_complete_job_idx = work.next() => {
+                    pausable_failpoint!("import-task-complete-pausable");
+
                     match maybe_complete_job_idx {
                         Some(Ok((job_idx, res))) => {
                             assert!(last_completed_job_idx.checked_add(1).unwrap() == job_idx);
@@ -440,6 +455,9 @@ impl Plan {
                                     import_plan_hash,
                                 };
 
+                                timeline.remote_client.schedule_index_upload_for_file_changes()?;
+                                timeline.remote_client.wait_completion().await?;
+
                                 storcon_client.put_timeline_import_status(
                                     timeline.tenant_shard_id,
                                     timeline.timeline_id,
@@ -640,7 +658,11 @@ impl Hash for ImportSingleKeyTask {
         let ImportSingleKeyTask { key, buf } = self;
 
         key.hash(state);
-        buf.hash(state);
+        // The key value might not have a stable binary representation.
+        // For instance, the db directory uses an unstable hash-map.
+        // To work around this we are a bit lax here and only hash the
+        // size of the buffer which must be consistent.
+        buf.len().hash(state);
     }
 }
 
@@ -915,7 +937,7 @@ impl ChunkProcessingJob {
                 let guard = timeline.layers.read().await;
                 let existing_layer = guard.try_get_from_key(&desc.key());
                 if let Some(layer) = existing_layer {
-                    if layer.metadata().generation != timeline.generation {
+                    if layer.metadata().generation == timeline.generation {
                         return Err(anyhow::anyhow!(
                             "Import attempted to rewrite layer file in the same generation: {}",
                             layer.local_path()

storage_controller/src/service.rs

@@ -3922,6 +3922,11 @@ impl Service {
         })
     }
 
+    #[instrument(skip_all, fields(
+        tenant_id=%req.tenant_shard_id.tenant_id,
+        shard_id=%req.tenant_shard_id.shard_slug(),
+        timeline_id=%req.timeline_id,
+    ))]
     pub(crate) async fn handle_timeline_shard_import_progress(
         self: &Arc<Self>,
         req: TimelineImportStatusRequest,
@@ -3971,6 +3976,11 @@ impl Service {
             })
     }
 
+    #[instrument(skip_all, fields(
+        tenant_id=%req.tenant_shard_id.tenant_id,
+        shard_id=%req.tenant_shard_id.shard_slug(),
+        timeline_id=%req.timeline_id,
+    ))]
     pub(crate) async fn handle_timeline_shard_import_progress_upcall(
         self: &Arc<Self>,
         req: PutTimelineImportStatusRequest,

test_runner/fixtures/neon_fixtures.py

@@ -404,6 +404,29 @@ def to_config_key_value(self) -> tuple[str, dict[str, Any]]:
         return ("tracing", value)
 
 
+@dataclass
+class PageserverImportConfig:
+    import_job_concurrency: int
+    import_job_soft_size_limit: int
+    import_job_checkpoint_threshold: int
+
+    @staticmethod
+    def default() -> PageserverImportConfig:
+        return PageserverImportConfig(
+            import_job_concurrency=4,
+            import_job_soft_size_limit=512 * 1024,
+            import_job_checkpoint_threshold=4,
+        )
+
+    def to_config_key_value(self) -> tuple[str, dict[str, Any]]:
+        value = {
+            "import_job_concurrency": self.import_job_concurrency,
+            "import_job_soft_size_limit": self.import_job_soft_size_limit,
+            "import_job_checkpoint_threshold": self.import_job_checkpoint_threshold,
+        }
+        return ("timeline_import_config", value)
+
+
 class NeonEnvBuilder:
     """
     Builder object to create a Neon runtime environment
@@ -454,6 +477,7 @@ def __init__(
         pageserver_wal_receiver_protocol: PageserverWalReceiverProtocol | None = None,
         pageserver_get_vectored_concurrent_io: str | None = None,
         pageserver_tracing_config: PageserverTracingConfig | None = None,
+        pageserver_import_config: PageserverImportConfig | None = None,
     ):
         self.repo_dir = repo_dir
         self.rust_log_override = rust_log_override
@@ -511,6 +535,7 @@ def __init__(
         )
 
         self.pageserver_tracing_config = pageserver_tracing_config
+        self.pageserver_import_config = pageserver_import_config
 
         self.pageserver_default_tenant_config_compaction_algorithm: dict[str, Any] | None = (
             pageserver_default_tenant_config_compaction_algorithm
@@ -1179,6 +1204,10 @@ def __init__(self, config: NeonEnvBuilder):
         self.pageserver_wal_receiver_protocol = config.pageserver_wal_receiver_protocol
         self.pageserver_get_vectored_concurrent_io = config.pageserver_get_vectored_concurrent_io
         self.pageserver_tracing_config = config.pageserver_tracing_config
+        if config.pageserver_import_config is None:
+            self.pageserver_import_config = PageserverImportConfig.default()
+        else:
+            self.pageserver_import_config = config.pageserver_import_config
 
         # Create the neon_local's `NeonLocalInitConf`
         cfg: dict[str, Any] = {
@@ -1258,12 +1287,6 @@ def __init__(self, config: NeonEnvBuilder):
                 "no_sync": True,
                 # Look for gaps in WAL received from safekeepeers
                 "validate_wal_contiguity": True,
-                # TODO(vlad): make these configurable through the builder
-                "timeline_import_config": {
-                    "import_job_concurrency": 4,
-                    "import_job_soft_size_limit": 512 * 1024,
-                    "import_job_checkpoint_threshold": 4,
-                },
             }
 
             # Batching (https://github.com/neondatabase/neon/issues/9377):
@@ -1325,6 +1348,12 @@ def __init__(self, config: NeonEnvBuilder):
 
                 ps_cfg[key] = value
 
+            if self.pageserver_import_config is not None:
+                key, value = self.pageserver_import_config.to_config_key_value()
+
+                if key not in ps_cfg:
+                    ps_cfg[key] = value
+
             # Create a corresponding NeonPageserver object
             ps = NeonPageserver(
                 self, ps_id, port=pageserver_port, az_id=ps_cfg["availability_zone"]