Leverage x-ms-cosmos-hub-region-processing-only for 404 Read Session Not Available cross-region retry scenarios.

[jeet1995]

Motivation

The pull request serves as the first iteration which integrates x-ms-cosmos-hub-region-processing-only header. Setting the value of this header to true will allow a Cosmos DB backend node to return a 403:3 in case the backend node belongs to a non-hub physical partition.

Using this setup, the CosmosClient instance can determine partition-set level hub which in the first iteration helps in region detection of 404 Read Session Not Available cross-region detection for Single-Writer accounts. This is needed in particular when failover happens in a rolling-manner partition-set by partition-set and in Per-Partition Automatic Failover cases where hub is a partition-set granular notion. Simply relying on LocationCache to provide account-level hub region is incorrect.

Scope

In this pull request, the focus is on how 404 Read Session Not Available cross-region retry handling happens for Single-Writer accounts.

Critical Changes

• The approach taken here is to pin the x-ms-cosmos-hub-region-processing-only once a request hits a 404 Read Session Not Available. This ensures an operation (a construct which encapsulates several I/O calls) is sticky to the hub region.

• IMPORTANT: An extra retry cycle is executed on the hub region for 404 Read Session Not Available prior to pinning the x-ms-cosmos-hub-region-processing-only header. This is done to avoid 403 Read/Write Forbidden retry loop for Single-Writer Non-PPAF enabled accounts until LocationCache reports the updated hub region.

Testing

Per-Partition Automatic Failover

The approach was to set a naming configuration (simulateRevokeLocalWriteStatusOfPartition) consumed by the service fabric process mapped to the original hub region (say North Central US) for a particular physical partition.

Post that, a 404 Read Session Not Available is injected into the same partition for which the write privilege was revoked (North Central US).

Using a "pure-read" workload, the goal is to assert whether the read (a readItem operation) gets a 200 status code from the partition-set specific hub region.

As "reads" can get a 403 Write Forbidden status code, these "reads" can update partition-set level hub which future reads and writes can use.

Single-Writer accounts with no PPAF enabled

Pending item: The expected test setup is to execute a write region change on an account with a physical partition-set count in the order of ~2000 (typical in our DR drills) and to subject the account to a "read-only" workload and see how hub-region stickiness holds up.

Key Changes

1. ClientRetryPolicy.java

Changes to shouldRetry() method:

• Added: Reset of per-partition automatic failover override for reads at the start of each retry evaluation
• Added: On 403/WriteForbidden, sets per-partition automatic failover override for reads before attempting failover

Changes to shouldRetryOnSessionNotAvailable() method:

Aspect	Azure/main (Before)	jeet1995:AzCosmos_AddHubRegionProcessingOnlyHeader (After)
Max retry count (single-master)	> 1 (2 retries max)	> 2 (3 retries max)
PPAF override for reads	Sets flag once	Extracted to helper method setPerPartitionAutomaticFailoverOverrideForReads()
Hub Region Header	Not present	On 2nd retry (sessionTokenRetryCount == 2), sets shouldAddHubRegionProcessingOnlyHeader = true

New method setPerPartitionAutomaticFailoverOverrideForReads():

• Encapsulates logic for setting PPAF override for read requests
• Only applies to read-only, non-metadata requests
• Sets resolvedPartitionKeyRangeForPerPartitionAutomaticFailover from resolvedPartitionKeyRange

Changes to onBeforeSendRequest() method:

• Added: Checks CrossRegionAvailabilityContext.shouldAddHubRegionProcessingOnlyHeader()
• Added: Sets HUB_REGION_PROCESSING_ONLY HTTP header to "true" when flag is set

---

2. GlobalPartitionEndpointManagerForPerPartitionAutomaticFailover.java

Renamed class:

• PartitionLevelFailoverInfo → PartitionLevelAutomaticFailoverInfo

Method signature changes:

Method	Azure/main	jeet1995:AzCosmos_AddHubRegionProcessingOnlyHeader
tryMarkEndpointAsUnavailableForPartitionKeyRange	(request, isEndToEndTimeoutHit)	Overloaded: (request, isEndToEndTimeoutHit) + (request, isEndToEndTimeoutHit, forceFailoverThroughReads)
isPerPartitionAutomaticFailoverApplicable	(request)	Overloaded: (request) + (request, forceFailoverThroughReads)

Changes to isPerPartitionAutomaticFailoverApplicable():

Aspect	Azure/main	After
Read request handling	Returns false for all read requests	Checks forceFailoverThroughReads flag; if true, allows PPAF for reads
PPAF override check	N/A	Validates shouldUsePerPartitionAutomaticFailoverOverrideForReadsIfApplicable() unless forceFailoverThroughReads is set

---

3. CrossRegionAvailabilityContextForRxDocumentServiceRequest.java

New field:

• shouldAddHubRegionProcessingOnlyHeader (AtomicBoolean)

New methods:

• shouldAddHubRegionProcessingOnlyHeader() - getter
• setShouldAddHubRegionProcessingOnlyHeader(boolean) - setter

---

Flow Diagrams

READ_SESSION_NOT_AVAILABLE (404/1002) Retry Flow - Before (Azure/main)

