Skip to content

Commit 1994b03

Browse files
olim7tolim7t
committed
Prevent race between cancellation and query completion (JAVA-614).
1 parent 77e47fd commit 1994b03

6 files changed

Lines changed: 169 additions & 12 deletions

File tree

‎driver-core/CHANGELOG.rst‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@ CHANGELOG
44
2.0.9.1:
55
--------
66

7+
- [bug] Prevent race between cancellation and query completion (JAVA-614)
8+
79

810
2.0.9:
911
------

‎driver-core/pom.xml‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,20 @@
8989
<scope>test</scope>
9090
</dependency>
9191

92+
<dependency>
93+
<groupId>org.scassandra</groupId>
94+
<artifactId>java-client</artifactId>
95+
<!-- N.B. later versions of scassandra require JDK 7 -->
96+
<version>0.4.1</version>
97+
<scope>test</scope>
98+
<exclusions>
99+
<exclusion>
100+
<groupId>ch.qos.logback</groupId>
101+
<artifactId>logback-classic</artifactId>
102+
</exclusion>
103+
</exclusions>
104+
</dependency>
105+
92106
</dependencies>
93107

94108
<build>

‎driver-core/src/main/java/com/datastax/driver/core/RequestHandler.java‎

Lines changed: 44 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ class RequestHandler implements Connection.ResponseCallback {
5656
private volatile List<Host> triedHosts;
5757
private volatile HostConnectionPool currentPool;
5858
private final AtomicReference<QueryState> queryStateRef;
59+
private volatile boolean shouldCancelConnectionHandler;
5960

6061
// This represents the number of times a retry has been triggered by the RetryPolicy (this is different from
6162
// queryStateRef.get().retryCount, because some retries don't involve the policy, for example after an
@@ -67,7 +68,6 @@ class RequestHandler implements Connection.ResponseCallback {
6768

6869
private volatile Map<InetSocketAddress, Throwable> errors;
6970

70-
private volatile boolean isCanceled;
7171
private volatile Connection.ResponseHandler connectionHandler;
7272

7373
private final Timer.Context timerContext;
@@ -99,7 +99,7 @@ private Metrics metrics() {
9999

100100
public void sendRequest() {
101101
try {
102-
while (queryPlan.hasNext() && !isCanceled) {
102+
while (queryPlan.hasNext() && queryStateRef.get() != QueryState.CANCELLED) {
103103
Host host = queryPlan.next();
104104
logger.trace("Querying node {}", host);
105105
if (query(host))
@@ -126,7 +126,7 @@ private boolean query(Host host) {
126126
triedHosts.add(current);
127127
}
128128
current = host;
129-
connectionHandler = connection.write(this);
129+
write(connection, this);
130130
return true;
131131
} catch (ConnectionException e) {
132132
// If we have any problem with the connection, move to the next node.
@@ -154,6 +154,33 @@ private boolean query(Host host) {
154154
}
155155
}
156156

157+
private void write(Connection connection, Connection.ResponseCallback responseCallback) throws ConnectionException, BusyConnectionException {
158+
// Make sure cancel() does not see a stale connectionHandler if it sees the new query state
159+
// before connection.write has completed
160+
connectionHandler = null;
161+
162+
// Set query state to "in progress" at next iteration
163+
while (true) {
164+
QueryState previous = queryStateRef.get();
165+
assert !previous.inProgress;
166+
if (previous == QueryState.CANCELLED) {
167+
return;
168+
}
169+
if (queryStateRef.compareAndSet(previous, previous.startNext()))
170+
break;
171+
}
172+
173+
connectionHandler = connection.write(responseCallback);
174+
175+
// N.B: onSet/onException could have already been called at this point
176+
177+
// If cancel was called after we set the state to "in progress", but before connection.write, it might have seen a null connectionHandler.
178+
// In this case we need to cancel it now to release the connection (because the next onSet/onException/onTimeout will see the CANCELLED
179+
// state and return immediately).
180+
if (shouldCancelConnectionHandler)
181+
connectionHandler.cancelHandler();
182+
}
183+
157184
private void logError(InetSocketAddress address, Throwable exception) {
158185
logger.debug("Error querying {}, trying next host (error is: {})", address, exception.toString());
159186
if (errors == null)
@@ -162,15 +189,15 @@ private void logError(InetSocketAddress address, Throwable exception) {
162189
}
163190

164191
private void retry(final boolean retryCurrent, ConsistencyLevel newConsistencyLevel) {
165-
queryStateRef.set(queryStateRef.get().startNext());
166-
167192
final Host h = current;
168193
this.retryConsistencyLevel = newConsistencyLevel;
169194

170195
// We should not retry on the current thread as this will be an IO thread.
171196
manager.executor().execute(new Runnable() {
172197
@Override
173198
public void run() {
199+
if (queryStateRef.get() == QueryState.CANCELLED)
200+
return;
174201
try {
175202
if (retryCurrent) {
176203
if (query(h))
@@ -185,9 +212,13 @@ public void run() {
185212
}
186213

187214
public void cancel() {
188-
isCanceled = true;
189-
if (connectionHandler != null)
190-
connectionHandler.cancelHandler();
215+
QueryState previous = queryStateRef.getAndSet(QueryState.CANCELLED);
216+
if (previous.inProgress) {
217+
if (connectionHandler != null)
218+
connectionHandler.cancelHandler();
219+
else
220+
shouldCancelConnectionHandler = true;
221+
}
191222
}
192223

193224
@Override
@@ -388,9 +419,8 @@ public void onSet(Connection connection, Message.Response response, long latency
388419
connection.setKeyspace(prepareKeyspace);
389420
}
390421

391-
queryStateRef.set(queryStateRef.get().startNext());
392422
try {
393-
connection.write(prepareAndRetry(toPrepare.getQueryString()));
423+
write(connection, prepareAndRetry(toPrepare.getQueryString()));
394424
} finally {
395425
// Always reset the previous keyspace if needed
396426
if (connection.keyspace() == null || !connection.keyspace().equals(currentKeyspace))
@@ -583,10 +613,12 @@ interface Callback extends Connection.ResponseCallback {
583613
}
584614

585615
// This is used to prevent races between request completion (either success or error) and timeout.
586-
// A retry is in progress once we have written the request to the connection and until we get back a response or a timeout.
616+
// A retry is in progress once we have written the request to the connection and until we get back a response (see onSet
617+
// or onException) or a timeout (see onTimeout).
587618
// The count increments on each retry.
588619
static class QueryState {
589-
static QueryState INITIAL = new QueryState(0, true);
620+
static final QueryState INITIAL = new QueryState(-1, false);
621+
static final QueryState CANCELLED = new QueryState(Integer.MIN_VALUE, false);
590622

591623
final int retryCount;
592624
final boolean inProgress;

‎driver-core/src/test/java/com/datastax/driver/core/RequestHandlerTest.java‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
package com.datastax.driver.core;
2+
3+
import java.util.concurrent.TimeUnit;
4+
import java.util.concurrent.TimeoutException;
5+
6+
import com.google.common.collect.ImmutableMap;
7+
import org.scassandra.Scassandra;
8+
import org.scassandra.http.client.PrimingRequest;
9+
import org.testng.annotations.Test;
10+
11+
import static org.assertj.core.api.Assertions.assertThat;
12+
13+
public class RequestHandlerTest {
14+
15+
@Test(groups = "long")
16+
public void should_handle_race_between_response_and_cancellation() {
17+
Scassandra scassandra = TestUtils.createScassandraServer();
18+
Cluster cluster = null;
19+
20+
try {
21+
// Use a mock server that takes a constant time to reply
22+
scassandra.start();
23+
scassandra.primingClient().prime(
24+
PrimingRequest.queryBuilder()
25+
.withQuery("mock query")
26+
.withRows(ImmutableMap.of("key", 1))
27+
.withFixedDelay(10)
28+
.build()
29+
);
30+
31+
cluster = Cluster.builder().addContactPoint("127.0.0.1").withPort(scassandra.getBinaryPort())
32+
.withPoolingOptions(new PoolingOptions()
33+
.setCoreConnectionsPerHost(HostDistance.LOCAL, 1)
34+
.setMaxConnectionsPerHost(HostDistance.LOCAL, 1))
35+
.build();
36+
37+
Session session = cluster.connect();
38+
39+
// To reproduce, we need to cancel the query exactly when the reply arrives.
40+
// Run a few queries to estimate how much that will take.
41+
int samples = 100;
42+
long start = System.currentTimeMillis();
43+
for (int i = 0; i < samples; i++) {
44+
session.execute("mock query");
45+
}
46+
long elapsed = System.currentTimeMillis() - start;
47+
long queryDuration = elapsed / samples;
48+
49+
// Now run queries and cancel them after that estimated time
50+
for (int i = 0; i < 2000; i++) {
51+
ResultSetFuture future = session.executeAsync("mock query");
52+
try {
53+
future.getUninterruptibly(queryDuration, TimeUnit.MILLISECONDS);
54+
} catch (TimeoutException e) {
55+
future.cancel(true);
56+
}
57+
}
58+
59+
PooledConnection connection = getSingleConnection(session);
60+
assertThat(connection.inFlight.get()).isEqualTo(0);
61+
} finally {
62+
if (cluster != null)
63+
cluster.close();
64+
scassandra.stop();
65+
}
66+
}
67+
68+
private PooledConnection getSingleConnection(Session session) {
69+
HostConnectionPool pool = ((SessionManager)session).pools.values().iterator().next();
70+
return pool.connections.get(0);
71+
}
72+
}

‎driver-core/src/test/java/com/datastax/driver/core/TestUtils.java‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,14 +15,18 @@
1515
*/
1616
package com.datastax.driver.core;
1717

18+
import java.io.IOException;
1819
import java.math.BigDecimal;
1920
import java.math.BigInteger;
2021
import java.net.InetAddress;
22+
import java.net.ServerSocket;
2123
import java.nio.ByteBuffer;
2224
import java.util.*;
2325

2426
import static org.testng.Assert.fail;
2527

28+
import org.scassandra.Scassandra;
29+
import org.scassandra.ScassandraFactory;
2630
import org.slf4j.Logger;
2731
import org.slf4j.LoggerFactory;
2832
import org.testng.SkipException;
@@ -382,4 +386,34 @@ public static Host findHost(Cluster cluster, int hostNumber) {
382386
fail(address + " not found in cluster metadata");
383387
return null; // never reached
384388
}
389+
390+
/**
391+
* @return A Scassandra instance with an arbitrarily chosen binary port from 8042-8142 and admin port from
392+
* 8052-8152.
393+
*/
394+
public static Scassandra createScassandraServer() {
395+
int binaryPort = findAvailablePort(8042);
396+
int adminPort = findAvailablePort(8052);
397+
return ScassandraFactory.createServer(binaryPort, adminPort);
398+
}
399+
400+
/**
401+
* @param startingWith The first port to try, if unused will keep trying the next port until one is found up to
402+
* 100 subsequent ports.
403+
* @return A local port that is currently unused.
404+
*/
405+
public static int findAvailablePort(int startingWith) {
406+
IOException last = null;
407+
for(int port = startingWith; port < startingWith+100; port++) {
408+
try {
409+
ServerSocket s = new ServerSocket(port);
410+
s.close();
411+
return port;
412+
} catch (IOException e) {
413+
last = e;
414+
}
415+
}
416+
// If for whatever reason a port could not be acquired throw the last encountered exception.
417+
throw new RuntimeException("Could not acquire an available port", last);
418+
}
385419
}

‎driver-core/src/test/resources/log4j.properties‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
# Set root logger level to DEBUG and its only appender to A1.
22
log4j.rootLogger=INFO, A1
33

4+
# Scassandra's info log is a bit verbose
5+
log4j.logger.org.scassandra=WARN
6+
47
# A1 is set to be a ConsoleAppender.
58
log4j.appender.A1=org.apache.log4j.ConsoleAppender
69

0 commit comments

Comments
 (0)