Fix: LWT routing enhancements to utilize LBP #784

nikagra · 2026-01-09T15:30:01Z

Addresses: #384

Summary

LWT statements previously bypassed the configured LoadBalancingPolicy and used raw replicas from the token map, which ignored DC/rack-aware policies. This PR routes LWT through LBP, introduces explicit routing hints, and adjusts replica handling to prefer local-DC replicas and avoid rack-local bias.

Motivation

In clusters using TokenAware with DC/rack-aware policies, LWT traffic was coordinated against replicas in a fixed order, sometimes in a different DC or rack than intended. This led to suboptimal coordination and unexpected cross-DC behavior.

Changes

Public API:
- Add RequestRoutingType enum: REGULAR, LWT.
- Add RequestRoutingMethod enum: REGULAR, PRESERVE_REPLICA_ORDER, TOKEN_BASED_REPLICA_SHUFFLING (hint; only PRESERVE_REPLICA_ORDER is honored in this PR).
- Extend Request:
  - @Nullable RequestRoutingType getRequestType()
  - default RequestRoutingMethod getRoutingMethod() { return RequestRoutingMethod.REGULAR; }
Statements:
- SimpleStatement, BoundStatement, BatchStatement, graph statements, and PrepareRequest now implement getRequestType() (LWT when applicable; otherwise REGULAR).
- Refactor BatchStatementBuilder and DefaultBatchStatement: replace isLWT plumbing with RequestRoutingType. isLWT() now derives from routing type or constituent statements.
Request execution:
- Remove LWT-specific replica lookup from CqlRequestHandler; routing for both LWT and non-LWT is delegated to LBP via DefaultLoadBalancingPolicy.newQueryPlan(...).
Load balancing (DefaultLoadBalancingPolicy):
- Derive replicas from the request when available.
- LWT:
  - Build the plan from replicas, preferring local-DC replicas when a local DC is configured; preserve original replica order.
  - Ignore local rack prioritization to avoid rack-local leader congestion; treat local-DC replicas uniformly.
- Non-LWT: preserve existing behavior (local-DC live nodes; replicas bubbled to the front).
- Honor RequestRoutingMethod.PRESERVE_REPLICA_ORDER by skipping shuffling.
- Factor replica movement/shuffling and slow-replica avoidance into helpers; apply health-based reordering only when avoidSlowReplicas is enabled.
- Accept @Nullable Request and @Nullable Session to align with updated call sites.
Nullability/test adjustments:
- Update NodeStateIT to accept nullable params in newQueryPlan.
- Update examples (KeyRequest) to implement getRequestType().

Tests

LWTLoadBalancingIT: assert coordinators are a subset of replicas (not a single fixed node).
LWTLoadBalancingMultiDcIT: new test verifying LWT routing prefers local-DC replicas in multi-DC setups.
BatchStatementIT: strengthen LWT batch coverage; validate serial consistency, RequestRoutingType, presence of routing key, and rerun behavior.
Minor test rule/constant cleanups.

Behavior

LWT requests:
- Routed through LBP using the computed replica set.
- Prefer local-DC replicas; do not enforce rack-local prioritization.
- Optionally preserve replica order when getRoutingMethod() == PRESERVE_REPLICA_ORDER.
Non-LWT requests:
- Unchanged; existing DC/rack-aware + token-aware behavior remains.

Backward Compatibility

API additions are backward-compatible (new enums and default methods).
Existing client code continues to work; advanced users can opt into routing hints via Request.

Notes for Reviewers

RequestRoutingMethod.TOKEN_BASED_REPLICA_SHUFFLING is introduced as a hint but not actively applied in DefaultLoadBalancingPolicy in this PR; scope limited to PRESERVE_REPLICA_ORDER.
Graph statements and PrepareRequest report REGULAR routing type by default.

core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java

core/src/main/java/com/datastax/oss/driver/internal/core/cql/DefaultBoundStatement.java

