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.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.metrics.SessionMetricUpdater;
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<ResultT> extends AdminRequestHandler<ResultT>
implements Throttled {
public static ThrottledAdminRequestHandler<AdminResult> query(
DriverChannel channel,
Message message,
Map<String, ByteBuffer> customPayload,
Duration timeout,
RequestThrottler throttler,
SessionMetricUpdater metricUpdater,
String logPrefix,
String debugString) {
return new ThrottledAdminRequestHandler<>(
channel,
message,
customPayload,
timeout,
throttler,
metricUpdater,
logPrefix,
debugString,
Rows.class);
}
public static ThrottledAdminRequestHandler<ByteBuffer> prepare(
DriverChannel channel,
Message message,
Map<String, ByteBuffer> customPayload,
Duration timeout,
RequestThrottler throttler,
SessionMetricUpdater metricUpdater,
String logPrefix) {
return new ThrottledAdminRequestHandler<>(
channel,
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,
Message message,
Map<String, ByteBuffer> customPayload,
Duration timeout,
RequestThrottler throttler,
SessionMetricUpdater metricUpdater,
String logPrefix,
String debugString,
Class<? extends Result> expectedResponseType) {
super(channel, message, customPayload, timeout, logPrefix, debugString, expectedResponseType);
this.startTimeNanos = System.nanoTime();
this.throttler = throttler;
this.metricUpdater = metricUpdater;
}
@Override
public CompletionStage<ResultT> start() {
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;
}
}