storage_controller: coordinate imports across shards in the storage controller (#11345)

## Problem

Pageservers notify control plane directly when a shard import has
completed.
Control plane has to download the status of each shard from S3 and
figure out if everything is truly done,
before proceeding with branch activation.

Issues with this approach are:
* We can't control shard split behaviour on the storage controller side.
It's unsafe to split
during import.
* Control plane needs to know about shards and implement logic to check
all timelines are indeed ready.

## Summary of changes

In short, storage controller coordinates imports, and, only when
everything is done, notifies control plane.

Big rocks:
1. Store timeline imports in the storage controller database. Each
import stores the status of its shards in the database.
We hook into the timeline creation call as our entry point for this.
2. Pageservers get a new upcall endpoint to notify the storage
controller of shard import updates.
3. Storage controller handles these updates by updating persisted state.
If an update finalizes the import,
then poll pageservers until timeline activation, and, then, notify the
control plane that the import is complete.

Cplane side change with new endpoint is in
https://github.com/neondatabase/cloud/pull/26166

Closes https://github.com/neondatabase/neon/issues/11566

test_runner/regress/test_import_pgdata.py

@@ -1,9 +1,9 @@
 import base64
 import json
-import re
 import time
 from enum import Enum
 from pathlib import Path
+from threading import Event
 
 import psycopg2
 import psycopg2.errors
@@ -14,12 +14,11 @@ from fixtures.log_helper import log
 from fixtures.neon_fixtures import NeonEnvBuilder, PgBin, PgProtocol, VanillaPostgres
 from fixtures.pageserver.http import (
     ImportPgdataIdemptencyKey,
-    PageserverApiException,
 )
 from fixtures.pg_version import PgVersion
 from fixtures.port_distributor import PortDistributor
 from fixtures.remote_storage import MockS3Server, RemoteStorageKind
-from fixtures.utils import shared_buffers_for_max_cu
+from fixtures.utils import shared_buffers_for_max_cu, skip_in_debug_build, wait_until
 from mypy_boto3_kms import KMSClient
 from mypy_boto3_kms.type_defs import EncryptResponseTypeDef
 from mypy_boto3_s3 import S3Client
@@ -44,6 +43,7 @@ smoke_params = [
 ]
 
 
+@skip_in_debug_build("MULTIPLE_RELATION_SEGMENTS has non trivial amount of data")
 @pytest.mark.parametrize("shard_count,stripe_size,rel_block_size", smoke_params)
 def test_pgdata_import_smoke(
     vanilla_pg: VanillaPostgres,
@@ -56,24 +56,29 @@ def test_pgdata_import_smoke(
     #
     # Setup fake control plane for import progress
     #
+    import_completion_signaled = Event()
+
     def handler(request: Request) -> Response:
-        log.info(f"control plane request: {request.json}")
+        log.info(f"control plane /import_complete request: {request.json}")
+        import_completion_signaled.set()
         return Response(json.dumps({}), status=200)
 
     cplane_mgmt_api_server = make_httpserver
-    cplane_mgmt_api_server.expect_request(re.compile(".*")).respond_with_handler(handler)
+    cplane_mgmt_api_server.expect_request(
+        "/storage/api/v1/import_complete", method="PUT"
+    ).respond_with_handler(handler)
 
     neon_env_builder.enable_pageserver_remote_storage(RemoteStorageKind.LOCAL_FS)
+
+    neon_env_builder.control_plane_hooks_api = (
+        f"http://{cplane_mgmt_api_server.host}:{cplane_mgmt_api_server.port}/storage/api/v1/"
+    )
+
     env = neon_env_builder.init_start()
 
     # The test needs LocalFs support, which is only built in testing mode.
     env.pageserver.is_testing_enabled_or_skip()
 
-    env.pageserver.patch_config_toml_nonrecursive(
-        {
-            "import_pgdata_upcall_api": f"http://{cplane_mgmt_api_server.host}:{cplane_mgmt_api_server.port}/path/to/mgmt/api"
-        }
-    )
     env.pageserver.stop()
     env.pageserver.start()
 
@@ -193,40 +198,11 @@ def test_pgdata_import_smoke(
     )
     env.neon_cli.mappings_map_branch(import_branch_name, tenant_id, timeline_id)
 
-    while True:
-        locations = env.storage_controller.locate(tenant_id)
-        active_count = 0
-        for location in locations:
-            shard_id = TenantShardId.parse(location["shard_id"])
-            ps = env.get_pageserver(location["node_id"])
-            try:
-                detail = ps.http_client().timeline_detail(shard_id, timeline_id)
-                state = detail["state"]
-                log.info(f"shard {shard_id} state: {state}")
-                if state == "Active":
-                    active_count += 1
-            except PageserverApiException as e:
-                if e.status_code == 404:
-                    log.info("not found, import is in progress")
-                    continue
-                elif e.status_code == 429:
-                    log.info("import is in progress")
-                    continue
-                else:
-                    raise
+    def cplane_notified():
+        assert import_completion_signaled.is_set()
 
-            shard_status_file = statusdir / f"shard-{shard_id.shard_index}"
-            if state == "Active":
-                shard_status_file_contents = (
-                    shard_status_file.read_text()
-                )  # Active state implies import is done
-                shard_status = json.loads(shard_status_file_contents)
-                assert shard_status["done"] is True
-
-        if active_count == len(locations):
-            log.info("all shards are active")
-            break
-        time.sleep(1)
+    # Generous timeout for the MULTIPLE_RELATION_SEGMENTS test variants
+    wait_until(cplane_notified, timeout=90)
 
     import_duration = time.monotonic() - start
     log.info(f"import complete; duration={import_duration:.2f}s")
@@ -372,19 +348,27 @@ def test_fast_import_with_pageserver_ingest(
     vanilla_pg.safe_psql("CREATE TABLE foo (a int); INSERT INTO foo SELECT generate_series(1, 10);")
 
     # Setup pageserver and fake cplane for import progress
+    import_completion_signaled = Event()
+
     def handler(request: Request) -> Response:
-        log.info(f"control plane request: {request.json}")
+        log.info(f"control plane /import_complete request: {request.json}")
+        import_completion_signaled.set()
         return Response(json.dumps({}), status=200)
 
     cplane_mgmt_api_server = make_httpserver
-    cplane_mgmt_api_server.expect_request(re.compile(".*")).respond_with_handler(handler)
+    cplane_mgmt_api_server.expect_request(
+        "/storage/api/v1/import_complete", method="PUT"
+    ).respond_with_handler(handler)
+
+    neon_env_builder.control_plane_hooks_api = (
+        f"http://{cplane_mgmt_api_server.host}:{cplane_mgmt_api_server.port}/storage/api/v1/"
+    )
 
     neon_env_builder.enable_pageserver_remote_storage(RemoteStorageKind.MOCK_S3)
     env = neon_env_builder.init_start()
 
     env.pageserver.patch_config_toml_nonrecursive(
         {
-            "import_pgdata_upcall_api": f"http://{cplane_mgmt_api_server.host}:{cplane_mgmt_api_server.port}/path/to/mgmt/api",
             # because import_pgdata code uses this endpoint, not the one in common remote storage config
             # TODO: maybe use common remote_storage config in pageserver?
             "import_pgdata_aws_endpoint_url": env.s3_mock_server.endpoint(),
@@ -476,42 +460,10 @@ def test_fast_import_with_pageserver_ingest(
         conn = PgProtocol(dsn=f"postgresql://cloud_admin@localhost:{pg_port}/neondb")
         validate_vanilla_equivalence(conn)
 
-    # Poll pageserver statuses in s3
-    while True:
-        locations = env.storage_controller.locate(tenant_id)
-        active_count = 0
-        for location in locations:
-            shard_id = TenantShardId.parse(location["shard_id"])
-            ps = env.get_pageserver(location["node_id"])
-            try:
-                detail = ps.http_client().timeline_detail(shard_id, timeline_id)
-                log.info(f"timeline {tenant_id}/{timeline_id} detail: {detail}")
-                state = detail["state"]
-                log.info(f"shard {shard_id} state: {state}")
-                if state == "Active":
-                    active_count += 1
-            except PageserverApiException as e:
-                if e.status_code == 404:
-                    log.info("not found, import is in progress")
-                    continue
-                elif e.status_code == 429:
-                    log.info("import is in progress")
-                    continue
-                else:
-                    raise
+    def cplane_notified():
+        assert import_completion_signaled.is_set()
 
-            if state == "Active":
-                key = f"{key_prefix}/status/shard-{shard_id.shard_index}"
-                shard_status_file_contents = (
-                    mock_s3_client.get_object(Bucket=bucket, Key=key)["Body"].read().decode("utf-8")
-                )
-                shard_status = json.loads(shard_status_file_contents)
-                assert shard_status["done"] is True
-
-        if active_count == len(locations):
-            log.info("all shards are active")
-            break
-        time.sleep(0.5)
+    wait_until(cplane_notified, timeout=60)
 
     import_duration = time.monotonic() - start
     log.info(f"import complete; duration={import_duration:.2f}s")