core/src/main/java/com/datastax/oss/driver/api/core/RequestRoutingMethod.java

core/src/main/java/com/datastax/oss/driver/api/core/session/Request.java

core/src/main/java/com/datastax/oss/driver/api/core/RequestRoutingMethod.java

core/src/main/java/com/datastax/oss/driver/internal/core/cql/DefaultBatchStatement.java

core/src/main/java/com/datastax/oss/driver/internal/core/cql/DefaultBoundStatement.java

core/src/main/java/com/datastax/oss/driver/internal/core/cql/DefaultSimpleStatement.java

...ain/java/com/datastax/oss/driver/internal/core/loadbalancing/DefaultLoadBalancingPolicy.java

avikivity · 2026-01-22T10:31:42Z

Fix: LWT routing enhancements to utilize LBP

What's LBP?

dkropachev · 2026-01-22T13:38:32Z

Fix: LWT routing enhancements to utilize LBP

What's LBP?

LoadBalancingPolicy

dkropachev

Last ask

...ain/java/com/datastax/oss/driver/internal/core/loadbalancing/DefaultLoadBalancingPolicy.java

core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java

core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java

- Added RequestRoutingType and RequestRoutingMethod enums to define request routing strategies. - Updated DefaultLoadBalancingPolicy to consider request routing type for replica selection, especially for LWT requests. - Updated various graph statement classes (BytecodeGraphStatement, DefaultBatchGraphStatement, DefaultFluentGraphStatement, DefaultScriptGraphStatement) to implement getRequestRoutingType method. - Modified BatchStatementBuilder to set request routing type based on LWT status. - Enhanced DefaultBatchStatement, DefaultBoundStatement, DefaultPrepareRequest, and DefaultSimpleStatement to include routing type in constructors and methods. - Added logic to avoid slow replicas based on health checks and in-flight requests.

core/src/main/java/com/datastax/oss/driver/api/core/cql/StatementBuilder.java

core/src/main/java/com/datastax/oss/driver/internal/core/cql/DefaultBatchStatement.java

core/src/main/java/com/datastax/oss/driver/internal/core/cql/DefaultBoundStatement.java

...ain/java/com/datastax/oss/driver/internal/core/loadbalancing/DefaultLoadBalancingPolicy.java

core/src/main/resources/reference.conf

integration-tests/src/test/java/com/datastax/oss/driver/core/metadata/NodeStateIT.java

...ain/java/com/datastax/oss/driver/internal/core/loadbalancing/DefaultLoadBalancingPolicy.java

Introduce RequestRoutingType to distinguish LWT from regular queries and add a new configuration option to control LWT routing behavior. The new `advanced.load-balancing-policy.default-lwt-request-routing-method` option allows choosing between: - REGULAR: Default shuffling and slow replica avoidance - PRESERVE_REPLICA_ORDER: Maintains replica order from partitioner Changes: - Add RequestRoutingType enum (REGULAR, LWT) to classify requests - Remove unused RequestRoutingMethod enum from Request interface - Thread RequestRoutingType through Statement builders and implementations - Update DefaultLoadBalancingPolicy to route LWT queries according to config - Add corresponding TypedDriverOption and OptionsMap support - Update prepared statement creation to detect and mark LWT queries - Remove RequestRoutingMethod.getRoutingMethod() default method This enables optimized LWT performance by avoiding unnecessary shuffling when replica order preservation is beneficial for linearizability. refactor: update LWT request routing method to preserve replica order

...ain/java/com/datastax/oss/driver/internal/core/loadbalancing/DefaultLoadBalancingPolicy.java

@nonnull