flowchart TD
    A[Read Request Fails with 404/1002] --> B{Endpoint Discovery Enabled?}
    B -->|No| C[No Retry]
    B -->|Yes| D{canUseMultipleWriteLocations?}
    
    D -->|Yes| E{sessionTokenRetryCount >= endpoints.size?}
    E -->|Yes| C
    E -->|No| F[Retry on next preferred location]
    
    D -->|No| G{sessionTokenRetryCount > 1?}
    G -->|Yes| C
    G -->|No| H{PPAF Enabled?}
    
    H -->|Yes| I[Set PPAF override for reads]
    H -->|No| J[Standard retry]
    
    I --> K[Retry with Duration.ZERO]
    J --> K
    
    style C fill:#f66
    style K fill:#6f6

Loading

READ_SESSION_NOT_AVAILABLE (404/1002) Retry Flow - After (jeet1995:AzCosmos_AddHubRegionProcessingOnlyHeader)

flowchart TD
    A[Read Request Fails with 404/1002] --> B{Endpoint Discovery Enabled?}
    B -->|No| C[No Retry]
    B -->|Yes| D{canUseMultipleWriteLocations?}
    
    D -->|Yes| E{sessionTokenRetryCount >= endpoints.size?}
    E -->|Yes| C
    E -->|No| F[Retry on next preferred location]
    
    D -->|No| G{sessionTokenRetryCount > 2?}
    G -->|Yes| C
    G -->|No| H[Set retryContext]
    
    H --> I{PPAF Enabled?}
    I -->|Yes| J[Set PPAF override for reads via helper method]
    I -->|No| K[Continue]
    
    J --> K
    K --> L{sessionTokenRetryCount == 2?}
    
    L -->|Yes| M[Set shouldAddHubRegionProcessingOnlyHeader = true]
    L -->|No| N[Retry with Duration.ZERO]
    M --> N
    
    style C fill:#f66
    style N fill:#6f6
    style M fill:#ff9

Loading

Hub Region Header Application in onBeforeSendRequest()

flowchart TD
    A[onBeforeSendRequest called] --> B[Set request metadata]
    B --> C[Clear previous routing directive]
    C --> D{retryContext != null?}
    
    D -->|Yes| E[Route based on retry context]
    D -->|No| F[Continue]
    
    E --> F
    F --> G{CrossRegionAvailabilityContext exists?}
    
    G -->|Yes| H{shouldAddHubRegionProcessingOnlyHeader?}
    G -->|No| I[Resolve endpoint]
    
    H -->|Yes| J[Add HUB_REGION_PROCESSING_ONLY header = true]
    H -->|No| I
    
    J --> I
    I --> K[Route to resolved endpoint]
    K --> L[Apply PPAF location override if applicable]
    
    style J fill:#ff9

Loading

PPAF Applicability Check for Reads - After

flowchart TD
    A[isPerPartitionAutomaticFailoverApplicable] --> B{PPAF Enabled?}
    B -->|No| C[Return false]
    B -->|Yes| D{Request is read-only?}
    
    D -->|Yes| E{forceFailoverThroughReads?}
    D -->|No| F[Continue validation]
    
    E -->|Yes| F
    E -->|No| G{PPAF override flag set?}
    
    G -->|Yes| F
    G -->|No| C
    
    F --> H{Multiple regions available?}
    H -->|No| C
    H -->|Yes| I{Valid resource/operation type?}
    
    I -->|No| C
    I -->|Yes| J{Single-master account?}
    
    J -->|Yes| K[Return true]
    J -->|No| C
    
    style C fill:#f66
    style K fill:#6f6

Loading

---

Behavior Summary

Scenario	Retry 1	Retry 2	Retry 3
Single-master, 404/1002 on read	Retry locally	Retry to write region	Retry to write region with HUB_REGION_PROCESSING_ONLY header
Multi-master, 404/1002 on read	Cycle through preferred locations	Continue cycling	Continue until exhausted

---

Testing Considerations

1. Verify HUB_REGION_PROCESSING_ONLY header is added only on the 3rd retry attempt (sessionTokenRetryCount == 2)
2. Verify header is only added for single-master accounts
3. Verify PPAF override for reads works correctly with the new helper method
4. Verify 403/WriteForbidden triggers PPAF failover through reads

All SDK Contribution checklist:

• The pull request does not introduce [breaking changes]
• CHANGELOG is updated for new features, bug fixes or other significant changes.
• I have read the contribution guidelines.

General Guidelines and Best Practices

• Title of the pull request is clear and informative.
• There are a small number of commits, each of which have an informative message. This means that previously merged commits do not appear in the history of the PR. For more information on cleaning up the commits in your PR, see this page.

Testing Guidelines

• Pull request includes test coverage for the included changes.

[FabianMeiswinkel]

LGTM except for one question whether one change is intended - and if so why?

[jeet1995]

/azp run java - cosmos - tests

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[jeet1995]

/azp run java - cosmos - tests

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[jeet1995]

/azp run java - cosmos - tests

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[jeet1995]

/azp run java - cosmos - tests

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[jeet1995]

/azp run java - cosmos - tests

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[jeet1995]

/azp run java - cosmos - tests

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[xinlian12]

LGTM, thanks

[jeet1995]

/azp run java - cosmos - tests

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[jeet1995]

/azp run java - cosmos - tests

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[jeet1995]

/azp run java - cosmos - tests

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[xinlian12]

LGTM, thanks

[jeet1995]

/check-enforcer override