com.datastax.oss.driver.internal.core.adminrequest.ThrottledAdminRequestHandler 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.RequestThrottlingException;
import com.datastax.oss.driver.api.core.metadata.Node;
import com.datastax.oss.driver.api.core.metrics.DefaultSessionMetric;
import com.datastax.oss.driver.api.core.session.throttling.RequestThrottler;
import com.datastax.oss.driver.api.core.session.throttling.Throttled;
import com.datastax.oss.driver.internal.core.channel.DriverChannel;
import com.datastax.oss.driver.internal.core.control.ControlConnection;
import com.datastax.oss.driver.internal.core.metrics.SessionMetricUpdater;
import com.datastax.oss.driver.internal.core.pool.ChannelPool;
import com.datastax.oss.driver.internal.core.session.DefaultSession;
import com.datastax.oss.protocol.internal.Message;
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 edu.umd.cs.findbugs.annotations.NonNull;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.util.Map;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.TimeUnit;
import net.jcip.annotations.ThreadSafe;
@ThreadSafe
public class ThrottledAdminRequestHandler extends AdminRequestHandler
implements Throttled {
/**
* @param shouldPreAcquireId whether to call {@link DriverChannel#preAcquireId()} before sending
* the request. This must be false if you obtained the connection from a pool ({@link
* ChannelPool#next()}, or {@link DefaultSession#getChannel(Node, String)}). It must be
* true if you are using a standalone channel (e.g. in {@link ControlConnection} or one of
* its auxiliary components).
*/
public static ThrottledAdminRequestHandler query(
DriverChannel channel,
boolean shouldPreAcquireId,
Message message,
Map customPayload,
Duration timeout,
RequestThrottler throttler,
SessionMetricUpdater metricUpdater,
String logPrefix,
String debugString) {
return new ThrottledAdminRequestHandler<>(
channel,
shouldPreAcquireId,
message,
customPayload,
timeout,
throttler,
metricUpdater,
logPrefix,
debugString,
Rows.class);
}
/**
* @param shouldPreAcquireId whether to call {@link DriverChannel#preAcquireId()} before sending
* the request. See {@link #query(DriverChannel, boolean, Message, Map, Duration,
* RequestThrottler, SessionMetricUpdater, String, String)} for more explanations.
*/
public static ThrottledAdminRequestHandler prepare(
DriverChannel channel,
boolean shouldPreAcquireId,
Message message,
Map customPayload,
Duration timeout,
RequestThrottler throttler,
SessionMetricUpdater metricUpdater,
String logPrefix) {
return new ThrottledAdminRequestHandler<>(
channel,
shouldPreAcquireId,
message,
customPayload,
timeout,
throttler,
metricUpdater,
logPrefix,
message.toString(),
Prepared.class);
}
private final long startTimeNanos;
private final RequestThrottler throttler;
private final SessionMetricUpdater metricUpdater;
protected ThrottledAdminRequestHandler(
DriverChannel channel,
boolean preAcquireId,
Message message,
Map customPayload,
Duration timeout,
RequestThrottler throttler,
SessionMetricUpdater metricUpdater,
String logPrefix,
String debugString,
Class expectedResponseType) {
super(
channel,
preAcquireId,
message,
customPayload,
timeout,
logPrefix,
debugString,
expectedResponseType);
this.startTimeNanos = System.nanoTime();
this.throttler = throttler;
this.metricUpdater = metricUpdater;
}
@Override
public CompletionStage start() {
// Don't write request yet, wait for green light from throttler
throttler.register(this);
return result;
}
@Override
public void onThrottleReady(boolean wasDelayed) {
if (wasDelayed) {
metricUpdater.updateTimer(
DefaultSessionMetric.THROTTLING_DELAY,
null,
System.nanoTime() - startTimeNanos,
TimeUnit.NANOSECONDS);
}
super.start();
}
@Override
public void onThrottleFailure(@NonNull RequestThrottlingException error) {
metricUpdater.incrementCounter(DefaultSessionMetric.THROTTLING_ERRORS, null);
setFinalError(error);
}
@Override
protected boolean setFinalResult(ResultT result) {
boolean wasSet = super.setFinalResult(result);
if (wasSet) {
throttler.signalSuccess(this);
}
return wasSet;
}
@Override
protected boolean setFinalError(Throwable error) {
boolean wasSet = super.setFinalError(error);
if (wasSet) {
if (error instanceof DriverTimeoutException) {
throttler.signalTimeout(this);
} else if (!(error instanceof RequestThrottlingException)) {
throttler.signalError(this, error);
}
}
return wasSet;
}
}