- Fix Javadoc positioning: Move @nonnull annotations after doc comments in DefaultLoadBalancingPolicy methods (per Java conventions) - Add missing @nonnull annotation to StatementBuilderTest mock builder - Add @nullable annotation to NodeStateIT query plan method signature - Standardize test infrastructure: * Add @RunWith(MockitoJUnitRunner.Silent.class) to 7 test classes * Update LoadBalancingPolicyTestBase to stub LWT routing config option * Convert base class from @RunWith to abstract (subclasses now declare runner) - Standardize integration test naming: ccmRule→CCM_RULE, sessionRule→SESSION_RULE - Update test mocks with RequestRoutingType.REGULAR parameter for compatibility - Improve LWT integration tests: * BatchStatementIT: Fix variable references, enhance LWT batch assertions * LWTLoadBalancingIT: Change from single-node to replica-set validation * Add LWTLoadBalancingMultiDcIT: New multi-DC LWT routing test coverage No functional changes to production code—purely code quality and test improvements. Apply suggestions from code review Co-authored-by: Dmitry Kropachev <dmitry.kropachev@gmail.com>

dkropachev · 2026-01-26T16:19:00Z

...ain/java/com/datastax/oss/driver/internal/core/loadbalancing/DefaultLoadBalancingPolicy.java

