Retry once when the batch signer returns a partial map

src/ndi/cloud/batch_signed_url.py

@@ -46,6 +46,13 @@
 # MATLAB counterpart's testFailingSignerReturnsEmpty.
 DEFAULT_FAILURE_TTL_SECONDS = 60
 
+# 500 ms: pause before re-fetching a scope whose first answer was a
+# partial map -- the scope was populated, but it did not name the uid the
+# caller asked about. On some cloud environments the batch endpoint lags
+# briefly behind ``waitForAllBulkUploads``, and the second call names the
+# whole set. See Waltham-Data-Science/NDI-python#309.
+DEFAULT_PARTIAL_MAP_RETRY_SECONDS = 0.5
+
 
 @dataclass
 class _CacheEntry:
@@ -75,7 +82,7 @@ class Stats:
 
     Fields:
         signer_calls: How many times the batch endpoint was actually
-            called (one per scope fetch).
+            called (one per scope fetch, plus one per partial-map retry).
         uid_hits: uids answered from a batch map.
         uid_misses: uids the batch could NOT answer, each of which sends
             the caller to the per-uid getFileDetails fallback.
@@ -88,6 +95,12 @@ class Stats:
             usable map. Four unrelated causes end there; reporting them as
             one silent miss leaves the endpoint owner with nothing to act
             on. See NDI-matlab#968.
