Skip to content

Commit 35d4a85

Browse files
pdabre12pdabre12
committed
Use a single instance of PlanCheckerRouterPluginPrestoClient for EXPLAIN calls
1 parent 53073bf commit 35d4a85

6 files changed

Lines changed: 66 additions & 46 deletions

File tree

‎presto-native-tests/pom.xml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,12 @@
117117
<scope>test</scope>
118118
</dependency>
119119

120+
<dependency>
121+
<groupId>com.facebook.airlift</groupId>
122+
<artifactId>testing</artifactId>
123+
<scope>test</scope>
124+
</dependency>
125+
120126
<dependency>
121127
<groupId>com.facebook.presto</groupId>
122128
<artifactId>presto-common</artifactId>

‎presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestPlanCheckerRouterPlugin.java‎

Lines changed: 24 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@
5252
import java.util.Optional;
5353

5454
import static com.facebook.airlift.json.JsonCodec.jsonCodec;
55+
import static com.facebook.airlift.testing.Closeables.closeAllRuntimeException;
5556
import static com.facebook.presto.router.scheduler.SchedulerType.CUSTOM_PLUGIN_SCHEDULER;
5657
import static java.lang.Boolean.parseBoolean;
5758
import static java.lang.String.format;
@@ -67,6 +68,7 @@ public class TestPlanCheckerRouterPlugin
6768
private boolean sidecarEnabled;
6869
// mock object only to check the redirect requests counters.
6970
private PlanCheckerRouterPluginPrestoClient planCheckerRouterPluginPrestoClient;
71+
private QueryRunner nativeQueryRunner;
7072

7173
@BeforeClass
7274
public void init()
@@ -77,20 +79,19 @@ public void init()
7779
super.init();
7880
Logging.initialize();
7981

82+
nativeQueryRunner = getQueryRunner();
83+
8084
// for testing purposes, we can skip the router chaining part and specify the native/java clusters directly here
8185
URI nativeClusterURI = ((DistributedQueryRunner) getQueryRunner()).getCoordinator().getBaseUrl();
8286
URI javaClusterURI = ((DistributedQueryRunner) getExpectedQueryRunner()).getCoordinator().getBaseUrl();
8387
PlanCheckerRouterPluginConfig planCheckerRouterConfig = new PlanCheckerRouterPluginConfig()
8488
.setPlanCheckClustersURIs(nativeClusterURI.toString())
8589
.setJavaRouterURI(javaClusterURI)
86-
.setNativeRouterURI(nativeClusterURI);
90+
.setNativeRouterURI(nativeClusterURI)
91+
.setJavaClusterFallbackEnabled(true);
8792

88-
planCheckerRouterPluginPrestoClient = new PlanCheckerRouterPluginPrestoClient(
89-
planCheckerRouterConfig.getPlanCheckClustersURIs().get(0),
90-
planCheckerRouterConfig.getJavaRouterURI(),
91-
planCheckerRouterConfig.getNativeRouterURI(),
92-
planCheckerRouterConfig.getClientRequestTimeout(),
93-
planCheckerRouterConfig.isJavaClusterFallbackEnabled());
93+
planCheckerRouterPluginPrestoClient =
94+
new PlanCheckerRouterPluginPrestoClient(planCheckerRouterConfig);
9495

9596
Path tempFile = Files.createTempFile("temp-config", ".json");
9697
File configFile = getConfigFile(singletonList(planCheckerRouterConfig.getNativeRouterURI()), tempFile.toFile());
@@ -174,6 +175,22 @@ public void testFailingQueriesOnBothClusters(String query, String exceptionMessa
174175
}
175176
}
176177

178+
@Test(dependsOnMethods =
179+
{"testFailingQueriesOnBothClusters",
180+
"testNativeCompatibleQueries",
181+
"testNativeIncompatibleQueries"})
182+
public void testPlanCheckerClusterNotAvailable()
183+
throws SQLException
184+
{
185+
if (sidecarEnabled) {
186+
closeAllRuntimeException(nativeQueryRunner);
187+
nativeQueryRunner = null;
188+
for (String query : getNativeIncompatibleQueries()) {
189+
runQuery(query, httpServerUri);
190+
}
191+
}
192+
}
193+
177194
@DataProvider(name = "failingQueriesOnBothClustersProvider")
178195
public Object[][] getFailingQueriesOnBothClustersProvider()
179196
{

‎presto-plan-checker-router-plugin/src/main/java/com/facebook/presto/router/scheduler/PlanCheckerRouterPluginModule.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ public class PlanCheckerRouterPluginModule
2626
public void configure(Binder binder)
2727
{
2828
configBinder(binder).bindConfig(PlanCheckerRouterPluginConfig.class);
29+
binder.bind(PlanCheckerRouterPluginPrestoClient.class).in(SINGLETON);
2930
binder.bind(PlanCheckerRouterPluginScheduler.class).in(SINGLETON);
3031
}
3132
}

‎presto-plan-checker-router-plugin/src/main/java/com/facebook/presto/router/scheduler/PlanCheckerRouterPluginPrestoClient.java‎

