com.datastax.oss.driver.internal.core.channel.InFlightHandler 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.channel;
import com.datastax.oss.driver.api.core.CqlIdentifier;
import com.datastax.oss.driver.api.core.DriverException;
import com.datastax.oss.driver.api.core.ProtocolVersion;
import com.datastax.oss.driver.api.core.connection.BusyConnectionException;
import com.datastax.oss.driver.api.core.connection.ClosedConnectionException;
import com.datastax.oss.driver.api.core.connection.HeartbeatException;
import com.datastax.oss.driver.internal.core.channel.DriverChannel.RequestMessage;
import com.datastax.oss.driver.internal.core.channel.DriverChannel.SetKeyspaceEvent;
import com.datastax.oss.driver.internal.core.protocol.FrameDecodingException;
import com.datastax.oss.driver.internal.core.util.Loggers;
import com.datastax.oss.driver.shaded.guava.common.collect.BiMap;
import com.datastax.oss.driver.shaded.guava.common.collect.HashBiMap;
import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableSet;
import com.datastax.oss.protocol.internal.Frame;
import com.datastax.oss.protocol.internal.Message;
import com.datastax.oss.protocol.internal.request.Query;
import com.datastax.oss.protocol.internal.response.result.SetKeyspace;
import io.netty.channel.ChannelDuplexHandler;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelPromise;
import io.netty.util.concurrent.Promise;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import net.jcip.annotations.NotThreadSafe;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/** Manages requests that are currently executing on a channel. */
@NotThreadSafe
public class InFlightHandler extends ChannelDuplexHandler {
private static final Logger LOG = LoggerFactory.getLogger(InFlightHandler.class);
private final ProtocolVersion protocolVersion;
private final StreamIdGenerator streamIds;
final ChannelPromise closeStartedFuture;
private final String ownerLogPrefix;
private final BiMap inFlight;
private final Map orphaned;
private volatile int orphanedSize; // thread-safe view for metrics
private final long setKeyspaceTimeoutMillis;
private final EventCallback eventCallback;
private final int maxOrphanStreamIds;
private boolean closingGracefully;
private SetKeyspaceRequest setKeyspaceRequest;
private String logPrefix;
InFlightHandler(
ProtocolVersion protocolVersion,
StreamIdGenerator streamIds,
int maxOrphanStreamIds,
long setKeyspaceTimeoutMillis,
ChannelPromise closeStartedFuture,
EventCallback eventCallback,
String ownerLogPrefix) {
this.protocolVersion = protocolVersion;
this.streamIds = streamIds;
this.maxOrphanStreamIds = maxOrphanStreamIds;
this.closeStartedFuture = closeStartedFuture;
this.ownerLogPrefix = ownerLogPrefix;
this.logPrefix = ownerLogPrefix + "|connecting...";
this.inFlight = HashBiMap.create(streamIds.getMaxAvailableIds());
this.orphaned = new HashMap<>(maxOrphanStreamIds);
this.setKeyspaceTimeoutMillis = setKeyspaceTimeoutMillis;
this.eventCallback = eventCallback;
}
@Override
public void channelActive(ChannelHandlerContext ctx) throws Exception {
super.channelActive(ctx);
String channelId = ctx.channel().toString();
this.logPrefix = ownerLogPrefix + "|" + channelId.substring(1, channelId.length() - 1);
}
@Override
public void write(ChannelHandlerContext ctx, Object in, ChannelPromise promise) throws Exception {
if (in == DriverChannel.GRACEFUL_CLOSE_MESSAGE) {
LOG.debug("[{}] Received graceful close request", logPrefix);
startGracefulShutdown(ctx);
} else if (in == DriverChannel.FORCEFUL_CLOSE_MESSAGE) {
LOG.debug("[{}] Received forceful close request, aborting pending queries", logPrefix);
abortAllInFlight(new ClosedConnectionException("Channel was force-closed"));
ctx.channel().close();
} else if (in instanceof HeartbeatException) {
abortAllInFlight(
new ClosedConnectionException("Heartbeat query failed", ((HeartbeatException) in)));
ctx.close();
} else if (in instanceof RequestMessage) {
write(ctx, (RequestMessage) in, promise);
} else if (in instanceof ResponseCallback) {
cancel(ctx, (ResponseCallback) in, promise);
} else {
promise.setFailure(
new IllegalArgumentException("Unsupported message type " + in.getClass().getName()));
}
}
private void write(ChannelHandlerContext ctx, RequestMessage message, ChannelPromise promise) {
if (closingGracefully) {
promise.setFailure(new IllegalStateException("Channel is closing"));
streamIds.cancelPreAcquire();
return;
}
int streamId = streamIds.acquire();
if (streamId < 0) {
// Should not happen with the preAcquire mechanism, but handle gracefully
promise.setFailure(
new BusyConnectionException(
String.format(
"Couldn't acquire a stream id from InFlightHandler on %s", ctx.channel())));
streamIds.cancelPreAcquire();
return;
}
if (inFlight.containsKey(streamId)) {
promise.setFailure(
new IllegalStateException("Found pending callback for stream id " + streamId));
streamIds.cancelPreAcquire();
return;
}
LOG.trace("[{}] Writing {} on stream id {}", logPrefix, message.responseCallback, streamId);
Frame frame =
Frame.forRequest(
protocolVersion.getCode(),
streamId,
message.tracing,
message.customPayload,
message.request);
inFlight.put(streamId, message.responseCallback);
ChannelFuture writeFuture = ctx.write(frame, promise);
writeFuture.addListener(
future -> {
if (future.isSuccess()) {
message.responseCallback.onStreamIdAssigned(streamId);
} else {
release(streamId, ctx);
}
});
}
private void cancel(
ChannelHandlerContext ctx, ResponseCallback responseCallback, ChannelPromise promise) {
Integer streamId = inFlight.inverse().remove(responseCallback);
if (streamId == null) {
LOG.trace(
"[{}] Received cancellation for unknown or already cancelled callback {}, skipping",
logPrefix,
responseCallback);
} else {
LOG.trace(
"[{}] Cancelled callback {} for stream id {}", logPrefix, responseCallback, streamId);
if (closingGracefully && inFlight.isEmpty()) {
LOG.debug("[{}] Last pending query was cancelled, closing channel", logPrefix);
ctx.channel().close();
} else {
// We can't release the stream id, because a response might still come back from the server.
// Keep track of those "orphaned" ids, to release them later if we get a response and the
// callback says it's the last one.
orphaned.put(streamId, responseCallback);
if (orphaned.size() > maxOrphanStreamIds) {
LOG.debug(
"[{}] Orphan stream ids exceeded the configured threshold ({}), closing gracefully",
logPrefix,
maxOrphanStreamIds);
startGracefulShutdown(ctx);
} else {
orphanedSize = orphaned.size();
}
}
}
promise.setSuccess();
}
private void startGracefulShutdown(ChannelHandlerContext ctx) {
if (inFlight.isEmpty()) {
LOG.debug("[{}] No pending queries, completing graceful shutdown now", logPrefix);
ctx.channel().close();
} else {
// remove heartbeat handler from pipeline if present.
ChannelHandler heartbeatHandler = ctx.pipeline().get(ChannelFactory.HEARTBEAT_HANDLER_NAME);
if (heartbeatHandler != null) {
ctx.pipeline().remove(heartbeatHandler);
}
LOG.debug("[{}] There are pending queries, delaying graceful shutdown", logPrefix);
closingGracefully = true;
closeStartedFuture.setSuccess();
}
}
@Override
@SuppressWarnings("NonAtomicVolatileUpdate")
public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
Frame responseFrame = (Frame) msg;
int streamId = responseFrame.streamId;
if (streamId < 0) {
Message event = responseFrame.message;
if (eventCallback == null) {
LOG.debug("[{}] Received event {} but no callback was registered", logPrefix, event);
} else {
LOG.debug("[{}] Received event {}, notifying callback", logPrefix, event);
try {
eventCallback.onEvent(event);
} catch (Throwable t) {
Loggers.warnWithException(
LOG, "[{}] Unexpected error while invoking event handler", logPrefix, t);
}
}
} else {
boolean wasInFlight = true;
ResponseCallback callback = inFlight.get(streamId);
if (callback == null) {
wasInFlight = false;
callback = orphaned.get(streamId);
if (callback == null) {
LOG.trace("[{}] Got response on unknown stream id {}, skipping", logPrefix, streamId);
return;
}
}
try {
if (callback.isLastResponse(responseFrame)) {
LOG.debug(
"[{}] Got last response on {} stream id {}, completing and releasing",
logPrefix,
wasInFlight ? "in-flight" : "orphaned",
streamId);
release(streamId, ctx);
} else {
LOG.trace(
"[{}] Got non-last response on {} stream id {}, still holding",
logPrefix,
wasInFlight ? "in-flight" : "orphaned",
streamId);
}
if (wasInFlight) {
callback.onResponse(responseFrame);
}
} catch (Throwable t) {
if (wasInFlight) {
fail(
callback,
new IllegalArgumentException("Unexpected error while invoking response handler", t));
} else {
// Assume the callback is already completed, so it's better to log
Loggers.warnWithException(
LOG,
"[{}] Unexpected error while invoking response handler on stream id {}",
logPrefix,
t,
streamId);
}
}
}
}
/** Called if an exception was thrown while processing an inbound event (i.e. a response). */
@Override
public void exceptionCaught(ChannelHandlerContext ctx, Throwable exception) throws Exception {
if (exception instanceof FrameDecodingException) {
int streamId = ((FrameDecodingException) exception).streamId;
LOG.debug("[{}] Error while decoding response on stream id {}", logPrefix, streamId);
if (streamId >= 0) {
// We know which request matches the failing response, fail that one only
ResponseCallback responseCallback = inFlight.get(streamId);
if (responseCallback != null) {
fail(responseCallback, exception.getCause());
}
release(streamId, ctx);
} else {
Loggers.warnWithException(
LOG,
"[{}] Unexpected error while decoding incoming event frame",
logPrefix,
exception.getCause());
}
} else {
// Otherwise fail all pending requests
abortAllInFlight(
(exception instanceof HeartbeatException)
? (HeartbeatException) exception
: new ClosedConnectionException("Unexpected error on channel", exception));
ctx.close();
}
}
@Override
public void userEventTriggered(ChannelHandlerContext ctx, Object event) throws Exception {
if (event instanceof SetKeyspaceEvent) {
SetKeyspaceEvent setKeyspaceEvent = (SetKeyspaceEvent) event;
if (this.setKeyspaceRequest != null) {
setKeyspaceEvent.promise.setFailure(
new IllegalStateException(
"Can't call setKeyspace while a keyspace switch is already in progress"));
} else {
LOG.debug(
"[{}] Switching to keyspace {}", logPrefix, setKeyspaceEvent.keyspaceName.asInternal());
this.setKeyspaceRequest = new SetKeyspaceRequest(ctx, setKeyspaceEvent);
this.setKeyspaceRequest.send();
}
} else {
super.userEventTriggered(ctx, event);
}
}
@Override
public void channelInactive(ChannelHandlerContext ctx) throws Exception {
// If the channel was closed normally (normal or forced shutdown), inFlight is already empty by
// the time we get here. So if it's not, it means the channel closed unexpectedly (e.g. the
// connection was dropped).
abortAllInFlight(new ClosedConnectionException("Lost connection to remote peer"));
super.channelInactive(ctx);
}
private void release(int streamId, ChannelHandlerContext ctx) {
LOG.trace("[{}] Releasing stream id {}", logPrefix, streamId);
if (inFlight.remove(streamId) != null) {
// If we're in the middle of an orderly close and this was the last request, actually close
// the channel now
if (closingGracefully && inFlight.isEmpty()) {
LOG.debug("[{}] Done handling the last pending query, closing channel", logPrefix);
ctx.channel().close();
}
} else if (orphaned.remove(streamId) != null) {
orphanedSize = orphaned.size();
}
// Note: it's possible that the callback is in neither map, if we get here after a call to
// abortAllInFlight that already cleared the map (see JAVA-2000)
streamIds.release(streamId);
}
private void abortAllInFlight(DriverException cause) {
abortAllInFlight(cause, null);
}
/**
* @param ignore the ResponseCallback that called this method, if applicable (avoids a recursive
* loop)
*/
private void abortAllInFlight(DriverException cause, ResponseCallback ignore) {
if (!inFlight.isEmpty()) {
// Create a local copy and clear the map immediately. This prevents
// ConcurrentModificationException if aborting one of the handlers recurses back into this
// method.
Set responseCallbacks = ImmutableSet.copyOf(inFlight.values());
inFlight.clear();
for (ResponseCallback responseCallback : responseCallbacks) {
if (responseCallback != ignore) {
fail(responseCallback, cause);
}
}
// It's not necessary to release the stream ids, since we always call this method right before
// closing the channel
}
}
private void fail(ResponseCallback callback, Throwable failure) {
try {
callback.onFailure(failure);
} catch (Throwable throwable) {
// Protect against unexpected errors. We don't have anywhere to report the error (since
// onFailure failed), so log as a last resort.
LOG.error("[{}] Unexpected error while failing {}", logPrefix, callback, throwable);
}
}
int getAvailableIds() {
return streamIds.getAvailableIds();
}
boolean preAcquireId() {
return streamIds.preAcquire();
}
int getInFlight() {
return streamIds.getMaxAvailableIds() - streamIds.getAvailableIds();
}
int getOrphanIds() {
return orphanedSize;
}
private class SetKeyspaceRequest extends ChannelHandlerRequest {
private final CqlIdentifier keyspaceName;
private final Promise promise;
SetKeyspaceRequest(ChannelHandlerContext ctx, SetKeyspaceEvent setKeyspaceEvent) {
super(ctx, setKeyspaceTimeoutMillis);
this.keyspaceName = setKeyspaceEvent.keyspaceName;
this.promise = setKeyspaceEvent.promise;
}
@Override
String describe() {
return "[" + logPrefix + "] Set keyspace request (USE " + keyspaceName.asCql(true) + ")";
}
@Override
Message getRequest() {
return new Query("USE " + keyspaceName.asCql(false));
}
@Override
void onResponse(Message response) {
if (response instanceof SetKeyspace) {
if (promise.trySuccess(null)) {
InFlightHandler.this.setKeyspaceRequest = null;
}
} else {
failOnUnexpected(response);
}
}
@Override
void fail(String message, Throwable cause) {
ClosedConnectionException setKeyspaceException =
new ClosedConnectionException(message, cause);
if (promise.tryFailure(setKeyspaceException)) {
InFlightHandler.this.setKeyspaceRequest = null;
// setKeyspace queries are not triggered directly by the user, but only as a response to a
// successful "USE... query", so the keyspace name should generally be valid. If the
// keyspace switch fails, this could be due to a schema disagreement or a more serious
// error. Rescheduling the switch is impractical, we can't do much better than closing the
// channel and letting it reconnect.
Loggers.warnWithException(
LOG, "[{}] Unexpected error while switching keyspace", logPrefix, setKeyspaceException);
abortAllInFlight(setKeyspaceException, this);
ctx.channel().close();
}
}
}
}