+        partial_map_retries: How many scopes were re-fetched because the
+            first answer was populated but did not name a requested uid.
+            One retry per scope, ever, whether it settled or not; a caller
+            that wants to know "does this scope EVER answer for this uid"
+            gets that answer after the retry. See Waltham-Data-Science/
+            NDI-python#309.
     """
 
     signer_calls: int = 0
@@ -96,6 +109,7 @@ class Stats:
     last_map_size: int = 0
     last_map_uids: list[str] = field(default_factory=list)
     last_failure_reason: str = ""
+    partial_map_retries: int = 0
 
 
 # The default signer walks every page. Injected for tests.
@@ -146,12 +160,19 @@ def __init__(
         signer: Callable[..., tuple[bool, dict[str, Any]]] | None = None,
         ttl_seconds: float = DEFAULT_TTL_SECONDS,
         failure_ttl_seconds: float = DEFAULT_FAILURE_TTL_SECONDS,
+        partial_map_retry_seconds: float = DEFAULT_PARTIAL_MAP_RETRY_SECONDS,
+        sleep: Callable[[float], None] | None = None,
     ) -> None:
         self._signer = signer or _default_signer
         self._ttl_seconds = ttl_seconds
         self._failure_ttl_seconds = failure_ttl_seconds
+        self._partial_map_retry_seconds = partial_map_retry_seconds
+        # Injected only for tests: the retry has a real wait, and a test
+        # that runs it 20 times should not pay 10 seconds for it.
+        self._sleep = sleep or time.sleep
         self._cache: dict[str, _CacheEntry] = {}
         self._failed_scopes: dict[str, _FailureEntry] = {}
+        self._retried_scopes: set[str] = set()
         self._warned: set[str] = set()
         self._stats = Stats()
         # A batch fetch is IO-bound and slow, and download handlers can be
@@ -235,9 +256,45 @@ def lookup(
                 self._stats.uid_hits += 1
                 return url
 
-            # The scope was fetched but does not name this uid -- data
-            # drift, or a scope that does not cover the file. The caller
-            # falls back per uid.
+            # The scope was fetched but does not name this uid. Two
+            # distinct causes: data drift (this file was never in the
+            # scope), or a lagged index (the file IS in the scope on the
+            # server, but the batch endpoint's map has not caught up yet).
+            # Only the second is worth a retry, and this side has no way
+            # to distinguish them a priori -- so retry ONCE per scope,
+            # sleep briefly first, and if the second answer still does not
+            # name the uid, take that as data drift and warn.
+            #
+            # Bounded and per-scope: a data-drift miss costs one extra
+            # signer call for the whole scope, not one per uid. A lagged
+            # index that recovers replaces the cache entry for every
+            # subsequent uid in the same scope, so a whole series' worth
+            # of misses becomes a whole series' worth of hits from one
+            # retry. See Waltham-Data-Science/NDI-python#309.
+            if cache_key not in self._retried_scopes:
+                self._retried_scopes.add(cache_key)
+                if self._partial_map_retry_seconds > 0:
+                    self._sleep(self._partial_map_retry_seconds)
+                # Drop the stale entry so _fetch_scope replaces it fresh.
+                self._cache.pop(cache_key, None)
+                self._stats.partial_map_retries += 1
+                refreshed = self._fetch_scope(
+                    cache_key,
+                    cloud_dataset_id,
+                    ndi_document_id,
+                    series_name,
+                    uid,
+                    client=client,
+                )
+                if refreshed is None:
+                    # _fetch_scope already recorded the miss and warned.
+                    # Do not double-count here.
+                    return ""
+                url = refreshed.files.get(uid, "")
+                if url:
+                    self._stats.uid_hits += 1
+                    return url
+
             self._stats.uid_misses += 1
             self._warn_once(cache_key)
             return ""
@@ -252,13 +309,15 @@ def stats(self) -> Stats:
                 last_map_size=self._stats.last_map_size,
                 last_map_uids=list(self._stats.last_map_uids),
                 last_failure_reason=self._stats.last_failure_reason,
+                partial_map_retries=self._stats.partial_map_retries,
             )
 
     def clear(self) -> None:
         """Drop every cached scope and reset the counters."""
         with self._lock:
             self._cache.clear()
             self._failed_scopes.clear()
+            self._retried_scopes.clear()
             self._warned.clear()
             self._stats = Stats()
 

tests/test_cloud_batch_signed_url.py

@@ -174,18 +174,28 @@ def test_the_same_uid_retried_still_gets_a_fresh_attempt(self):
         assert len(signer.calls) == 2
 
     def test_a_uid_absent_from_the_batch_map_falls_back(self):
-        """Data drift: the scope was fetched, but this uid is not in it."""
+        """Data drift: the scope was fetched, but this uid is not in it.
+
+        A missing uid triggers ONE re-fetch (see
+        :class:`TestBatchLookupPartialMapRetry` below). When the second
+        answer still does not name the uid -- which is the case here,
+        the fake signer serves the same map every time -- the caller
+        falls back per uid, exactly as it did before the retry existed.
+        """
         from ndi.cloud.batch_signed_url import BatchSignedUrlLookup
 
         files = {"known": "https://s3.example.com/known"}
         signer = _FakeSigner({("ds1", "doc1", ""): files})
-        lookup = BatchSignedUrlLookup(signer=signer)
+        # Skip the sleep for offline test speed.
+        lookup = BatchSignedUrlLookup(signer=signer, partial_map_retry_seconds=0)
 
         assert lookup.lookup("ds1", "doc1", "", "known") == "https://s3.example.com/known"
         assert lookup.lookup("ds1", "doc1", "", "unknown") == ""
         stats = lookup.stats()
         assert stats.uid_hits == 1
         assert stats.uid_misses == 1
+        assert stats.partial_map_retries == 1, "the miss should have triggered one retry"
+        assert stats.signer_calls == 2, "one call for the initial fetch, one for the retry"
 
     def test_failure_reason_names_the_cause(self):
         """Reporting one silent miss leaves the endpoint owner with nothing.
@@ -202,6 +212,147 @@ def test_failure_reason_names_the_cause(self):
         assert "no 'files'" in lookup.stats().last_failure_reason
 
 
