[integration_test] stabilize multi-location lazy prune test

integration_test/reclaimer/multi_location_test.py

@@ -40,6 +40,8 @@ class MultiLocationTest(abc.ABC, TestBase, unittest.TestCase):
     # Two spec groups that partition the location_spec_infos
     GROUP_A = "GroupA"  # contains spec "tp0"
     GROUP_B = "GroupB"  # contains spec "tp1"
+    POLL_ATTEMPTS = 10
+    POLL_INTERVAL_SECONDS = 1
 
     def setUp(self):
         self.init_default()
@@ -128,58 +130,32 @@ def test_partial_location_loss_still_queryable(self):
         # Delete all GroupA location files
         self._delete_cache_locations(locs_a, list(range(len(locs_a))))
 
-        # Poll until tp0 is fully pruned (query triggers lazy detection).
-        # Each query submits async prune requests; retry to wait for them
-        # to take effect in the meta indexer.
-        max_attempts = 5
-        for attempt in range(max_attempts):
-            resp = self._prefix_query(block_keys)
-            if len(resp["locations"]) != 3:
-                if attempt < max_attempts - 1:
-                    time.sleep(2)
-                    continue
-                self.fail(
-                    f"blocks should still be queryable via GroupB location, "
-                    f"got {len(resp['locations'])} after {max_attempts} attempts")
-
-            all_specs = set()
-            for loc in resp["locations"]:
-                for s in loc["location_specs"]:
-                    all_specs.add(s["name"])
-
-            if "tp0" not in all_specs:
-                # tp0 fully pruned — success
-                break
-
-            if attempt < max_attempts - 1:
-                logging.info(
-                    "attempt %d: tp0 still present, waiting for async "
-                    "prune to complete...", attempt + 1)
-                time.sleep(2)
-            else:
-                self.fail(
-                    f"tp0 specs still present after {max_attempts} attempts; "
-                    f"async prune did not complete in time")
-
-        # Verify only tp1 specs remain (tp0 was pruned)
-        for loc in resp["locations"]:
-            spec_names = [s["name"] for s in loc["location_specs"]]
-            self.assertIn("tp1", spec_names,
-                          "tp1 should survive since GroupB data is intact")
-
-        # Rewrite with GroupA → should allocate new locations
-        new_locs_a = self._write_blocks_with_group(block_keys, token_ids,
-                                                   self.GROUP_A)
-        self.assertEqual(len(new_locs_a), 3,
-                         "all blocks need new GroupA locations")
+        # Query performs lazy stale-location detection and then submits async
+        # metadata pruning.  Wait for the read path to stop exposing tp0.
+        self._wait_for_prefix_specs(block_keys,
+                                    expected_location_count=3,
+                                    required_specs={"tp1"},
+                                    absent_specs={"tp0"})
+
+        # Dummy storage uses deterministic paths for the same key/spec.  Put
+        # marker files back on the old GroupA paths: if stale metadata still
+        # exists, StartWriteCache treats it as covered until the queued prune
+        # deletes the markers and removes the old metadata; if prune already
+        # finished, these markers simply become the rewritten data files.
+        self._touch_cache_locations(locs_a)
+
+        # Rewrite with GroupA.  StartWriteCache may still observe async
+        # deleting metadata, so retry until write filtering allocates all
+        # missing GroupA locations.  Failed attempts are explicitly aborted.
+        new_locs_a = self._write_blocks_with_group_retry(
+            block_keys, token_ids, self.GROUP_A, expected_location_count=3)
+        self._verify_block_keys(new_locs_a, block_keys)
 
         # Query: full coverage restored
         resp = self._prefix_query(block_keys)
         self.assertEqual(len(resp["locations"]), 3)
