Skip to content

Commit 0125be7

Browse files
CopilotCopilot
committed
Scope metadata hedging to cold-start cache reads + skeptic-review fixes
Self-review (Seon thorough) fixes: - Gate hedging on an explicit metadataCachePopulation request flag set only by the container-properties refresh, cold partition-key read, and routing-map fetch callsites, so public container reads and hybrid-search PK reads are no longer hedged. - Size the hedge thread pool to max(cpu_count, 2*budget) so each in-flight hedge always has its primary+hedge threads. - Remove the dead record_failure no-op (metadata cache health is account-global, not per-partition). - Clarify the threshold rationale and the async budget non-blocking check. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent d681614 commit 0125be7

12 files changed

Lines changed: 101 additions & 54 deletions

‎sdk/cosmos/azure-cosmos/azure/cosmos/_availability_strategy_config.py‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,11 @@
2828
DEFAULT_THRESHOLD_STEPS_MS = 100
2929

3030
# Defaults for cold-start metadata cache hedging. These are SDK-derived and not
31-
# customer-configurable. The threshold mirrors the .NET design (first-attempt
32-
# control-plane timeout ~1s + a 500ms step => 1500ms). The concurrency budget
33-
# caps the number of in-flight metadata hedges per client.
31+
# customer-configurable. The threshold is an aggressive tail-latency trigger: metadata
32+
# (control-plane) reads use a short read timeout (DBAReadTimeout, 3s by default), so the
33+
# 1.5s hedge fires well before the primary's own timeout, giving a slow region a chance
34+
# to be raced by a second region while the primary attempt is still outstanding. The
35+
# concurrency budget caps the number of in-flight metadata hedges per client.
3436
DEFAULT_METADATA_HEDGING_THRESHOLD_MS = 1500
3537
DEFAULT_METADATA_HEDGING_THRESHOLD_STEPS_MS = 500
3638
DEFAULT_METADATA_HEDGING_CONCURRENCY_BUDGET = 8

‎sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@
7272
from ._cosmos_responses import CosmosDict, CosmosList, CosmosItemPaged
7373
from ._range_partition_resolver import RangePartitionResolver
7474
from ._read_items_helper import ReadItemsHelperSync
75-
from ._request_object import RequestObject
75+
from ._request_object import RequestObject, METADATA_CACHE_POPULATION_OPTION
7676
from ._retry_utility import ConnectionRetryPolicy
7777
from ._routing import routing_map_provider
7878
from ._query_advisor import get_query_advice_info
@@ -2983,6 +2983,7 @@ def Read(
29832983
headers,
29842984
options.get("partitionKey", None))
29852985
request_params.set_excluded_location_from_options(options)
2986+
request_params.set_metadata_cache_population_from_options(options)
29862987
request_params.set_availability_strategy(options, self.availability_strategy)
29872988
request_params.availability_strategy_executor = self.availability_strategy_executor
29882989
base.set_session_token_header(self, headers, path, request_params, options)
@@ -3318,6 +3319,7 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]:
33183319
options.get("partitionKey", None)
33193320
)
33203321
request_params.set_excluded_location_from_options(options)
3322+
request_params.set_metadata_cache_population_from_options(options)
33213323
request_params.set_availability_strategy(options, self.availability_strategy)
33223324
request_params.availability_strategy_executor = self.availability_strategy_executor
33233325

