diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyRangeCache.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyRangeCache.java index 98c26381285d..6a3ee0aaeb7a 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyRangeCache.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyRangeCache.java @@ -813,43 +813,49 @@ private TabletSnapshot selectTablet( List skippedTabletDetails, Map resolvedEndpoints, SelectionState selectionStats) { - if (!preferLeader || hintBuilder.getOperationUid() > 0L) { - TabletSnapshot preferredLeader = - preferLeader ? localLeaderForScoreBias(snapshot, hasDirectedReadOptions) : null; - return selectScoreAwareTablet( + TabletSnapshot preferredLeader = + preferLeader ? localLeaderForScoreBias(snapshot, hasDirectedReadOptions) : null; + boolean retryAfterInitialRequest = !skippedTabletUids.isEmpty(); + boolean attemptedLeaderSelection = false; + + if (preferLeader + && !retryAfterInitialRequest + && preferredLeader != null + && preferredLeader.matches(directedReadOptions)) { + attemptedLeaderSelection = true; + selectionStats.sawMatchingReplica = true; + if (!shouldSkip( snapshot, - preferLeader, - directedReadOptions, + preferredLeader, hintBuilder, excludedEndpoints, skippedTabletUids, skippedTabletDetails, resolvedEndpoints, - selectionStats, - preferredLeader); + selectionStats)) { + return preferredLeader; + } } - boolean checkedLeader = false; - if (preferLeader - && !hasDirectedReadOptions - && snapshot.hasLeader() - && snapshot.leader().distance <= MAX_LOCAL_REPLICA_DISTANCE) { - checkedLeader = true; - selectionStats.sawMatchingReplica = true; - if (!shouldSkip( + if (!preferLeader + || retryAfterInitialRequest + || attemptedLeaderSelection + || hintBuilder.getOperationUid() > 0L) { + return selectScoreAwareTablet( snapshot, - snapshot.leader(), + preferLeader, + directedReadOptions, hintBuilder, excludedEndpoints, skippedTabletUids, skippedTabletDetails, resolvedEndpoints, - selectionStats)) { - return snapshot.leader(); - } + selectionStats, + preferredLeader); } + for (int index = 0; index < snapshot.tablets.size(); index++) { - if (checkedLeader && index == snapshot.leaderIndex) { + if (preferredLeader != null && index == snapshot.leaderIndex) { continue; } TabletSnapshot tablet = snapshot.tablets.get(index); diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyRangeCacheTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyRangeCacheTest.java index b7645f044a13..96c27d9c8e34 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyRangeCacheTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyRangeCacheTest.java @@ -693,7 +693,7 @@ public void preferLeaderFalseUsesLowestLatencyReplicaWhenScoresAvailable() { } @Test - public void preferLeaderTrueUsesLatencyScoresWhenOperationUidAvailable() { + public void preferLeaderTrueWithOperationUidStillPrefersReadyLeaderOnInitialRequest() { FakeEndpointCache endpointCache = new FakeEndpointCache(); KeyRangeCache cache = new KeyRangeCache(endpointCache); cache.useDeterministicRandom(); @@ -719,6 +719,70 @@ public void preferLeaderTrueUsesLatencyScoresWhenOperationUidAvailable() { DirectedReadOptions.getDefaultInstance(), hint); + assertNotNull(server); + assertEquals("server1", server.getAddress()); + } + + @Test + public void preferLeaderTrueFallsBackToScoreAwareSelectionWhenLeaderIsExcluded() { + FakeEndpointCache endpointCache = new FakeEndpointCache(); + KeyRangeCache cache = new KeyRangeCache(endpointCache); + cache.useDeterministicRandom(); + cache.addRanges(threeReplicaUpdate()); + + endpointCache.get("server1"); + endpointCache.get("server2"); + endpointCache.get("server3"); + + EndpointLatencyRegistry.recordLatency( + null, TEST_OPERATION_UID, true, "server1", Duration.ofNanos(300_000L)); + EndpointLatencyRegistry.recordLatency( + null, TEST_OPERATION_UID, true, "server2", Duration.ofNanos(200_000L)); + EndpointLatencyRegistry.recordLatency( + null, TEST_OPERATION_UID, true, "server3", Duration.ofNanos(100_000L)); + + RoutingHint.Builder hint = + RoutingHint.newBuilder().setKey(bytes("a")).setOperationUid(TEST_OPERATION_UID); + ChannelEndpoint server = + cache.fillRoutingHint( + true, + KeyRangeCache.RangeMode.COVERING_SPLIT, + DirectedReadOptions.getDefaultInstance(), + hint, + "server1"::equals); + + assertNotNull(server); + assertEquals("server3", server.getAddress()); + } + + @Test + public void preferLeaderTrueUsesScoreAwareSelectionOnRetryPath() { + FakeEndpointCache endpointCache = new FakeEndpointCache(); + KeyRangeCache cache = new KeyRangeCache(endpointCache); + cache.useDeterministicRandom(); + cache.addRanges(threeReplicaUpdate()); + + endpointCache.get("server1"); + endpointCache.get("server2"); + endpointCache.get("server3"); + + EndpointLatencyRegistry.recordLatency( + null, TEST_OPERATION_UID, true, "server1", Duration.ofNanos(300_000L)); + EndpointLatencyRegistry.recordLatency( + null, TEST_OPERATION_UID, true, "server2", Duration.ofNanos(100_000L)); + EndpointLatencyRegistry.recordLatency( + null, TEST_OPERATION_UID, true, "server3", Duration.ofNanos(200_000L)); + + RoutingHint.Builder hint = + RoutingHint.newBuilder().setKey(bytes("a")).setOperationUid(TEST_OPERATION_UID); + hint.addSkippedTabletUidBuilder().setTabletUid(1).setIncarnation(bytes("1")); + ChannelEndpoint server = + cache.fillRoutingHint( + true, + KeyRangeCache.RangeMode.COVERING_SPLIT, + DirectedReadOptions.getDefaultInstance(), + hint); + assertNotNull(server); assertEquals("server2", server.getAddress()); }