-        for loc in resp["locations"]:
-            spec_names = [s["name"] for s in loc["location_specs"]]
-            self.assertIn("tp0", spec_names)
-            self.assertIn("tp1", spec_names)
+        self._assert_locations_have_specs(resp["locations"],
+                                          required_specs={"tp0", "tp1"})
 
     def test_all_locations_lost_breaks_prefix(self):
         """When all locations' data is lost, prefix match breaks.
@@ -229,16 +205,58 @@ def _get_manager_client(self):
             MetaServiceHttpClient(self._http_url),
         )
 
-    def _write_blocks_with_group(self, block_keys, token_ids, group_name):
+    def _write_blocks_with_group(self, block_keys, token_ids, group_name,
+                                 expected_location_count=None):
         """Write blocks with a specific LocationSpecGroup."""
         resp = self._start_write_blocks(block_keys, token_ids,
                                         group_name=group_name)
         write_session_id = resp["write_session_id"]
         locations = resp["locations"]
+        if expected_location_count is not None:
+            self.assertEqual(
+                len(locations),
+                expected_location_count,
+                self._format_start_write_mismatch(resp, locations,
+                                                  expected_location_count))
         self._touch_cache_locations(locations)
         self._finish_write_blocks(write_session_id, len(locations))
         return locations
 
+    def _write_blocks_with_group_retry(self, block_keys, token_ids, group_name,
+                                       expected_location_count):
+        last_resp = None
+        for attempt in range(1, self.POLL_ATTEMPTS + 1):
+            resp = self._start_write_blocks(block_keys, token_ids,
+                                            group_name=group_name)
+            last_resp = resp
+            write_session_id = resp["write_session_id"]
+            locations = resp["locations"]
+            if len(locations) == expected_location_count:
+                self._touch_cache_locations(locations)
+                self._finish_write_blocks(write_session_id, len(locations))
+                return locations
+
+            self._touch_cache_locations(locations)
+            self._finish_write_blocks(write_session_id, len(locations),
+                                      success=False)
+            self._wait_for_locations_absent(locations)
+            mismatch = self._format_start_write_mismatch(
+                resp, locations, expected_location_count)
+            logging.info(
+                "attempt %d/%d: rewrite did not allocate all locations, "
+                "aborted session and will retry: %s",
+                attempt,
+                self.POLL_ATTEMPTS,
+                mismatch)
+            if attempt < self.POLL_ATTEMPTS:
+                time.sleep(self.POLL_INTERVAL_SECONDS)
+
+        mismatch = self._format_start_write_mismatch(
+            last_resp, last_resp["locations"], expected_location_count)
+        self.fail(
+            "rewrite did not allocate expected locations after retries: "
+            f"{mismatch}")
+
     def _start_write_blocks(self, block_keys, token_ids, group_name=None):
         req = {
             "trace_id": self._trace_id,
@@ -251,13 +269,13 @@ def _start_write_blocks(self, block_keys, token_ids, group_name=None):
             req["location_spec_group_names"] = [group_name] * len(block_keys)
         return self._client.start_write_cache(req)
 
-    def _finish_write_blocks(self, write_session_id, loc_sz):
+    def _finish_write_blocks(self, write_session_id, loc_sz, success=True):
         return self._client.finish_write_cache({
             "trace_id": self._trace_id,
             "instance_id": self._instance_id,
             "write_session_id": write_session_id,
             "success_blocks": {
-                "bool_masks": {"values": [True] * loc_sz},
+                "bool_masks": {"values": [success] * loc_sz},
             },
         })
 
@@ -270,18 +288,110 @@ def _prefix_query(self, block_keys):
             "block_mask": {"offset": 0},
         })
 
+    def _wait_for_prefix_specs(self, block_keys, expected_location_count,
+                               required_specs=None, absent_specs=None):
+        required_specs = required_specs or set()
+        absent_specs = absent_specs or set()
+        last_resp = None
+        for attempt in range(1, self.POLL_ATTEMPTS + 1):
+            last_resp = self._prefix_query(block_keys)
+            locations = last_resp["locations"]
+            if (len(locations) == expected_location_count
+                    and self._locations_match_specs(locations,
+                                                    required_specs,
+                                                    absent_specs)):
+                return last_resp
+
+            logging.info(
+                "attempt %d/%d: waiting for prefix specs, locations=%s",
+                attempt,
+                self.POLL_ATTEMPTS,
+                self._summarize_locations(locations))
+            time.sleep(self.POLL_INTERVAL_SECONDS)
+
+        self.fail(
+            "prefix query did not reach expected specs: "
+            f"expected_location_count={expected_location_count}, "
+            f"required_specs={sorted(required_specs)}, "
+            f"absent_specs={sorted(absent_specs)}, "
+            f"last_locations={self._summarize_locations(last_resp['locations'])}")
+
+    def _assert_locations_have_specs(self, locations, required_specs=None,
+                                     absent_specs=None):
+        required_specs = required_specs or set()
+        absent_specs = absent_specs or set()
+        self.assertTrue(
+            self._locations_match_specs(locations, required_specs,
+                                        absent_specs),
+            f"locations specs mismatch: "
+            f"required_specs={sorted(required_specs)}, "
+            f"absent_specs={sorted(absent_specs)}, "
+            f"locations={self._summarize_locations(locations)}")
+
+    @staticmethod
+    def _locations_match_specs(locations, required_specs, absent_specs):
+        for loc in locations:
+            spec_names = {s["name"] for s in loc.get("location_specs", [])}
+            if not required_specs.issubset(spec_names):
+                return False
+            if absent_specs.intersection(spec_names):
+                return False
+        return True
+
+    @staticmethod
+    def _summarize_locations(locations):
+        return [
+            [s.get("name") for s in loc.get("location_specs", [])]
+            for loc in locations
+        ]
+
+    @staticmethod
+    def _format_start_write_mismatch(resp, locations, expected_count):
+        uris = [
+            spec.get("uri")
+            for loc in locations
+            for spec in loc.get("location_specs", [])
+        ]
+        return (
+            f"start_write_cache returned {len(locations)} locations, "
+            f"expected {expected_count}; "
+            f"block_mask={resp.get('block_mask')}, uris={uris}")
+
+    def _wait_for_locations_absent(self, locations):
+        paths = self._location_paths(locations)
+        for attempt in range(1, self.POLL_ATTEMPTS + 1):
+            remaining = [path for path in paths if os.path.exists(path)]
+            if not remaining:
+                return
+            logging.info(
+                "attempt %d/%d: waiting for aborted write cleanup, "
+                "remaining_paths=%s",
+                attempt,
+                self.POLL_ATTEMPTS,
+                remaining)
+            if attempt < self.POLL_ATTEMPTS:
+                time.sleep(self.POLL_INTERVAL_SECONDS)
+
+        self.fail(f"aborted write cleanup did not delete paths: {remaining}")
+
+    @staticmethod
+    def _location_paths(locations):
+        return [
+            urlparse(spec["uri"]).path
+            for loc in locations
+            for spec in loc.get("location_specs", [])
+        ]
+
     @staticmethod
     def _touch_cache_locations(locations):
         """Simulate the cache data write by creating data files."""
-        for loc in locations:
-            for spec in loc.get("location_specs", []):
-                file_path = urlparse(spec["uri"]).path
-                try:
-                    os.utime(file_path)
-                except FileNotFoundError:
-                    os.makedirs(os.path.dirname(file_path), exist_ok=True)
-                    with open(file_path, 'x') as _:
-                        pass
+        for file_path in MultiLocationTest._location_paths(locations):
+            try:
+                os.utime(file_path)
+            except FileNotFoundError:
+                os.makedirs(os.path.dirname(file_path), exist_ok=True)
+                with open(file_path, 'x') as _:
+                    pass
 
     @staticmethod
     def _delete_cache_locations(locations, indices):
@@ -293,6 +403,14 @@ def _delete_cache_locations(locations, indices):
                 if os.path.exists(file_path):
                     os.remove(file_path)
 
+    def _verify_block_keys(self, locations, block_keys):
+        self.assertEqual(len(locations), len(block_keys))
+        for loc, key in zip(locations, block_keys):
+            for spec in loc.get("location_specs", []):
+                file_path = urlparse(spec["uri"]).path
+                self.assertEqual(int(os.path.basename(file_path), base=16),
+                                 key)
+
     def _make_dummy_storage(self):
         dummy_root_path = f"{self.get_workdir()}/{self._storage_name}/data/"
         add_storage_req = {