Skip to content

Commit 0031453

Browse files
jeet1995Copilot
andauthored
[GatewayV2] Fix over-spanning issue when partial partition keys are provided. (#49688)
* Fix thin-client MULTI_HASH prefix partition key EPK over-span For a partial (prefix) hierarchical (MULTI_HASH) partition key, the thin-client (Gateway V2) store model previously sent only StartEpkHash/EndEpkHash routing headers and no doc-level EPK filter. The proxy therefore resolved the request to the owning physical partition and returned every co-located document (an over-span) instead of only those matching the prefix. ThinClientStoreModel.wrapInHttpRequest now detects a prefix MULTI_HASH key and sets READ_FEED_KEY_TYPE=EffectivePartitionKeyRange, START_EPK and END_EPK on the request headers before RntbdRequest.from(), so RntbdRequestHeaders serializes the prefix EPK sub-range [hash(prefix), hash(prefix) + "FF") as the backend doc-level filter (hex string as UTF-8 bytes), mirroring the Direct/RNTBD FeedRangeEpkImpl path. StartEpkHash/EndEpkHash routing headers are retained. QueryPlan requests and full (non-prefix) keys are unchanged. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Add thin-client query E2E parity regression tests Ports ThinClientQueryE2ETest and its ThinClientTestBase helper. The suite self-baselines every query against a Direct (RNTBD) client and the thin-client (:10250 HTTP/2) path, asserting identical results. This directly guards the MULTI_HASH prefix partition-key EPK over-span fix: prefix (single-component) hierarchical-partition-key queries must return only the narrow matching set, matching Direct, rather than over-spanning to all co-located documents. Both files are runtime self-enabling (enableThinClientForTest() + builder-driven HTTP/2 via setEnabled(true)); no module property or TestSuiteBase changes needed. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Reuse prefix EPK range on thin-client parallel prefix-query path The earlier fix (499f928) only scoped the prefix MULTI_HASH read when the request carried a non-null partition key. The parallel prefix-query path intentionally nulls the partial hierarchical partition key and instead carries the narrow prefix effective-partition-key range on the request's feedRange, so that guard never fired there and the thin-client read still over-spanned to the whole physical partition. Mark such requests (prefixPartitionKeyQuery) in ParallelDocumentQueryExecutionContextBase and, in ThinClientStoreModel, reuse the already-computed FeedRangeEpkImpl range as the doc-level EPK filter (START_EPK/END_EPK + StartEpkHash/EndEpkHash) instead of recomputing getEPKRangeForPrefixPartitionKey from a partition key we no longer have. This mirrors the Direct/GatewayV1 decision points and honors DRY/SRP. Also relax assertThinClientEndpointUsed so an all-QueryPlan diagnostics context passes: in thin-client mode QueryPlan calls resolve via the classic gateway and are the only requests permitted on a non-thin-client endpoint; any non-QueryPlan (data) request on a non-thin-client endpoint still fails the assertion. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Consolidate thin-client prefix EPK RNTBD headers, propagate prefix flag on clone, add readAllItems/readMany prefix HPK tests ThinClientStoreModel.wrapInHttpRequest now sets all five EPK-specific RNTBD headers (ReadFeedKeyType, StartEpk/EndEpk, StartEpkHash/EndEpkHash) in a single block directly on the RNTBD request, replacing the split request-header-map + post-construction approach. Wire encoding is unchanged (hex-string UTF-8 bytes for StartEpk/EndEpk, decoded hash bytes for the hash headers). RxDocumentServiceRequest.clone() now copies prefixPartitionKeyQuery so hedged/availability-strategy/retry clones keep the prefix signal and do not over-span. ThinClientQueryE2ETest adds prefix-HPK regression coverage for readAllItems (ReadFeed) and readManyByPartitionKeys, asserting Direct-vs-thin-client parity and per-tenant scoping (no over-span). * Simplify prefix-query flag assignment in ParallelDocumentQueryExecutionContextBase Assign isPrefixPartitionKeyQuery directly from isPartialPartitionKeyQuery(...) and guard only the PARTITION_KEY-setting branch with !isPrefixPartitionKeyQuery, removing the if/else. Behavior is unchanged: a partial (prefix) key still leaves partitionKeyInternal null so the prefix FeedRangeEpkImpl survives createDocumentServiceRequestWithFeedRange. * Add CHANGELOG entry for thin-client prefix hierarchical partition key over-span fix * Minimize prefix-query flag change vs main in ParallelDocumentQueryExecutionContextBase Keep main's single 'if (!isPartialPartitionKeyQuery)' structure; only capture the predicate in isPrefixPartitionKeyQuery and pass it to request.setPrefixPartitionKeyQuery. Removes the if/else. * Order thin-client prefix detection to prefer the cached feed-range path In ThinClientStoreModel.wrapInHttpRequest, check the cached prefix feed-range path (isPrefixPartitionKeyQuery + FeedRangeEpkImpl) first and fall back to recomputing from the partition key only for the PK-bearing prefix case. Behavior is unchanged (the two cases are mutually exclusive). * Address Copilot review: validate all requests in thin-client endpoint assertion TestSuiteBase.assertThinClientEndpointUsed now validates every request instead of early-returning on the first thin-client match, so a mixed scenario (some data requests via the classic gateway) fails; QueryPlan requests remain exempt and null endpoints are handled. ThinClientTestBase.assertThinClientEndpointUsed delegates to the shared TestSuiteBase implementation, removing duplicated logic that could NPE on a null endpoint and reject valid QueryPlan-via-gateway scenarios. * Simplify thin-client prefix EPK handling to reuse HTTP EPK headers Replace the explicit MULTI_HASH prefix-detection and pre-serialization range computation with a header-presence branch: when the request already carries START_EPK/END_EPK HTTP headers (populated upstream by FeedRangeEpkImpl for prefix hierarchical partition keys and explicit EPK-range reads), set only the additive StartEpkHash/EndEpkHash RNTBD tokens. The ReadFeedKeyType/StartEpk/ EndEpk string tokens are copied verbatim from those HTTP headers by RntbdRequestHeaders#addStartAndEndKeys during RntbdRequest.from(), so they no longer need to be set here. Removes now-unused imports. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Remove dead prefix-query flag and expand HPK thin-client test coverage Following the simplification of the prefix EPK fix to reuse the existing StartEpk/EndEpk HTTP headers (a56db9a), the RxDocumentServiceRequest prefixPartitionKeyQuery flag is no longer read anywhere. Remove the field, its accessors, the clone copy, and the caller assignment in ParallelDocumentQueryExecutionContextBase. Test changes: - Retrofit CosmosMultiHashTest into the thinclient TestNG group so the MULTI_HASH prefix over-span coverage also runs against GatewayV2 (proxy :10250) over HTTP/2, not just the emulator gateway. - ThinClientQueryE2ETest: correct the prefix-scope Javadoc to describe the actual header-copy behavior and add a deterministic prefix-partition-key readAllItems regression. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Revert ParallelDocumentQueryExecutionContextBase to main The prefix-query flag on RxDocumentServiceRequest was removed in 03a7c67 after the fix was simplified to reuse the StartEpk/EndEpk HTTP headers. The only functional edit in this file (the setter call) went away with it, leaving behind unnecessary refactor noise. Restore the file to its upstream/main state so the hotfix touches no query execution code. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Address final review comments on thin-client EPK over-span fix Comment 1: Tighten the StartEpkHash/EndEpkHash guard in ThinClientStoreModel#wrapInHttpRequest to require READ_FEED_KEY_TYPE == EffectivePartitionKeyRange in addition to the StartEpk/EndEpk headers. RntbdRequestHeaders#addStartAndEndKeys only emits the ReadFeedKeyType RNTBD token when that header is present, so the consumer guard now matches the producer contract and keeps the string and binary EPK tokens coherent. Comment 2: Add testExplicitFeedRangeShardedReadParity, which exercises the explicit user-supplied FeedRange (EPK-range) read path directly by sharding the multi-physical-partition container into its physical feed ranges and asserting per-shard parity with the Direct baseline plus a clean, non-overlapping union over the full fan-out. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Add mode-aware thin-client endpoint routing assertions to CosmosMultiHashTest Verify data requests route through the thin-client endpoint (:10250) when the thinclient group is active, and through the gateway otherwise. Also removes redundant in-code thin-client enablement (owned by the TestNG config / Maven profile). Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Address xinlian12 review comments: docstring accuracy, dead scaffolding cleanup, container-leak fix - ThinClientQueryE2ETest: clarify prefix-HPK test docstring as routing+SQL-predicate parity guard; cross-reference genuine EPK-header guards - ThinClientTestBase: convert to stateless final helper holder (drop unused lifecycle scaffolding) - ThinClientQueryE2ETest: best-effort container cleanup in setup catch block to avoid leaking provisioned containers Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
1 parent 0d3040f commit 0031453

6 files changed

Lines changed: 2405 additions & 16 deletions

File tree

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosMultiHashTest.java

Lines changed: 63 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66

77
package com.azure.cosmos;
88

9+
import com.azure.cosmos.implementation.TestConfigurations;
910
import com.azure.cosmos.models.CosmosContainerProperties;
1011
import com.azure.cosmos.models.CosmosItemIdentity;
1112
import com.azure.cosmos.models.CosmosItemRequestOptions;
@@ -42,6 +43,8 @@
4243
import java.util.concurrent.atomic.AtomicReference;
4344
import java.util.stream.Collectors;
4445

46+
import static com.azure.cosmos.rx.ThinClientTestBase.assertGatewayEndpointUsed;
47+
import static com.azure.cosmos.rx.ThinClientTestBase.assertThinClientEndpointUsed;
4548
import static org.assertj.core.api.Assertions.assertThat;
4649

4750
public class CosmosMultiHashTest extends TestSuiteBase {
@@ -53,6 +56,9 @@ public class CosmosMultiHashTest extends TestSuiteBase {
5356
private CosmosDatabase createdDatabase;
5457
private CosmosContainer createdMultiHashContainer;
5558
private CosmosContainer createdNestedPathContainer;
59+
// Reflects the active @BeforeClass lifecycle: true under the thinclient group, false under emulator.
60+
// Drives mode-aware endpoint-routing assertions and the continuation-token client selection below.
61+
private boolean thinClientEnabled;
5662

5763
@Factory(dataProvider = "clientBuilders")
5864
public CosmosMultiHashTest(CosmosClientBuilder clientBuilder) {
@@ -61,7 +67,25 @@ public CosmosMultiHashTest(CosmosClientBuilder clientBuilder) {
6167

6268
@BeforeClass(groups = {"emulator"}, timeOut = SETUP_TIMEOUT)
6369
public void before_CosmosMultiHashTest() {
70+
thinClientEnabled = false;
6471
client = getClientBuilder().buildClient();
72+
initDatabaseAndContainers();
73+
}
74+
75+
// Enrolls the MULTI_HASH prefix over-span coverage into the thin-client (GatewayV2, proxy :10250) group.
76+
// The class-level @Factory yields a plain gateway builder without HTTP/2, which cannot route to the thin
77+
// client, so this lifecycle builds an explicit HTTP/2 gateway client (mirrors clientBuildersWithGatewayAndHttp2)
78+
// so the same test bodies exercise the RNTBD prefix-EPK header path. COSMOS.THINCLIENT_ENABLED is supplied by
79+
// the thin-client CI lane (see sdk/cosmos/tests.yml); IDE/standalone runs must pass -DCOSMOS.THINCLIENT_ENABLED=true.
80+
@BeforeClass(groups = {"thinclient"}, timeOut = SETUP_TIMEOUT)
81+
public void before_CosmosMultiHashTest_thinClient() {
82+
thinClientEnabled = true;
83+
client = createGatewayRxDocumentClient(
84+
TestConfigurations.HOST, null, true, null, true, true, true).buildClient();
85+
initDatabaseAndContainers();
86+
}
87+
88+
private void initDatabaseAndContainers() {
6589
createdDatabase = createSyncDatabase(client, preExistingDatabaseId);
6690
String collectionName = UUID.randomUUID().toString();
6791

@@ -103,7 +127,24 @@ public void afterClass() {
103127
safeCloseSyncClient(client);
104128
}
105129

106-
@Test(groups = {"emulator"}, timeOut = TIMEOUT)
130+
@AfterClass(groups = {"thinclient"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true)
131+
public void afterClass_thinClient() {
132+
logger.info("starting cleanup (thin client)....");
133+
safeDeleteSyncDatabase(createdDatabase);
134+
safeCloseSyncClient(client);
135+
}
136+
137+
// Mode-aware endpoint routing assertion. Under the thinclient group every non-QueryPlan (data) request
138+
// must route through the thin-client proxy endpoint (:10250); under the emulator group no request may.
139+
private void assertEndpointRouting(CosmosDiagnostics diagnostics) {
140+
if (thinClientEnabled) {
141+
assertThinClientEndpointUsed(diagnostics);
142+
} else {
143+
assertGatewayEndpointUsed(diagnostics);
144+
}
145+
}
146+
147+
@Test(groups = {"emulator", "thinclient"}, timeOut = TIMEOUT)
107148
public void itemCRUD() {
108149
CityItem cityItem = new CityItem(UUID.randomUUID().toString(), "Redmond", "98052", 1);
109150

@@ -120,6 +161,7 @@ public void itemCRUD() {
120161
cityItem.getId(), partitionKey, CityItem.class);
121162

122163
assertThat(readResponse.getItem().toString()).isEqualTo(cityItem.toString());
164+
assertEndpointRouting(readResponse.getDiagnostics());
123165
createdMultiHashContainer.deleteItem(cityItem.getId(), partitionKey, new CosmosItemRequestOptions());
124166
}
125167

@@ -170,7 +212,7 @@ private void validateResponse(FeedResponse<ObjectNode> response,
170212
.collect(Collectors.toList())
171213
);
172214
}
173-
@Test(groups = { "emulator" }, timeOut = TIMEOUT)
215+
@Test(groups = { "emulator", "thinclient" }, timeOut = TIMEOUT)
174216
public void readManySupportsNestedPartitionKeyPaths() {
175217
String city = "nested-readmany-" + UUID.randomUUID();
176218

@@ -186,9 +228,10 @@ public void readManySupportsNestedPartitionKeyPaths() {
186228

187229
FeedResponse<ObjectNode> documentFeedResponse = createdNestedPathContainer.readMany(itemList, ObjectNode.class);
188230
validateResponse(documentFeedResponse, itemList);
231+
assertEndpointRouting(documentFeedResponse.getCosmosDiagnostics());
189232
}
190233

191-
@Test(groups = { "emulator" }, timeOut = TIMEOUT)
234+
@Test(groups = { "emulator", "thinclient" }, timeOut = TIMEOUT)
192235
public void readAllItemsSupportsNestedPartitionKeyPaths() {
193236
String city = "nested-readall-" + UUID.randomUUID();
194237

@@ -203,7 +246,13 @@ public void readAllItemsSupportsNestedPartitionKeyPaths() {
203246
CosmosPagedIterable<ObjectNode> readAllResults =
204247
createdNestedPathContainer.readAllItems(new PartitionKey(city), ObjectNode.class);
205248

206-
assertThat(readAllResults.stream().map(item -> item.get("id").asText()).collect(Collectors.toList()))
249+
List<String> readAllIds = new ArrayList<>();
250+
for (FeedResponse<ObjectNode> page : readAllResults.iterableByPage()) {
251+
page.getResults().forEach(item -> readAllIds.add(item.get("id").asText()));
252+
assertEndpointRouting(page.getCosmosDiagnostics());
253+
}
254+
255+
assertThat(readAllIds)
207256
.containsExactlyInAnyOrder(firstItem.get("id").asText(), secondItem.get("id").asText());
208257
}
209258

@@ -445,7 +494,7 @@ private void validateDocCRUDAndQuery() throws InterruptedException {
445494
deleteAllItems();
446495
}
447496

448-
@Test(groups = { "emulator" }, timeOut = TIMEOUT)
497+
@Test(groups = { "emulator", "thinclient" }, timeOut = TIMEOUT)
449498
private void multiHashQueryTests() {
450499
ArrayList<CityItem> docs = createItems();
451500

@@ -530,6 +579,7 @@ private void multiHashQueryTests() {
530579
FeedResponse<ObjectNode> feedResponse = feedResponses.iterator().next();
531580
assertThat(feedResponse.getResults().size()).isEqualTo(2);
532581
assertThat(feedResponse.getResults().get(0).get("zipcode").asInt() < feedResponse.getResults().get(1).get("zipcode").asInt()).isTrue();
582+
assertEndpointRouting(feedResponse.getCosmosDiagnostics());
533583

534584
//Using continuation token
535585
testPartialPKContinuationToken();
@@ -645,7 +695,13 @@ private void deleteAllItems() {
645695
private void testPartialPKContinuationToken() {
646696
String requestContinuation = null;
647697
List<ObjectNode> receivedDocuments = new ArrayList<>();
648-
CosmosAsyncClient asyncClient = getClientBuilder().buildAsyncClient();
698+
// Under the thinclient group build an explicit HTTP/2 gateway client so this prefix continuation-token
699+
// pass also routes to the thin client (:10250) instead of the plain gateway; the emulator group uses the
700+
// factory-provided builder. thinClientEnabled mirrors the active @BeforeClass lifecycle.
701+
CosmosAsyncClient asyncClient =
702+
thinClientEnabled
703+
? createGatewayRxDocumentClient(TestConfigurations.HOST, null, true, null, true, true, true).buildAsyncClient()
704+
: getClientBuilder().buildAsyncClient();
649705
CosmosAsyncDatabase cosmosAsyncDatabase = new CosmosAsyncDatabase(createdDatabase.getId(), asyncClient);
650706
CosmosAsyncContainer cosmosAsyncContainer = new CosmosAsyncContainer(createdMultiHashContainer.getId(), cosmosAsyncDatabase);
651707
String query = "SELECT * FROM c ORDER BY c.zipcode ASC";
@@ -671,6 +727,7 @@ private void testPartialPKContinuationToken() {
671727
requestContinuation = firstPage.getContinuationToken();
672728
receivedDocuments.addAll(firstPage.getResults());
673729
assertThat(firstPage.getResults().size()).isEqualTo(1);
730+
assertEndpointRouting(firstPage.getCosmosDiagnostics());
674731
} while (requestContinuation != null);
675732
assertThat(receivedDocuments.size()).isEqualTo(3);
676733
asyncClient.close();

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java

Lines changed: 17 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2710,16 +2710,26 @@ protected static void assertThinClientEndpointUsed(CosmosDiagnosticsContext ctx)
27102710
assertThat(requests).isNotNull();
27112711
assertThat(requests.size()).isPositive();
27122712

2713+
// Validate every request rather than early-returning on the first thin-client match: a mixed
2714+
// scenario (some data requests via the thin-client endpoint, some via the classic gateway) must
2715+
// fail. Every non-QueryPlan (data) request must route through the thin-client proxy endpoint;
2716+
// QueryPlan calls are resolved via the classic gateway in thin-client mode, so they are the only
2717+
// requests allowed to target a non-thin-client endpoint.
27132718
for (CosmosDiagnosticsRequestInfo requestInfo : requests) {
2714-
if (requestInfo.getEndpoint() != null
2715-
&& requestInfo.getEndpoint().contains(THIN_CLIENT_ENDPOINT_INDICATOR)) {
2716-
return;
2719+
// requestType has the form "<ResourceType>:<OperationType>" (OperationType.QueryPlan
2720+
// stringifies to "QueryPlan").
2721+
String requestType = requestInfo.getRequestType();
2722+
if (requestType != null && requestType.endsWith(":QueryPlan")) {
2723+
continue;
27172724
}
2718-
}
27192725

2720-
assertThat(false)
2721-
.as("No request targeting thin client proxy endpoint (" + THIN_CLIENT_ENDPOINT_INDICATOR + ")")
2722-
.isTrue();
2726+
String endpoint = requestInfo.getEndpoint();
2727+
assertThat(endpoint != null && endpoint.contains(THIN_CLIENT_ENDPOINT_INDICATOR))
2728+
.as("Non-QueryPlan request must target the thin client proxy endpoint ("
2729+
+ THIN_CLIENT_ENDPOINT_INDICATOR + "), but was: " + endpoint
2730+
+ " (requestType: " + requestType + ")")
2731+
.isTrue();
2732+
}
27232733
}
27242734

27252735
protected static void safeClose(AsyncDocumentClient client) {

0 commit comments

Comments
 (0)