+    if (request == null) {
+      return RequestRoutingMethod.REGULAR;
+    }
+    switch (request.getRequestRoutingType()) {


switch does not handle null, it will throw NullPointerException

dkropachev · 2026-01-26T16:21:50Z

core/src/main/java/com/datastax/oss/driver/internal/core/cql/DefaultBatchStatement.java

  private final Node node;
  private final int nowInSeconds;
-  private final Boolean isLWT;
+  private final Boolean isExplicitLWTSet;


Can you please store RoutingType instead

New impl: keeping @NonNull and @Nullable for requestRoutingType, additionally moving requestRoutingType field to concrete builder classes implementing StatementBuilder abstract class since it can't fulfill both requrements.

nikagra · 2026-01-26T19:38:49Z

core/src/main/java/com/datastax/oss/driver/api/core/cql/BatchStatementBuilder.java

  @NonNull private ImmutableList.Builder<BatchableStatement<?>> statementsBuilder;
  private int statementsCount;
-  @Nullable private Boolean isLWT = null;
+  @Nullable protected RequestRoutingType requestRoutingType = null;


Note @Nullable with default null value.

nikagra · 2026-01-26T19:38:59Z

core/src/main/java/com/datastax/oss/driver/api/core/cql/BoundStatementBuilder.java

  @NonNull private final ByteBuffer[] values;
  @NonNull private final CodecRegistry codecRegistry;
  @NonNull private final ProtocolVersion protocolVersion;
+  @NonNull protected RequestRoutingType requestRoutingType = RequestRoutingType.REGULAR;


Note @NonNull with default RequestRoutingType.REGULAR value.

nikagra · 2026-01-26T19:40:51Z

core/src/main/java/com/datastax/oss/driver/api/core/cql/StatementBuilder.java

  @Nullable protected Duration timeout;
  @Nullable protected Node node;
  protected int nowInSeconds = Statement.NO_NOW_IN_SECONDS;
-  @NonNull protected RequestRoutingType requestRoutingType = RequestRoutingType.REGULAR;


Removing from base class since @NonNull annotation can not serve every concrete implementation (nullable for batch statements)

dkropachev requested changes Jan 9, 2026

View reviewed changes

core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java Outdated Show resolved Hide resolved

nikagra force-pushed the fix/lwt-routing-lbp-filtering-384 branch 4 times, most recently from c4bdac5 to 5bff3d3 Compare January 15, 2026 20:43

nikagra changed the title ~~Fix/lwt routing lbp filtering 384~~ Fix: LWT routing lbp filtering 384 Jan 19, 2026

nikagra changed the title ~~Fix: LWT routing lbp filtering 384~~ Fix: LWT routing enhancements to utilize LBP Jan 19, 2026

nikagra force-pushed the fix/lwt-routing-lbp-filtering-384 branch 5 times, most recently from f02d437 to 3c7f0db Compare January 20, 2026 18:26

nikagra marked this pull request as ready for review January 20, 2026 18:26

nikagra requested a review from dkropachev January 20, 2026 18:26

nikagra force-pushed the fix/lwt-routing-lbp-filtering-384 branch from 3c7f0db to 5d8ec44 Compare January 20, 2026 18:33

dkropachev requested changes Jan 20, 2026

View reviewed changes

core/src/main/java/com/datastax/oss/driver/internal/core/cql/DefaultBoundStatement.java Outdated Show resolved Hide resolved

nikagra requested a review from dkropachev January 20, 2026 19:00

dkropachev reviewed Jan 21, 2026

View reviewed changes

nikagra force-pushed the fix/lwt-routing-lbp-filtering-384 branch 3 times, most recently from c74510a to 5f6c11d Compare January 22, 2026 08:34

nikagra requested a review from dkropachev January 22, 2026 14:02

nikagra force-pushed the fix/lwt-routing-lbp-filtering-384 branch from 059281b to a4f26e1 Compare January 22, 2026 19:32

dkropachev requested changes Jan 22, 2026

View reviewed changes

dkropachev reviewed Jan 22, 2026

View reviewed changes

core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java Outdated Show resolved Hide resolved

core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java Outdated Show resolved Hide resolved

nikagra force-pushed the fix/lwt-routing-lbp-filtering-384 branch 2 times, most recently from 6a26ea9 to 8456130 Compare January 23, 2026 14:03

nikagra force-pushed the fix/lwt-routing-lbp-filtering-384 branch from 8456130 to b2df37b Compare January 23, 2026 14:07

dkropachev requested changes Jan 23, 2026

View reviewed changes

dkropachev reviewed Jan 23, 2026

View reviewed changes

...ain/java/com/datastax/oss/driver/internal/core/loadbalancing/DefaultLoadBalancingPolicy.java Outdated Show resolved Hide resolved

nikagra force-pushed the fix/lwt-routing-lbp-filtering-384 branch 2 times, most recently from 40f7257 to d227fc4 Compare January 23, 2026 15:13

nikagra force-pushed the fix/lwt-routing-lbp-filtering-384 branch 2 times, most recently from c3889cd to 30f3b4b Compare January 26, 2026 13:58

nikagra requested a review from dkropachev January 26, 2026 15:20

dkropachev requested changes Jan 26, 2026

View reviewed changes

nikagra force-pushed the fix/lwt-routing-lbp-filtering-384 branch from 30f3b4b to 1baa9e8 Compare January 26, 2026 19:27

refactor: Streamline LWT handling in batch statements

a481bc2

nikagra force-pushed the fix/lwt-routing-lbp-filtering-384 branch from 1baa9e8 to a481bc2 Compare January 26, 2026 19:31

nikagra commented Jan 26, 2026

View reviewed changes

nikagra requested a review from dkropachev January 26, 2026 21:12

Fix: LWT routing enhancements to utilize LBP #784

Are you sure you want to change the base?

Fix: LWT routing enhancements to utilize LBP #784

Uh oh!

Conversation

nikagra commented Jan 9, 2026 • edited by dkropachev Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Motivation

Changes

Tests

Behavior

Backward Compatibility

Notes for Reviewers

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

avikivity commented Jan 22, 2026

Uh oh!

dkropachev commented Jan 22, 2026

Uh oh!

dkropachev left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dkropachev Jan 26, 2026

Choose a reason for hiding this comment

Uh oh!

dkropachev Jan 26, 2026

Choose a reason for hiding this comment

Uh oh!

nikagra Jan 26, 2026

Choose a reason for hiding this comment

Uh oh!

nikagra Jan 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nikagra Jan 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nikagra Jan 26, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

nikagra commented Jan 9, 2026 •

edited by dkropachev

Loading

nikagra Jan 26, 2026 •

edited

Loading

nikagra Jan 26, 2026 •

edited

Loading