Lines changed: 24 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -25,18 +25,22 @@
2525
import org.weakref.jmx.Managed;
2626
import org.weakref.jmx.Nested;
2727

28+
import javax.inject.Inject;
29+
2830
import java.net.URI;
2931
import java.security.Principal;
3032
import java.util.List;
3133
import java.util.Locale;
3234
import java.util.Map;
3335
import java.util.Optional;
36+
import java.util.concurrent.atomic.AtomicInteger;
3437

3538
import static com.facebook.presto.client.PrestoHeaders.PRESTO_TRANSACTION_ID;
3639
import static com.facebook.presto.client.StatementClientFactory.newStatementClient;
3740
import static com.facebook.presto.router.scheduler.HttpRequestSessionContext.getResourceEstimates;
3841
import static com.facebook.presto.router.scheduler.HttpRequestSessionContext.getSerializedSessionFunctions;
3942
import static com.google.common.base.Verify.verify;
43+
import static java.util.Objects.requireNonNull;
4044

4145
public class PlanCheckerRouterPluginPrestoClient
4246
{
@@ -45,19 +49,25 @@ public class PlanCheckerRouterPluginPrestoClient
4549
private static final CounterStat javaClusterRedirectRequests = new CounterStat();
4650
private static final CounterStat nativeClusterRedirectRequests = new CounterStat();
4751
private final OkHttpClient httpClient = new OkHttpClient();
48-
private final URI planCheckerClusterURI;
52+
private final AtomicInteger planCheckerClusterCandidateIndex = new AtomicInteger(0);
53+
private final List<URI> planCheckerClusterCandidates;
4954
private final URI javaRouterURI;
5055
private final URI nativeRouterURI;
5156
private final Duration clientRequestTimeout;
5257
private final boolean javaClusterFallbackEnabled;
5358

54-
public PlanCheckerRouterPluginPrestoClient(URI planCheckerClusterURI, URI javaRouterURI, URI nativeRouterURI, Duration clientRequestTimeout, boolean javaClusterFallbackEnabled)
59+
@Inject
60+
public PlanCheckerRouterPluginPrestoClient(PlanCheckerRouterPluginConfig planCheckerRouterPluginConfig)
5561
{
56-
this.planCheckerClusterURI = planCheckerClusterURI;
57-
this.javaRouterURI = javaRouterURI;
58-
this.nativeRouterURI = nativeRouterURI;
59-
this.clientRequestTimeout = clientRequestTimeout;
60-
this.javaClusterFallbackEnabled = javaClusterFallbackEnabled;
62+
requireNonNull(planCheckerRouterPluginConfig, "planCheckerRouterPluginConfig is null");
63+
this.planCheckerClusterCandidates =
64+
requireNonNull(planCheckerRouterPluginConfig.getPlanCheckClustersURIs(), "planCheckerClusterCandidates is null");
65+
this.javaRouterURI =
66+
requireNonNull(planCheckerRouterPluginConfig.getJavaRouterURI(), "javaRouterURI is null");
67+
this.nativeRouterURI =
68+
requireNonNull(planCheckerRouterPluginConfig.getNativeRouterURI(), "nativeRouterURI is null");
69+
this.clientRequestTimeout = planCheckerRouterPluginConfig.getClientRequestTimeout();
70+
this.javaClusterFallbackEnabled = planCheckerRouterPluginConfig.isJavaClusterFallbackEnabled();
6171
}
6272

6373
public Optional<URI> getCompatibleClusterURI(Map<String, List<String>> headers, String statement, Principal principal)
@@ -137,7 +147,7 @@ private ClientSession parseHeadersToClientSession(Map<String, List<String>> head
137147
principal);
138148

139149
return new ClientSession(
140-
planCheckerClusterURI,
150+
getPlanCheckerClusterDestination(),
141151
sessionContext.getIdentity().getUser(),
142152
sessionContext.getSource(),
143153
Optional.empty(),
@@ -159,4 +169,10 @@ private ClientSession parseHeadersToClientSession(Map<String, List<String>> head
159169
ImmutableMap.of(), // todo: do we need custom headers?
160170
true);
161171
}
172+
173+
private URI getPlanCheckerClusterDestination()
174+
{
175+
int currentIndex = planCheckerClusterCandidateIndex.getAndUpdate(i -> (i + 1) % planCheckerClusterCandidates.size());
176+
return planCheckerClusterCandidates.get(currentIndex);
177+
}
162178
}

‎presto-plan-checker-router-plugin/src/main/java/com/facebook/presto/router/scheduler/PlanCheckerRouterPluginScheduler.java‎

Lines changed: 5 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -15,59 +15,37 @@
1515

1616
import com.facebook.presto.spi.router.RouterRequestInfo;
1717
import com.facebook.presto.spi.router.Scheduler;
18-
import io.airlift.units.Duration;
1918

2019
import javax.inject.Inject;
2120

2221
import java.net.URI;
2322
import java.util.List;
2423
import java.util.Optional;
25-
import java.util.concurrent.atomic.AtomicInteger;
2624

