com.datastax.oss.driver.internal.core.adminrequest.AdminRequestHandler Maven / Gradle / Ivy
The newest version!
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you 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.oss.driver.internal.core.adminrequest;
import com.datastax.oss.driver.api.core.DriverTimeoutException;
import com.datastax.oss.driver.api.core.ProtocolVersion;
import com.datastax.oss.driver.api.core.connection.BusyConnectionException;
import com.datastax.oss.driver.api.core.cql.Statement;
import com.datastax.oss.driver.api.core.type.codec.TypeCodecs;
import com.datastax.oss.driver.internal.core.channel.DriverChannel;
import com.datastax.oss.driver.internal.core.channel.ResponseCallback;
import com.datastax.oss.driver.internal.core.util.concurrent.UncaughtExceptions;
import com.datastax.oss.driver.shaded.guava.common.collect.Maps;
import com.datastax.oss.protocol.internal.Frame;
import com.datastax.oss.protocol.internal.Message;
import com.datastax.oss.protocol.internal.ProtocolConstants;
import com.datastax.oss.protocol.internal.request.Query;
import com.datastax.oss.protocol.internal.request.query.QueryOptions;
import com.datastax.oss.protocol.internal.response.Result;
import com.datastax.oss.protocol.internal.response.result.Prepared;
import com.datastax.oss.protocol.internal.response.result.Rows;
import io.netty.util.concurrent.Future;
import io.netty.util.concurrent.ScheduledFuture;
import java.net.InetAddress;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.TimeUnit;
import net.jcip.annotations.ThreadSafe;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/** Handles the lifecycle of an admin request (such as a node refresh or schema refresh query). */
@ThreadSafe
public class AdminRequestHandler implements ResponseCallback {
private static final Logger LOG = LoggerFactory.getLogger(AdminRequestHandler.class);
public static AdminRequestHandler call(
DriverChannel channel, Query query, Duration timeout, String logPrefix) {
return new AdminRequestHandler<>(
channel,
true,
query,
Frame.NO_PAYLOAD,
timeout,
logPrefix,
"call '" + query.query + "'",
com.datastax.oss.protocol.internal.response.result.Void.class);
}
public static AdminRequestHandler query(
DriverChannel channel,
String query,
Map parameters,
Duration timeout,
int pageSize,
String logPrefix) {
Query message =
new Query(
query,
buildQueryOptions(pageSize, serialize(parameters, channel.protocolVersion()), null));
String debugString = "query '" + message.query + "'";
if (!parameters.isEmpty()) {
debugString += " with parameters " + parameters;
}
return new AdminRequestHandler<>(
channel, true, message, Frame.NO_PAYLOAD, timeout, logPrefix, debugString, Rows.class);
}
public static AdminRequestHandler query(
DriverChannel channel, String query, Duration timeout, int pageSize, String logPrefix) {
return query(channel, query, Collections.emptyMap(), timeout, pageSize, logPrefix);
}
private final DriverChannel channel;
private final boolean shouldPreAcquireId;
private final Message message;
private final Map customPayload;
private final Duration timeout;
private final String logPrefix;
private final String debugString;
private final Class expectedResponseType;
protected final CompletableFuture result = new CompletableFuture<>();
// This is only ever accessed on the channel's event loop, so it doesn't need to be volatile
private ScheduledFuture timeoutFuture;
protected AdminRequestHandler(
DriverChannel channel,
boolean shouldPreAcquireId,
Message message,
Map customPayload,
Duration timeout,
String logPrefix,
String debugString,
Class expectedResponseType) {
this.channel = channel;
this.shouldPreAcquireId = shouldPreAcquireId;
this.message = message;
this.customPayload = customPayload;
this.timeout = timeout;
this.logPrefix = logPrefix;
this.debugString = debugString;
this.expectedResponseType = expectedResponseType;
}
public CompletionStage start() {
LOG.debug("[{}] Executing {}", logPrefix, this);
if (shouldPreAcquireId && !channel.preAcquireId()) {
setFinalError(
new BusyConnectionException(
String.format(
"%s has reached its maximum number of simultaneous requests", channel)));
} else {
channel.write(message, false, customPayload, this).addListener(this::onWriteComplete);
}
return result;
}
private void onWriteComplete(Future future) {
if (future.isSuccess()) {
LOG.debug("[{}] Successfully wrote {}, waiting for response", logPrefix, this);
if (timeout.toNanos() > 0) {
timeoutFuture =
channel
.eventLoop()
.schedule(this::fireTimeout, timeout.toNanos(), TimeUnit.NANOSECONDS);
timeoutFuture.addListener(UncaughtExceptions::log);
}
} else {
setFinalError(future.cause());
}
}
private void fireTimeout() {
setFinalError(
new DriverTimeoutException(String.format("%s timed out after %s", debugString, timeout)));
if (!channel.closeFuture().isDone()) {
channel.cancel(this);
}
}
@Override
public void onFailure(Throwable error) {
if (timeoutFuture != null) {
timeoutFuture.cancel(true);
}
setFinalError(error);
}
@Override
public void onResponse(Frame responseFrame) {
if (timeoutFuture != null) {
timeoutFuture.cancel(true);
}
Message message = responseFrame.message;
LOG.debug("[{}] Got response {}", logPrefix, responseFrame.message);
if (!expectedResponseType.isInstance(message)) {
// Note that this also covers error responses, no need to get too fancy here
setFinalError(new UnexpectedResponseException(debugString, message));
} else if (expectedResponseType == Rows.class) {
Rows rows = (Rows) message;
ByteBuffer pagingState = rows.getMetadata().pagingState;
AdminRequestHandler nextHandler = (pagingState == null) ? null : this.copy(pagingState);
// The public factory methods guarantee that expectedResponseType and ResultT always match:
@SuppressWarnings("unchecked")
ResultT result = (ResultT) new AdminResult(rows, nextHandler, channel.protocolVersion());
setFinalResult(result);
} else if (expectedResponseType == Prepared.class) {
Prepared prepared = (Prepared) message;
@SuppressWarnings("unchecked")
ResultT result = (ResultT) ByteBuffer.wrap(prepared.preparedQueryId);
setFinalResult(result);
} else if (expectedResponseType
== com.datastax.oss.protocol.internal.response.result.Void.class) {
setFinalResult(null);
} else {
setFinalError(new AssertionError("Unhandled response type" + expectedResponseType));
}
}
protected boolean setFinalResult(ResultT result) {
return this.result.complete(result);
}
protected boolean setFinalError(Throwable error) {
return result.completeExceptionally(error);
}
private AdminRequestHandler copy(ByteBuffer pagingState) {
assert message instanceof Query;
Query current = (Query) this.message;
QueryOptions currentOptions = current.options;
QueryOptions newOptions =
buildQueryOptions(currentOptions.pageSize, currentOptions.namedValues, pagingState);
return new AdminRequestHandler<>(
channel,
// This is called for next page queries, so we always need to reacquire an id:
true,
new Query(current.query, newOptions),
customPayload,
timeout,
logPrefix,
debugString,
expectedResponseType);
}
private static QueryOptions buildQueryOptions(
int pageSize, Map serialize, ByteBuffer pagingState) {
return new QueryOptions(
ProtocolConstants.ConsistencyLevel.ONE,
Collections.emptyList(),
serialize,
false,
pageSize,
pagingState,
ProtocolConstants.ConsistencyLevel.SERIAL,
Statement.NO_DEFAULT_TIMESTAMP,
null,
Statement.NO_NOW_IN_SECONDS);
}
private static Map serialize(
Map parameters, ProtocolVersion protocolVersion) {
Map result = Maps.newHashMapWithExpectedSize(parameters.size());
for (Map.Entry entry : parameters.entrySet()) {
result.put(entry.getKey(), serialize(entry.getValue(), protocolVersion));
}
return result;
}
private static ByteBuffer serialize(Object parameter, ProtocolVersion protocolVersion) {
if (parameter instanceof String) {
return TypeCodecs.TEXT.encode((String) parameter, protocolVersion);
} else if (parameter instanceof InetAddress) {
return TypeCodecs.INET.encode((InetAddress) parameter, protocolVersion);
} else if (parameter instanceof List && ((List) parameter).get(0) instanceof String) {
@SuppressWarnings("unchecked")
List l = (List) parameter;
return AdminRow.LIST_OF_TEXT.encode(l, protocolVersion);
} else if (parameter instanceof Integer) {
return TypeCodecs.INT.encode((Integer) parameter, protocolVersion);
} else {
throw new IllegalArgumentException(
"Unsupported variable type for admin query: " + parameter.getClass());
}
}
@Override
public String toString() {
return debugString;
}
}