@@ -3956,7 +3958,8 @@ def refresh_routing_map_provider(
39563958

39573959
def _refresh_container_properties_cache(self, container_link: str):
39583960
# If container properties cache is stale, refresh it by reading the container.
3959-
container = self.ReadContainer(container_link, options=None)
3961+
container = self.ReadContainer(
3962+
container_link, options={METADATA_CACHE_POPULATION_OPTION: True})
39603963
# Only cache Container Properties that will not change in the lifetime of the container
39613964
self._set_container_properties_cache(container_link, _build_properties_cache(container, container_link))
39623965

@@ -4003,7 +4006,9 @@ def _get_partition_key_definition(
40034006
partition_key_definition = cached_container.get("partitionKey")
40044007
# Else read the collection from backend and add it to the cache
40054008
else:
4006-
container = self.ReadContainer(collection_link, options)
4009+
read_options = {**options} if options else {}
4010+
read_options[METADATA_CACHE_POPULATION_OPTION] = True
4011+
container = self.ReadContainer(collection_link, read_options)
40074012
partition_key_definition = container.get("partitionKey")
40084013
self._set_container_properties_cache(collection_link, _build_properties_cache(container, collection_link))
40094014
return partition_key_definition

‎sdk/cosmos/azure-cosmos/azure/cosmos/_metadata_hedging.py‎

Lines changed: 14 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,6 @@
3535
import time
3636
from concurrent.futures import CancelledError, Future, ThreadPoolExecutor, as_completed
3737
from threading import Event, Semaphore
38-
from types import SimpleNamespace
3938
from typing import Any, Callable, Dict, List, Optional, Tuple
4039

4140
from azure.core.exceptions import ServiceRequestError, ServiceResponseError # pylint: disable=no-legacy-azure-core-http-response-import
@@ -142,9 +141,13 @@ class MetadataCrossRegionHedgingHandler(AvailabilityStrategyHandlerMixin):
142141
"""
143142

144143
def __init__(self, concurrency_budget: int = DEFAULT_METADATA_HEDGING_CONCURRENCY_BUDGET) -> None:
145-
self._budget = Semaphore(max(1, concurrency_budget))
146-
# Long-lived executor shared across this client's metadata hedges.
147-
self._executor = ThreadPoolExecutor(max_workers=os.cpu_count()) # pylint: disable=consider-using-with
144+
budget = max(1, concurrency_budget)
145+
self._budget = Semaphore(budget)
146+
# Long-lived executor shared across this client's metadata hedges. Every in-flight
147+
# hedge needs two worker threads (primary + hedge), so size the pool to at least
148+
# 2 * budget; otherwise hedge tasks can starve behind primaries on low-core hosts.
149+
max_workers = max(os.cpu_count() or 1, 2 * budget)
150+
self._executor = ThreadPoolExecutor(max_workers=max_workers) # pylint: disable=consider-using-with
148151

149152
def is_acceptable_winner( # pylint: disable=unused-argument
150153
self,
@@ -192,7 +195,6 @@ def _send_with_delay(
192195
location_index: int,
193196
available_locations: List[str],
194197
complete_status: Event,
195-
first_request_params_holder: SimpleNamespace,
196198
) -> ResponseType:
197199
strategy = request_params.availability_strategy
198200
if strategy is None:
@@ -212,8 +214,6 @@ def _send_with_delay(
212214
)
213215

214216
req = copy.deepcopy(request)
215-
if location_index == 0:
216-
first_request_params_holder.request_params = params
217217

218218
if complete_status.is_set():
219219
raise CancelledError("The request has been cancelled")
@@ -252,14 +252,13 @@ def execute_request(
252252
return execute_request_fn(request_params, request)
253253

254254
completion_status = Event()
255-
first_request_params_holder: SimpleNamespace = SimpleNamespace(request_params=None)
256255
try:
257256
primary_future = self._executor.submit(
258257
self._send_with_delay, request_params, request, execute_request_fn,
259-
0, available_locations, completion_status, first_request_params_holder)
258+
0, available_locations, completion_status)
260259
hedge_future = self._executor.submit(
261260
self._send_with_delay, request_params, request, execute_request_fn,
262-
1, available_locations, completion_status, first_request_params_holder)
261+
1, available_locations, completion_status)
263262
futures: List[Future] = [primary_future, hedge_future]
264263

265264
for completed_future in as_completed(futures):
@@ -269,10 +268,11 @@ def execute_request(
269268

270269
if self.is_acceptable_winner(result, exception, is_hedge=not is_primary):
271270
completion_status.set()
272-
if not is_primary:
273-
# Hedge won; record a failure for the primary region so health
274-
# tracking can react to the slow/failed primary.
275-
self._record_cancel_for_first_request(first_request_params_holder, global_endpoint_manager)
271+
# Note: no per-region failure is recorded for the primary when the
272+
# hedge wins. Metadata cache reads are account-global, not
273+
# per-partition, so the circuit-breaker health tracker (keyed on
274+
# partition-key ranges) does not apply here; a slow primary is not
275+
# necessarily an unhealthy one.
276276
if exception is not None:
277277
raise exception
278278
return result # type: ignore[return-value]
@@ -288,14 +288,6 @@ def execute_request(
288288
completion_status.set()
289289
self._budget.release()
290290

291-
def _record_cancel_for_first_request(
292-
self,
293-
request_params_holder: SimpleNamespace,
294-
global_endpoint_manager: Any,
295-
) -> None:
296-
if request_params_holder.request_params is not None:
297-
global_endpoint_manager.record_failure(request_params_holder.request_params)
298-
299291

300292
def execute_metadata_hedging(
301293
handler: MetadataCrossRegionHedgingHandler,

‎sdk/cosmos/azure-cosmos/azure/cosmos/_request_object.py‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,11 @@
3131
from .documents import _OperationType
3232
from .http_constants import ResourceType
3333

34+
# Internal options key used to mark a request as a cold-start metadata-cache
35+
# population read (container / partition-key-range cache). Only set by the
36+
# cache-population callsites; gates cold-start metadata hedging.
37+
METADATA_CACHE_POPULATION_OPTION = "metadataCachePopulation"
38+
3439

3540
class RequestObject(object): # pylint: disable=too-many-instance-attributes
3641
def __init__(
@@ -59,6 +64,7 @@ def __init__(
5964
self.pk_val = pk_val
6065
self.retry_write: int = 0
6166
self.is_hedging_request: bool = False # Flag to track if this is a hedged request
67+
self.is_metadata_cache_population: bool = False # Cold-start metadata-cache read
6268
self.completion_status: Optional[Union[threading.Event, asyncio.Event]] = None
6369

6470
def route_to_location_with_preferred_location_flag( # pylint: disable=name-too-long
@@ -97,6 +103,17 @@ def set_excluded_location_from_options(self, options: Mapping[str, Any]) -> None
97103
if self._can_set_excluded_location(options):
98104
self.excluded_locations = options['excludedLocations']
99105

106+
def set_metadata_cache_population_from_options(self, options: Mapping[str, Any]) -> None: # pylint: disable=name-too-long
107+
"""Mark this request as a cold-start metadata-cache population read when the
108+
internal ``metadataCachePopulation`` options flag is set.
109+
110+
:param options: The request options that may contain the internal flag.
111+
:type options: Mapping[str, Any]
112+
:return: None
113+
"""
114+
if options and options.get(METADATA_CACHE_POPULATION_OPTION):
115+
self.is_metadata_cache_population = True
116+
100117
def set_retry_write(self, request_options: Mapping[str, Any], client_retry_write: int) -> None:
101118
if self.resource_type == ResourceType.Document:
102119
if request_options and request_options.get(Constants.Kwargs.RETRY_WRITE):
@@ -149,7 +166,7 @@ def set_availability_strategy(
149166

150167
def should_cancel_request(self) -> bool:
151168
"""Check if this request should be cancelled due to parallel request completion.
152-
169+
153170
:return: True if request should be cancelled, False otherwise
154171
:rtype: bool
155172
"""

‎sdk/cosmos/azure-cosmos/azure/cosmos/_routing/aio/routing_map_provider.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
from typing import Dict, Any, Optional, List, TYPE_CHECKING
2929
from azure.core.utils import CaseInsensitiveDict
3030
from ... import _base, http_constants
31+
from ..._request_object import METADATA_CACHE_POPULATION_OPTION
3132
from ..collection_routing_map import CollectionRoutingMap
3233
from ...exceptions import CosmosHttpResponseError
3334
from .._routing_map_provider_common import (
@@ -406,6 +407,10 @@ async def _fetch_routing_map(
406407
change_feed_options = prepare_fetch_options_and_headers(
407408
current_previous_map, feed_options, base_kwargs_for_headers
408409
)
410+
# Mark these reads as cold-start metadata-cache population so the request
411+
# layer can hedge them cross-region; other PartitionKeyRange readers
412+
# (e.g. hybrid search) are not flagged and so are not hedged.
413+
change_feed_options[METADATA_CACHE_POPULATION_OPTION] = True
409414
base_headers: Dict[str, Any] = base_kwargs_for_headers['headers']
410415

411416
while True:

‎sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
from typing import Dict, Any, Optional, List, TYPE_CHECKING
2929
from azure.core.utils import CaseInsensitiveDict
3030
from .. import _base, http_constants
31+
from .._request_object import METADATA_CACHE_POPULATION_OPTION
3132
from .collection_routing_map import CollectionRoutingMap
3233
from ..exceptions import CosmosHttpResponseError
3334
from ._routing_map_provider_common import (
@@ -373,6 +374,10 @@ def _fetch_routing_map(
373374
change_feed_options = prepare_fetch_options_and_headers(
374375
current_previous_map, feed_options, base_kwargs_for_headers
375376
)
377+
# Mark these reads as cold-start metadata-cache population so the request
378+
# layer can hedge them cross-region; other PartitionKeyRange readers
379+
# (e.g. hybrid search) are not flagged and so are not hedged.
380+
change_feed_options[METADATA_CACHE_POPULATION_OPTION] = True
376381
base_headers: Dict[str, Any] = base_kwargs_for_headers['headers']
377382

378383
while True:

‎sdk/cosmos/azure-cosmos/azure/cosmos/_synchronized_request.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -264,6 +264,8 @@ def _is_metadata_hedging_applicable(client, request_params: RequestObject, globa
264264
return False
265265
if request_params.is_hedging_request or request_params.availability_strategy is not None:
266266
return False
267+
if not request_params.is_metadata_cache_population:
268+
return False
267269
if not is_supported_metadata_request(request_params):
268270
return False
269271
opt_in = getattr(client, "_metadata_hedging_opt_in", None)

‎sdk/cosmos/azure-cosmos/azure/cosmos/aio/_asynchronous_request.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -235,6 +235,8 @@ def _is_metadata_hedging_applicable(client, request_params: RequestObject, globa
235235
return False
236236
if request_params.is_hedging_request or request_params.availability_strategy is not None:
237237
return False
238+
if not request_params.is_metadata_cache_population:
239+
return False
238240
if not is_supported_metadata_request(request_params):
239241
return False
240242
opt_in = getattr(client, "_metadata_hedging_opt_in", None)

‎sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1312,6 +1312,7 @@ async def Read(
13121312
headers,
13131313
options.get("partitionKey", None))
13141314
request_params.set_excluded_location_from_options(options)
1315+
request_params.set_metadata_cache_population_from_options(options)
13151316
await base.set_session_token_header_async(self, headers, path, request_params, options)
13161317
request_params.set_availability_strategy(options, self.availability_strategy)
13171318
request_params.availability_strategy_max_concurrency = self.availability_strategy_max_concurrency
@@ -3116,6 +3117,7 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]:
31163117
options.get("partitionKey", None),
31173118
)
31183119
request_params.set_excluded_location_from_options(options)
3120+
request_params.set_metadata_cache_population_from_options(options)
31193121
request_params.set_availability_strategy(options, self.availability_strategy)
31203122
request_params.availability_strategy_max_concurrency = self.availability_strategy_max_concurrency
31213123
headers = base.GetHeaders(self, initial_headers, "get", path, id_, resource_type,
@@ -3819,7 +3821,8 @@ async def refresh_routing_map_provider(
38193821

38203822
async def _refresh_container_properties_cache(self, container_link: str):
38213823
# If container properties cache is stale, refresh it by reading the container.
3822-
container = await self.ReadContainer(container_link, options=None)
3824+
container = await self.ReadContainer(
3825+
container_link, options={_request_object.METADATA_CACHE_POPULATION_OPTION: True})
38233826
# Only cache Container Properties that will not change in the lifetime of the container
38243827
self._set_container_properties_cache(container_link, _build_properties_cache(container, container_link))
38253828

@@ -3916,7 +3919,9 @@ async def _get_partition_key_definition(
39163919
partition_key_definition = cached_container.get("partitionKey")
39173920
# Else read the collection from backend and add it to the cache
39183921
else:
3919-
container = await self.ReadContainer(collection_link, options)
3922+
read_options = {**options} if options else {}
3923+
read_options[_request_object.METADATA_CACHE_POPULATION_OPTION] = True
3924+
container = await self.ReadContainer(collection_link, read_options)
39203925
partition_key_definition = container.get("partitionKey")
39213926
self._set_container_properties_cache(collection_link, _build_properties_cache(container, collection_link))
39223927
return partition_key_definition

‎sdk/cosmos/azure-cosmos/azure/cosmos/aio/_metadata_hedging.py‎

Lines changed: 10 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,6 @@
3030
import copy
3131
import logging
3232
from asyncio import CancelledError, Event, Task
33-
from types import SimpleNamespace
3433
from typing import Any, Awaitable, Callable, Dict, List, Optional, Tuple
3534

3635
from azure.core.pipeline.transport import HttpRequest # pylint: disable=no-legacy-azure-core-http-response-import
@@ -106,7 +105,6 @@ async def _send_with_delay(
106105
location_index: int,
107106
available_locations: List[str],
108107
complete_status: Event,
109-
first_request_params_holder: SimpleNamespace,
110108
) -> ResponseType:
111109
strategy = request_params.availability_strategy
112110
if strategy is None:
@@ -126,8 +124,6 @@ async def _send_with_delay(
126124
)
127125

128126
req = copy.deepcopy(request)
129-
if location_index == 0:
130-
first_request_params_holder.request_params = params
131127

132128
if complete_status.is_set():
133129
raise CancelledError("The request has been cancelled")
@@ -160,21 +156,24 @@ async def execute_request(
160156
logger.debug("Metadata hedge skipped: SingleRegion")
161157
return await execute_request_fn(request_params, request)
162158

159+
# Non-blocking budget check: locked() is True only when the budget is fully
160+
# exhausted (value == 0). There is no await between this check and acquire(),
161+
# so in single-threaded asyncio no other coroutine can consume the slot in
162+
# between and acquire() completes immediately without blocking.
163163
if self._budget.locked():
164164
logger.debug("Metadata hedge skipped: BudgetExhausted")
165165
return await execute_request_fn(request_params, request)
166166

167167
await self._budget.acquire()
168168
completion_status = Event()
169-
first_request_params_holder: SimpleNamespace = SimpleNamespace(request_params=None)
170169
active_tasks: List[Task] = []
171170
try:
172171
primary_task = asyncio.create_task(self._send_with_delay(
173172
request_params, request, execute_request_fn,
174-
0, available_locations, completion_status, first_request_params_holder))
173+
0, available_locations, completion_status))
175174
hedge_task = asyncio.create_task(self._send_with_delay(
176175
request_params, request, execute_request_fn,
177-
1, available_locations, completion_status, first_request_params_holder))
176+
1, available_locations, completion_status))
178177
active_tasks = [primary_task, hedge_task]
179178

180179
pending = set(active_tasks)
@@ -187,9 +186,10 @@ async def execute_request(
187186

188187
if self.is_acceptable_winner(result, exception, is_hedge=not is_primary):
189188
completion_status.set()
190-
if not is_primary:
191-
await self._record_cancel_for_first_request(
192-
first_request_params_holder, global_endpoint_manager)
189+
# Note: no per-region failure is recorded for the primary when the
190+
# hedge wins. Metadata cache reads are account-global, not
191+
# per-partition, so the circuit-breaker health tracker (keyed on
192+
# partition-key ranges) does not apply here.
193193
if exception is not None:
194194
raise exception
195195
return result # type: ignore[return-value]
@@ -208,14 +208,6 @@ async def execute_request(
208208
await asyncio.gather(*active_tasks, return_exceptions=True)
209209
self._budget.release()
210210

211-
async def _record_cancel_for_first_request(
212-
self,
213-
request_params_holder: SimpleNamespace,
214-
global_endpoint_manager: Any,
215-
) -> None:
216-
if request_params_holder.request_params is not None:
217-
await global_endpoint_manager.record_failure(request_params_holder.request_params)
218-
219211

220212
async def execute_metadata_hedging(
221213
handler: MetadataCrossRegionAsyncHedgingHandler,

0 commit comments

Comments
 (0)