/*
* Copyright (C) 2012-2015 DataStax Inc.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.datastax.driver.core;
import com.datastax.driver.core.exceptions.InvalidQueryException;
import com.datastax.driver.core.utils.CassandraVersion;
import com.google.common.base.Function;
import com.google.common.base.Throwables;
import com.google.common.collect.Lists;
import com.google.common.util.concurrent.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
import java.nio.ByteBuffer;
import java.util.Collection;
import java.util.Map;
import java.util.concurrent.*;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.fail;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertTrue;
public class AsyncQueryTest extends CCMTestsSupport {
Logger logger = LoggerFactory.getLogger(AsyncQueryTest.class);
@DataProvider(name = "keyspace")
public static Object[][] keyspace() {
return new Object[][]{{"asyncquerytest"}, {"\"AsyncQueryTest\""}};
}
@Override
public void onTestContextInitialized() {
for (Object[] objects : keyspace()) {
String keyspace = (String) objects[0];
execute(
String.format("create keyspace %s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}", keyspace),
String.format("create table %s.foo(k int, v int, primary key (k, v))", keyspace));
for (int v = 1; v <= 100; v++)
execute(String.format("insert into %s.foo (k, v) values (1, %d)", keyspace, v));
}
}
/**
* Checks that a cancelled query releases the connection (JAVA-407).
*/
@Test(groups = "short")
public void cancelled_query_should_release_the_connection() throws InterruptedException {
ResultSetFuture future = session().executeAsync("select release_version from system.local");
future.cancel(true);
assertTrue(future.isCancelled());
TimeUnit.MILLISECONDS.sleep(100);
HostConnectionPool pool = getPool(session());
for (Connection connection : pool.connections) {
assertEquals(connection.inFlight.get(), 0);
}
}
@Test(groups = "short")
public void should_init_cluster_and_session_if_needed() throws Exception {
// For this test we need an uninitialized cluster, so we can't reuse the one provided by the
// parent class. Rebuild a new one with the same (unique) host.
Host host = cluster().getMetadata().allHosts().iterator().next();
Cluster cluster2 = register(Cluster.builder()
.addContactPointsWithPorts(Lists.newArrayList(host.getSocketAddress()))
.build());
Session session2 = cluster2.newSession();
// Neither cluster2 nor session2 are initialized at this point
assertThat(cluster2.manager.metadata).isNull();
ResultSetFuture future = session2.executeAsync("select release_version from system.local");
Row row = Uninterruptibles.getUninterruptibly(future).one();
assertThat(row.getString(0)).isNotEmpty();
}
@Test(groups = "short", dataProvider = "keyspace", enabled = false,
description = "disabled because the blocking USE call in the current pool implementation makes it deadlock")
public void should_chain_query_on_async_session_init_with_same_executor(String keyspace) throws Exception {
ListenableFuture<Integer> resultFuture = connectAndQuery(keyspace, GuavaCompatibility.INSTANCE.sameThreadExecutor());
Integer result = Uninterruptibles.getUninterruptibly(resultFuture);
assertThat(result).isEqualTo(1);
}
@Test(groups = "short", dataProvider = "keyspace")
public void should_chain_query_on_async_session_init_with_different_executor(String keyspace) throws Exception {
ExecutorService executor = Executors.newFixedThreadPool(1);
ListenableFuture<Integer> resultFuture = connectAndQuery(keyspace, executor);
Integer result = Uninterruptibles.getUninterruptibly(resultFuture);
assertThat(result).isEqualTo(1);
executor.shutdownNow();
}
@Test(groups = "short")
public void should_propagate_error_to_chained_query_if_session_init_fails() throws Exception {
ListenableFuture<Integer> resultFuture = connectAndQuery("wrong_keyspace", GuavaCompatibility.INSTANCE.sameThreadExecutor());
try {
Uninterruptibles.getUninterruptibly(resultFuture);
} catch (ExecutionException e) {
assertThat(e.getCause())
.isInstanceOf(InvalidQueryException.class)
.hasMessage("Keyspace 'wrong_keyspace' does not exist");
}
}
@Test(groups = "short")
public void should_fail_when_synchronous_call_on_io_thread() throws Exception {
for (int i = 0; i < 1000; i++) {
ResultSetFuture f = session().executeAsync("select release_version from system.local");
ListenableFuture<Thread> f2 = Futures.transform(f, new Function<ResultSet, Thread>() {
@Override
public Thread apply(ResultSet input) {
session().execute("select release_version from system.local");
return Thread.currentThread();
}
});
if (isFailed(f2, IllegalStateException.class, "Detected a synchronous call on an I/O thread")) {
return;
}
}
fail("callback was not executed on io thread in 1000 attempts, something may be wrong.");
}
@Test(groups = "short")
public void should_fail_when_synchronous_call_on_io_thread_with_session_wrapper() throws Exception {
final Session session = new SessionWrapper(session());
for (int i = 0; i < 1000; i++) {
ResultSetFuture f = session.executeAsync("select release_version from system.local");
ListenableFuture<Thread> f2 = Futures.transform(f, new Function<ResultSet, Thread>() {
@Override
public Thread apply(ResultSet input) {
session.execute("select release_version from system.local");
return Thread.currentThread();
}
});
if (isFailed(f2, IllegalStateException.class, "Detected a synchronous call on an I/O thread")) {
return;
}
}
fail("callback was not executed on io thread in 1000 attempts, something may be wrong.");
}
@Test(groups = "short")
@CassandraVersion(value = "2.0.0", description = "Paging is not supported until 2.0")
public void should_fail_when_auto_paging_on_io_thread() throws Exception {
for (int i = 0; i < 1000; i++) {
Statement statement = new SimpleStatement("select v from asyncquerytest.foo where k = 1");
// Ensure results will be paged (there are 100 rows in the test table)
statement.setFetchSize(10);
ResultSetFuture f = session().executeAsync(statement);
ListenableFuture<Thread> f2 = Futures.transform(f, new Function<ResultSet, Thread>() {
@Override
public Thread apply(ResultSet rs) {
rs.all(); // Consume the whole result set
return Thread.currentThread();
}
});
if (isFailed(f2, IllegalStateException.class, "Detected a synchronous call on an I/O thread")) {
return;
}
}
fail("callback was not executed on io thread in 1000 attempts, something may be wrong.");
}
private boolean isFailed(ListenableFuture<Thread> future, Class<?> expectedException, String expectedMessage) {
try {
Thread executedThread = future.get();
if (executedThread != Thread.currentThread()) {
fail("Expected a failed future, callback was executed on " + executedThread);
} else {
// Callback was invoked on the same thread, which indicates that the future completed
// before the transform callback was registered. Try again to produce case where callback
// is called on io thread.
logger.warn("Future completed before transform callback registered, will try again.");
}
} catch (Exception e) {
assertThat(Throwables.getRootCause(e))
.isInstanceOf(expectedException)
.hasMessageContaining(expectedMessage);
return true;
}
return false;
}
private ListenableFuture<Integer> connectAndQuery(String keyspace, Executor executor) {
ListenableFuture<Session> sessionFuture = cluster().connectAsync(keyspace);
ListenableFuture<ResultSet> queryFuture = GuavaCompatibility.INSTANCE.transformAsync(sessionFuture, new AsyncFunction<Session, ResultSet>() {
@Override
public ListenableFuture<ResultSet> apply(Session session) throws Exception {
return session.executeAsync("select v from foo where k = 1");
}
}, executor);
return Futures.transform(queryFuture, new Function<ResultSet, Integer>() {
@Override
public Integer apply(ResultSet rs) {
return rs.one().getInt(0);
}
}, executor);
}
private static HostConnectionPool getPool(Session session) {
Collection<HostConnectionPool> pools = ((SessionManager) session).pools.values();
assertEquals(pools.size(), 1);
return pools.iterator().next();
}
private static class SessionWrapper extends AbstractSession {
private final Session session;
public SessionWrapper(Session session) {
this.session = session;
}
@Override
public ResultSet execute(Statement statement) {
// test a custom call to checkNotInEventLoop()
checkNotInEventLoop();
return executeAsync(statement).getUninterruptibly();
}
@Override
public String getLoggedKeyspace() {
return session.getLoggedKeyspace();
}
@Override
public Session init() {
return session.init();
}
@Override
public ListenableFuture<Session> initAsync() {
return session.initAsync();
}
@Override
public ResultSetFuture executeAsync(Statement statement) {
return session.executeAsync(statement);
}
@Override
public CloseFuture closeAsync() {
return session.closeAsync();
}
@Override
public boolean isClosed() {
return session.isClosed();
}
@Override
public Cluster getCluster() {
return session.getCluster();
}
@Override
public State getState() {
return session.getState();
}
@Override
protected ListenableFuture<PreparedStatement> prepareAsync(String query, Map<String, ByteBuffer> customPayload) {
return ((SessionManager) session).prepareAsync(query, customPayload);
}
}
}