+class TestBatchLookupPartialMapRetry:
+    """The batch endpoint sometimes lags behind ``waitForAllBulkUploads``.
+
+    A cluster that has not fully indexed a series' members answers with a
+    populated map that is missing some of them. Without a retry, the caller
+    then falls back per uid for those members and defeats the batch design.
+    NDI-python#309 saw this on TEST_USER_1/prod immediately after upload:
+    the batch answered with 2 of 4 member uids, so 2 members took a per-uid
+    ``getFileDetails`` round trip apiece.
+
+    Fix: on the miss, sleep briefly and re-fetch ONCE. If the second answer
+    covers the uid, take it. If not, warn and fall back. Bounded to one
+    retry per scope so a genuine data-drift miss costs one extra signer
+    call for the whole scope, not one per uid.
+    """
+
+    class _EventualSigner:
+        """Returns a partial map on call 1, the full map on call 2."""
+
+        def __init__(self, scope, partial_files, full_files):
+            self._scope = scope
+            self._partial = partial_files
+            self._full = full_files
+            self.calls = 0
+
+        def __call__(self, dataset_id, document_id, *, file_series="", client=None):
+            self.calls += 1
+            if (dataset_id, document_id, file_series) != self._scope:
+                return False, {"message": "no such scope"}
+            files = self._partial if self.calls == 1 else self._full
+            return True, {"files": dict(files)}
+
+    def _lookup(self, signer, sleeps=None):
+        """A lookup with the sleep short-circuited so tests stay fast.
+
+        The retry sleep exists on the live path to give a cloud endpoint
+        time to catch up; a test does not need to wait. If ``sleeps`` is
+        given, each recorded sleep duration is appended to it so the test
+        can assert that the retry did (or did not) pause.
+        """
+        from ndi.cloud.batch_signed_url import BatchSignedUrlLookup
+
+        recorded = sleeps if sleeps is not None else []
+        return BatchSignedUrlLookup(
+            signer=signer,
+            partial_map_retry_seconds=0.5,
+            sleep=recorded.append,
+        )
+
+    def test_a_lagged_scope_settles_on_the_retry(self):
+        """First map has 2 of 4; the retry names all 4."""
+        scope = ("ds1", "doc1", "chunks")
+        partial = {"u_1": "https://s3/u_1", "u_2": "https://s3/u_2"}
+        full = {
+            "u_1": "https://s3/u_1",
+            "u_2": "https://s3/u_2",
+            "u_3": "https://s3/u_3",
+            "u_4": "https://s3/u_4",
+        }
+        signer = self._EventualSigner(scope, partial, full)
+        lookup = self._lookup(signer)
+
+        # Open all four members in order. The first two are in the partial
+        # map; the third triggers the retry; the fourth uses the refreshed
+        # map.
+        urls = [lookup.lookup("ds1", "doc1", "chunks", f"u_{i}") for i in (1, 2, 3, 4)]
+
+        assert urls == [f"https://s3/u_{i}" for i in (1, 2, 3, 4)]
+        stats = lookup.stats()
+        assert stats.uid_hits == 4, f"every uid should resolve; got {stats!r}"
+        assert stats.uid_misses == 0
+        assert stats.partial_map_retries == 1, "one scope re-fetch, not one per uid"
+        assert stats.signer_calls == 2, "initial fetch + one retry"
+
+    def test_the_retry_pauses_before_re_fetching(self):
+        """The retry exists to give a lagged endpoint time to catch up.
+        Fire the second call too fast and the map is still stale."""
+        scope = ("ds1", "doc1", "chunks")
+        partial = {"u_1": "https://s3/u_1"}
+        full = {"u_1": "https://s3/u_1", "u_2": "https://s3/u_2"}
+        signer = self._EventualSigner(scope, partial, full)
+        sleeps: list[float] = []
+        lookup = self._lookup(signer, sleeps=sleeps)
+
+        lookup.lookup("ds1", "doc1", "chunks", "u_1")
+        lookup.lookup("ds1", "doc1", "chunks", "u_2")
+
+        assert sleeps == [0.5], f"expected one 500 ms pause before the retry; got {sleeps!r}"
+
+    def test_a_scope_is_only_retried_once(self):
+        """A second miss in the same scope must not fire another retry.
+
+        Without this guard a series whose scope is genuinely missing
+        several uids would cost N retries instead of one.
+        """
+        scope = ("ds1", "doc1", "chunks")
+        # Signer that always returns the same partial map -- retry never helps.
+        signer = _FakeSigner({scope: {"u_2": "https://s3/u_2"}})
+        lookup = self._lookup(signer)
+
+        lookup.lookup("ds1", "doc1", "chunks", "u_1")  # miss + retry
+        lookup.lookup("ds1", "doc1", "chunks", "u_3")  # miss, NO retry
+        lookup.lookup("ds1", "doc1", "chunks", "u_4")  # miss, NO retry
+
+        stats = lookup.stats()
+        assert stats.partial_map_retries == 1, "retry must be once per scope"
+        assert stats.signer_calls == 2, "one initial fetch, one retry, then no more"
+        assert stats.uid_misses == 3
+        # The one uid that IS in the map still resolves.
+        assert lookup.lookup("ds1", "doc1", "chunks", "u_2") == "https://s3/u_2"
+
+    def test_a_retry_that_returns_no_map_at_all_still_falls_back(self):
+        """The retry can itself fail (a fetch that raises, a bad payload).
+        The caller must still get an empty string, not a crash."""
+
+        class _FailingRetry:
+            """Partial map first, exception on retry."""
+
+            def __init__(self):
+                self.calls = 0
+
+            def __call__(self, dataset_id, document_id, *, file_series="", client=None):
+                self.calls += 1
+                if self.calls == 1:
+                    return True, {"files": {"u_1": "https://s3/u_1"}}
+                raise RuntimeError("retry blew up")
+
+        signer = _FailingRetry()
+        lookup = self._lookup(signer)
+
+        # u_1 resolves from the initial map; u_2 misses, retry raises.
+        assert lookup.lookup("ds1", "doc1", "chunks", "u_1") == "https://s3/u_1"
+        assert lookup.lookup("ds1", "doc1", "chunks", "u_2") == ""
+        stats = lookup.stats()
+        assert stats.uid_hits == 1
+        # One miss for u_2 (the retry recorded it as the scope-failure miss).
+        assert stats.uid_misses == 1
+        assert stats.partial_map_retries == 1
+        assert "RuntimeError" in stats.last_failure_reason
+
+
 class TestBatchLookupClear:
     def test_clear_drops_everything(self):
         from ndi.cloud.batch_signed_url import BatchSignedUrlLookup