2725
import static java.util.Objects.requireNonNull;
2826

2927
public class PlanCheckerRouterPluginScheduler
3028
implements Scheduler
3129
{
32-
private final AtomicInteger planCheckerClusterCandidateIndex = new AtomicInteger(0);
33-
3430
private List<URI> candidates;
35-
private final List<URI> planCheckerClusterCandidates;
36-
private final URI javaRouterURI;
37-
private final URI nativeRouterURI;
38-
private final Duration clientRequestTimeout;
39-
private final boolean javaClusterFallbackEnabled;
31+
private final PlanCheckerRouterPluginPrestoClient planCheckerRouterPluginPrestoClient;
4032

4133
@Inject
42-
public PlanCheckerRouterPluginScheduler(PlanCheckerRouterPluginConfig planCheckerRouterConfig)
34+
public PlanCheckerRouterPluginScheduler(PlanCheckerRouterPluginPrestoClient planCheckerRouterPluginPrestoClient)
4335
{
44-
requireNonNull(planCheckerRouterConfig, "PlanCheckerRouterPluginConfig is null");
45-
this.planCheckerClusterCandidates =
46-
requireNonNull(planCheckerRouterConfig.getPlanCheckClustersURIs(), "validatorCandidates is null");
47-
this.javaRouterURI =
48-
requireNonNull(planCheckerRouterConfig.getJavaRouterURI(), "javaRouterURI is null");
49-
this.nativeRouterURI =
50-
requireNonNull(planCheckerRouterConfig.getNativeRouterURI(), "nativeRouterURI is null");
51-
this.clientRequestTimeout = planCheckerRouterConfig.getClientRequestTimeout();
52-
this.javaClusterFallbackEnabled = planCheckerRouterConfig.isJavaClusterFallbackEnabled();
36+
this.planCheckerRouterPluginPrestoClient =
37+
requireNonNull(planCheckerRouterPluginPrestoClient, "planCheckerRouterPluginPrestoClient is null");
5338
}
5439

5540
@Override
5641
public Optional<URI> getDestination(RouterRequestInfo routerRequestInfo)
5742
{
58-
PlanCheckerRouterPluginPrestoClient planCheckerPrestoClient = new PlanCheckerRouterPluginPrestoClient(getValidatorDestination(), javaRouterURI, nativeRouterURI, clientRequestTimeout, javaClusterFallbackEnabled);
59-
return planCheckerPrestoClient.getCompatibleClusterURI(routerRequestInfo.getHeaders(), routerRequestInfo.getQuery(), routerRequestInfo.getPrincipal());
43+
return planCheckerRouterPluginPrestoClient.getCompatibleClusterURI(routerRequestInfo.getHeaders(), routerRequestInfo.getQuery(), routerRequestInfo.getPrincipal());
6044
}
6145

6246
@Override
6347
public void setCandidates(List<URI> candidates)
6448
{
6549
this.candidates = candidates;
6650
}
67-
68-
private URI getValidatorDestination()
69-
{
70-
int currentIndex = planCheckerClusterCandidateIndex.getAndUpdate(i -> (i + 1) % planCheckerClusterCandidates.size());
71-
return planCheckerClusterCandidates.get(currentIndex);
72-
}
7351
}

‎presto-plan-checker-router-plugin/src/test/java/com/facebook/presto/router/scheduler/TestPlanCheckerRouterPlugin.java‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -103,8 +103,9 @@ protected QueryRunner createExpectedQueryRunner()
103103
@Test
104104
public void testPlanCheckerPluginWithNativeCompatibleQueries()
105105
{
106-
Scheduler scheduler = new PlanCheckerRouterPluginScheduler(planCheckerRouterConfig);
107-
scheduler.setCandidates(planCheckerRouterConfig.getPlanCheckClustersURIs());
106+
PlanCheckerRouterPluginPrestoClient planCheckerRouterPluginPrestoClient =
107+
new PlanCheckerRouterPluginPrestoClient(planCheckerRouterConfig);
108+
Scheduler scheduler = new PlanCheckerRouterPluginScheduler(planCheckerRouterPluginPrestoClient);
108109

109110
// native compatible query
110111
Optional<URI> target = scheduler.getDestination(
@@ -122,8 +123,9 @@ public void testPlanCheckerPluginWithNativeCompatibleQueries()
122123
@Test
123124
public void testPlanCheckerPluginWithNativeIncompatibleQueries()
124125
{
125-
Scheduler scheduler = new PlanCheckerRouterPluginScheduler(planCheckerRouterConfig);
126-
scheduler.setCandidates(planCheckerRouterConfig.getPlanCheckClustersURIs());
126+
PlanCheckerRouterPluginPrestoClient planCheckerRouterPluginPrestoClient =
127+
new PlanCheckerRouterPluginPrestoClient(planCheckerRouterConfig);
128+
Scheduler scheduler = new PlanCheckerRouterPluginScheduler(planCheckerRouterPluginPrestoClient);
127129

128130
// native incompatible query
129131
Optional<URI> target = scheduler.getDestination(

0 commit comments

Comments
 (0)