Skip to content

Commit 26fb4c6

Browse files
authored
chore(spanner): add ReplicaSelector interface and impl (#12728)
Adds a ReplicaSelector interface and a default implementation for a ReplicaSelector that uses the power-of-2-random-choices selection strategy. Other strategies can be implemented in the future.
1 parent 4029de2 commit 26fb4c6

File tree

3 files changed

+227
-0
lines changed

3 files changed

+227
-0
lines changed
Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.spanner.spi.v1;
18+
19+
import com.google.api.core.BetaApi;
20+
import com.google.api.core.InternalApi;
21+
import com.google.common.base.MoreObjects;
22+
import java.util.List;
23+
import java.util.Random;
24+
import java.util.concurrent.ThreadLocalRandom;
25+
import java.util.function.Function;
26+
27+
/** Implementation of {@link ReplicaSelector} using the "Power of 2 Random Choices" strategy. */
28+
@InternalApi
29+
@BetaApi
30+
public class PowerOfTwoReplicaSelector implements ReplicaSelector {
31+
32+
@Override
33+
public ChannelEndpoint select(
34+
List<ChannelEndpoint> candidates, Function<ChannelEndpoint, Double> scoreLookup) {
35+
if (candidates == null || candidates.isEmpty()) {
36+
return null;
37+
}
38+
if (candidates.size() == 1) {
39+
return candidates.get(0);
40+
}
41+
42+
Random random = ThreadLocalRandom.current();
43+
int index1 = random.nextInt(candidates.size());
44+
int index2 = random.nextInt(candidates.size() - 1);
45+
if (index2 >= index1) {
46+
index2++;
47+
}
48+
49+
ChannelEndpoint c1 = candidates.get(index1);
50+
ChannelEndpoint c2 = candidates.get(index2);
51+
52+
Double score1 = scoreLookup.apply(c1);
53+
Double score2 = scoreLookup.apply(c2);
54+
55+
// Handle null scores by treating them as Double.MAX_VALUE (lowest priority)
56+
double s1 = MoreObjects.firstNonNull(score1, Double.MAX_VALUE);
57+
double s2 = MoreObjects.firstNonNull(score2, Double.MAX_VALUE);
58+
59+
return s1 <= s2 ? c1 : c2;
60+
}
61+
}
Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.spanner.spi.v1;
18+
19+
import com.google.api.core.BetaApi;
20+
import com.google.api.core.InternalApi;
21+
import java.util.List;
22+
import java.util.function.Function;
23+
24+
/** Interface for selecting a replica from a list of candidates. */
25+
@InternalApi
26+
@BetaApi
27+
public interface ReplicaSelector {
28+
29+
/**
30+
* Selects a replica from the given list of candidates.
31+
*
32+
* @param candidates the list of eligible candidates.
33+
* @param scoreLookup a function to look up the latency score for a candidate.
34+
* @return the selected candidate, or null if the list is empty.
35+
*/
36+
ChannelEndpoint select(
37+
List<ChannelEndpoint> candidates, Function<ChannelEndpoint, Double> scoreLookup);
38+
}
Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.spanner.spi.v1;
18+
19+
import static org.junit.Assert.assertEquals;
20+
import static org.junit.Assert.assertNull;
21+
import static org.junit.Assert.assertTrue;
22+
23+
import java.util.Arrays;
24+
import java.util.HashMap;
25+
import java.util.List;
26+
import java.util.Map;
27+
import org.junit.Test;
28+
import org.junit.runner.RunWith;
29+
import org.junit.runners.JUnit4;
30+
31+
@RunWith(JUnit4.class)
32+
public class PowerOfTwoReplicaSelectorTest {
33+
34+
private static class TestEndpoint implements ChannelEndpoint {
35+
private final String address;
36+
37+
TestEndpoint(String address) {
38+
this.address = address;
39+
}
40+
41+
@Override
42+
public String getAddress() {
43+
return address;
44+
}
45+
46+
@Override
47+
public boolean isHealthy() {
48+
return true;
49+
}
50+
51+
@Override
52+
public boolean isTransientFailure() {
53+
return false;
54+
}
55+
56+
@Override
57+
public io.grpc.ManagedChannel getChannel() {
58+
return null;
59+
}
60+
}
61+
62+
@Test
63+
public void testEmptyList() {
64+
PowerOfTwoReplicaSelector selector = new PowerOfTwoReplicaSelector();
65+
assertNull(selector.select(null, endpoint -> 1.0));
66+
assertNull(selector.select(Arrays.asList(), endpoint -> 1.0));
67+
}
68+
69+
@Test
70+
public void testSingleElement() {
71+
PowerOfTwoReplicaSelector selector = new PowerOfTwoReplicaSelector();
72+
ChannelEndpoint endpoint = new TestEndpoint("a");
73+
assertEquals(endpoint, selector.select(Arrays.asList(endpoint), e -> 1.0));
74+
}
75+
76+
@Test
77+
public void testTwoElementsPicksBetter() {
78+
PowerOfTwoReplicaSelector selector = new PowerOfTwoReplicaSelector();
79+
ChannelEndpoint better = new TestEndpoint("better");
80+
ChannelEndpoint worse = new TestEndpoint("worse");
81+
82+
Map<ChannelEndpoint, Double> scores = new HashMap<>();
83+
scores.put(better, 10.0);
84+
scores.put(worse, 20.0);
85+
86+
List<ChannelEndpoint> candidates = Arrays.asList(better, worse);
87+
88+
for (int i = 0; i < 100; i++) {
89+
assertEquals(better, selector.select(candidates, scores::get));
90+
}
91+
}
92+
93+
@Test
94+
public void testThreeElementsNeverPicksWorst() {
95+
PowerOfTwoReplicaSelector selector = new PowerOfTwoReplicaSelector();
96+
ChannelEndpoint best = new TestEndpoint("best");
97+
ChannelEndpoint middle = new TestEndpoint("middle");
98+
ChannelEndpoint worst = new TestEndpoint("worst");
99+
100+
Map<ChannelEndpoint, Double> scores = new HashMap<>();
101+
scores.put(best, 10.0);
102+
scores.put(middle, 20.0);
103+
scores.put(worst, 30.0);
104+
105+
List<ChannelEndpoint> candidates = Arrays.asList(best, middle, worst);
106+
107+
for (int i = 0; i < 100; i++) {
108+
ChannelEndpoint selected = selector.select(candidates, scores::get);
109+
assertTrue("Should not pick worst", selected != worst);
110+
}
111+
}
112+
113+
@Test
114+
public void testNullScoresTreatedAsMax() {
115+
PowerOfTwoReplicaSelector selector = new PowerOfTwoReplicaSelector();
116+
ChannelEndpoint withScore = new TestEndpoint("withScore");
117+
ChannelEndpoint withoutScore = new TestEndpoint("withoutScore");
118+
119+
Map<ChannelEndpoint, Double> scores = new HashMap<>();
120+
scores.put(withScore, 100.0);
121+
122+
List<ChannelEndpoint> candidates = Arrays.asList(withScore, withoutScore);
123+
124+
for (int i = 0; i < 100; i++) {
125+
assertEquals(withScore, selector.select(candidates, scores::get));
126+
}
127+
}
128+
}

0 commit comments

Comments
 (0)