tests/test_cloud_file_series_round_trip.py

@@ -345,12 +345,24 @@ def test_members_survive_a_download_from_the_cloud(
             f"a member fell back to per-uid getFileDetails. "
             f"Stats: signer_calls={stats.signer_calls}, "
             f"uid_hits={stats.uid_hits}, uid_misses={stats.uid_misses}, "
+            f"partial_map_retries={stats.partial_map_retries}, "
             f"failure_reason={stats.last_failure_reason!r}"
         )
-        assert stats.signer_calls == 1, (
-            f"the members of one series should cost ONE presign call: "
-            f"the first fetch populates the scope and the rest resolve "
-            f"from the in-process cache. Got {stats.signer_calls}."
+        # One presign call for the scope, plus at most one retry if the
+        # first answer was a partial map (Waltham-Data-Science/
+        # NDI-python#309 -- some environments lag briefly after the bulk
+        # upload). Anything beyond that would be per-uid fallback, which
+        # the uid_misses assertion above already rules out; this is just
+        # for a clearer failure message.
+        assert stats.signer_calls == 1 + stats.partial_map_retries, (
+            f"the members of one series should cost ONE presign call, "
+            f"plus at most one retry if the first map was partial. "
+            f"Got signer_calls={stats.signer_calls}, "
+            f"partial_map_retries={stats.partial_map_retries}."
+        )
+        assert stats.partial_map_retries <= 1, (
+            f"the retry is bounded to once per scope; "
+            f"got partial_map_retries={stats.partial_map_retries}."
         )
 
     def test_members_survive_a_download_from_the_cloud_with_sync_files_false(