extends UnaryCallSettings<BulkMutation, MutateRowsAttemptResult> {

// This settings is just a simple wrapper for BatchingCallSettings to allow us to add

// additional functionality.

private final BatchingCallSettings<RowMutationEntry, Void, BulkMutation, MutateRowsAttemptResult>

batchingCallSettings;

private final boolean isLatencyBasedThrottlingEnabled;

private final Long targetRpcLatencyMs;

private final DynamicFlowControlSettings dynamicFlowControlSettings;

private final boolean isServerInitiatedFlowControlEnabled;

private BigtableBatchingCallSettings(Builder builder) {

super(builder);

batchingCallSettings =

BatchingCallSettings.newBuilder(builder.batchingDescriptor)

.setBatchingSettings(builder.batchingSettings)

.setRetrySettings(builder.getRetrySettings())

.setRetryableCodes(builder.getRetryableCodes())

.build();

this.isLatencyBasedThrottlingEnabled = builder.isLatencyBasedThrottlingEnabled;

this.targetRpcLatencyMs = builder.targetRpcLatencyMs;

this.dynamicFlowControlSettings = builder.dynamicFlowControlSettings;

this.isServerInitiatedFlowControlEnabled = builder.isServerInitiatedFlowControlEnabled;

}

/** Returns batching settings which contains multiple batch threshold levels. */

public BatchingSettings getBatchingSettings() {

return batchingCallSettings.getBatchingSettings();

}

/** Returns an adapter that packs and unpacks batching elements. */

BatchingDescriptor<RowMutationEntry, Void, BulkMutation, MutateRowsAttemptResult>

getBatchingDescriptor() {

return batchingCallSettings.getBatchingDescriptor();

}

/** Gets if latency based throttling is enabled. */

public boolean isLatencyBasedThrottlingEnabled() {

return isLatencyBasedThrottlingEnabled;

}

/** Gets target rpc latency if latency based throttling is enabled. Otherwise returns null. */

@Nullable

public Long getTargetRpcLatencyMs() {

return targetRpcLatencyMs;

}

/**

* Gets {@link DynamicFlowControlSettings}.

*

* @see Builder#getDynamicFlowControlSettings()

*/

DynamicFlowControlSettings getDynamicFlowControlSettings() {

return dynamicFlowControlSettings;

}

/** Gets if flow control is enabled. */

@InternalApi("Intended for use by the Bigtable dataflow connectors only")

public boolean isServerInitiatedFlowControlEnabled() {

return isServerInitiatedFlowControlEnabled;

}

static Builder newBuilder(

BatchingDescriptor<RowMutationEntry, Void, BulkMutation, MutateRowsAttemptResult>

batchingDescriptor) {

return new Builder(batchingDescriptor);

}

/**

* Get a builder with the same values as this object. See the class documentation of {@link

* BigtableBatchingCallSettings} for a sample settings configuration.

*/

@Override

public final Builder toBuilder() {

return new Builder(this);

}

@Override

public String toString() {

return MoreObjects.toStringHelper(this)

.add("batchingCallSettings", batchingCallSettings)

.add("isLatencyBasedThrottlingEnabled", isLatencyBasedThrottlingEnabled)

.add("targetRpcLatency", targetRpcLatencyMs)

.add("dynamicFlowControlSettings", dynamicFlowControlSettings)

.add("isServerInitiatedFlowControlEnabled", isServerInitiatedFlowControlEnabled)

.toString();

}

/**

* A base builder class for {@link BigtableBatchingCallSettings}. See the class documentation of

* {@link BigtableBatchingCallSettings} for a description of the different values that can be set.

*/

public static class Builder

extends UnaryCallSettings.Builder<BulkMutation, MutateRowsAttemptResult> {

private BatchingDescriptor<RowMutationEntry, Void, BulkMutation, MutateRowsAttemptResult>

batchingDescriptor;

private BatchingSettings batchingSettings;

private boolean isLatencyBasedThrottlingEnabled;

private Long targetRpcLatencyMs;

private DynamicFlowControlSettings dynamicFlowControlSettings;

private boolean isServerInitiatedFlowControlEnabled;

private Builder(

@Nonnull

BatchingDescriptor<RowMutationEntry, Void, BulkMutation, MutateRowsAttemptResult>

batchingDescriptor) {

this.batchingDescriptor =

Preconditions.checkNotNull(batchingDescriptor, "batching descriptor can't be null");

}

private Builder(@Nonnull BigtableBatchingCallSettings settings) {

super(settings);

this.batchingDescriptor = settings.getBatchingDescriptor();

this.batchingSettings = settings.getBatchingSettings();

this.isLatencyBasedThrottlingEnabled = settings.isLatencyBasedThrottlingEnabled();

this.targetRpcLatencyMs = settings.getTargetRpcLatencyMs();

this.dynamicFlowControlSettings = settings.getDynamicFlowControlSettings();

this.isServerInitiatedFlowControlEnabled = settings.isServerInitiatedFlowControlEnabled();

}

/** Sets the batching settings with various thresholds. */

public Builder setBatchingSettings(@Nonnull BatchingSettings batchingSettings) {

Preconditions.checkNotNull(batchingSettings, "batching settings can't be null");

this.batchingSettings = batchingSettings;

return this;

}

/** Returns the {@link BatchingSettings}. */

public BatchingSettings getBatchingSettings() {

return batchingSettings;

}

/** Sets the rpc failure {@link StatusCode.Code code}, for which retries should be performed. */

@Override

public Builder setRetryableCodes(StatusCode.Code... codes) {

super.setRetryableCodes(codes);

return this;

}

/** Sets the rpc failure {@link StatusCode.Code code}, for which retries should be performed. */

@Override

public Builder setRetryableCodes(Set<StatusCode.Code> retryableCodes) {

super.setRetryableCodes(retryableCodes);

return this;

}

/** Sets the {@link RetrySettings} values for each retry attempts. */

@Override

public Builder setRetrySettings(@Nonnull RetrySettings retrySettings) {

super.setRetrySettings(retrySettings);

return this;

}

/**

* Enable latency based throttling. The number of allowed in-flight requests will be adjusted to

* reach the target rpc latency.

*/

public Builder enableLatencyBasedThrottling(long targetRpcLatency) {

Preconditions.checkArgument(

targetRpcLatency > 0, "target RPC latency must be greater than 0");

this.isLatencyBasedThrottlingEnabled = true;

this.targetRpcLatencyMs = targetRpcLatency;

return this;

}

/** Disable latency based throttling. */

public Builder disableLatencyBasedThrottling() {

this.isLatencyBasedThrottlingEnabled = false;

this.targetRpcLatencyMs = null;

return this;

}

/** Gets target rpc latency if latency based throttling is enabled. Otherwise returns null. */

@Nullable

public Long getTargetRpcLatencyMs() {

return isLatencyBasedThrottlingEnabled ? targetRpcLatencyMs : null;

}

/** Gets if latency based throttling is enabled. */

public boolean isLatencyBasedThrottlingEnabled() {

return this.isLatencyBasedThrottlingEnabled;

}

/**

* Gets the {@link DynamicFlowControlSettings} that'll be used to set up a {@link

* FlowController} for throttling.

*

* <p>By default, this will allow a maximum of 1000 entries per channel of {@link

* FlowControlSettings.Builder#setMaxOutstandingElementCount request count} and 100MB of {@link

* FlowControlSettings.Builder#setMaxOutstandingRequestBytes accumulated size} in-flight

* requests. Once the limits are reached, pending operations will by default be {@link

* FlowControlSettings.Builder#setLimitExceededBehavior blocked} until some of the in-flight

* requests are resolved.

*

* <p>If latency based throttling is enabled, number of entries allowed by {@link

* FlowController} will be adjusted to reach {@link Builder#getTargetRpcLatencyMs()}.

*

* <ul>

* <li>{@link FlowController} will be set to allow Math.max({@link BatchingSettings.Builder

* #setElementCountThreshold batch size}, {@link

* FlowControlSettings.Builder#setMaxOutstandingElementCount request count} / 4) entries

* to start with.

* <li>If bulk mutation rpc latency is higher than target latency, decrease allowed entries to

* a minimum of Math.max({@link BatchingSettings.Builder#setElementCountThreshold batch

* size}, {@link FlowControlSettings.Builder#setMaxOutstandingElementCount request count}

* / 100).

* <li>If bulk mutation rpc latency is lower than target latency and there was throttling,

* increase allowed entries to a maximum of {@link

* FlowControlSettings.Builder#setMaxOutstandingElementCount request count}.

* </ul>

*

* If latency based throttling is disabled, {@link FlowController} will always allow {@link

* FlowControlSettings.Builder#setMaxOutstandingElementCount request count}.

*

* <p>Latency based throttling only updates outstanding entries count. {@link FlowController}

* will always allow {@link FlowControlSettings.Builder#setMaxOutstandingRequestBytes

* accumulated size}.

*/

DynamicFlowControlSettings getDynamicFlowControlSettings() {

return this.dynamicFlowControlSettings;

}

/** Configure flow control based on the current load of the Bigtable server. */

@InternalApi("Intended for use by the Bigtable dataflow connectors only")

public Builder setServerInitiatedFlowControl(boolean isEnable) {

this.isServerInitiatedFlowControlEnabled = isEnable;

return this;

}

/** Gets if flow control is enabled based on the load of the Bigtable server. */

@InternalApi("Intended for use by the Bigtable dataflow connectors only")

public boolean isServerInitiatedFlowControlEnabled() {

return this.isServerInitiatedFlowControlEnabled;

}

/** Builds the {@link BigtableBatchingCallSettings} object with provided configuration. */

@Override

public BigtableBatchingCallSettings build() {

Preconditions.checkState(batchingSettings != null, "batchingSettings must be set");

FlowControlSettings flowControlSettings = batchingSettings.getFlowControlSettings();

Preconditions.checkState(

flowControlSettings.getMaxOutstandingElementCount() != null,

"maxOutstandingElementCount must be set in BatchingSettings#FlowControlSettings");

Preconditions.checkState(

flowControlSettings.getMaxOutstandingRequestBytes() != null,

"maxOutstandingRequestBytes must be set in BatchingSettings#FlowControlSettings");

Preconditions.checkArgument(

batchingSettings.getElementCountThreshold() == null

|| flowControlSettings.getMaxOutstandingElementCount()

> batchingSettings.getElementCountThreshold(),

"if batch elementCountThreshold is set in BatchingSettings, flow control maxOutstandingElementCount must be > elementCountThreshold");

Preconditions.checkArgument(

batchingSettings.getRequestByteThreshold() == null

|| flowControlSettings.getMaxOutstandingRequestBytes()

> batchingSettings.getRequestByteThreshold(),

"if batch requestByteThreshold is set in BatchingSettings, flow control maxOutstandingRequestBytes must be > getRequestByteThreshold");

// Combine static FlowControlSettings with latency based throttling settings to create

// DynamicFlowControlSettings.

if (isLatencyBasedThrottlingEnabled()) {

long maxThrottlingElementCount = flowControlSettings.getMaxOutstandingElementCount();

long maxThrottlingRequestByteCount = flowControlSettings.getMaxOutstandingRequestBytes();

// The maximum in flight element count is pretty high. Set the initial parallelism to 25%

// of the maximum and then work up or down. This reduction should reduce the

// impacts of a bursty job, such as those found in Dataflow.

long initialElementCount = maxThrottlingElementCount / 4;

// Decreases are floored at 1% of the maximum so that there is some level of

// throughput.

long minElementCount = maxThrottlingElementCount / 100;

// Make sure initialOutstandingElementCount and minOutstandingElementCount element count are

// greater or equal to batch size to avoid deadlocks.

if (batchingSettings.getElementCountThreshold() != null) {

initialElementCount =

Math.max(initialElementCount, batchingSettings.getElementCountThreshold());

minElementCount = Math.max(minElementCount, batchingSettings.getElementCountThreshold());

}

dynamicFlowControlSettings =

DynamicFlowControlSettings.newBuilder()

.setLimitExceededBehavior(flowControlSettings.getLimitExceededBehavior())

.setInitialOutstandingElementCount(initialElementCount)

.setMaxOutstandingElementCount(maxThrottlingElementCount)

.setMinOutstandingElementCount(minElementCount)

.setInitialOutstandingRequestBytes(maxThrottlingRequestByteCount)

.setMinOutstandingRequestBytes(maxThrottlingRequestByteCount)

.setMaxOutstandingRequestBytes(maxThrottlingRequestByteCount)

.build();

} else {

dynamicFlowControlSettings =

DynamicFlowControlSettings.newBuilder()

.setLimitExceededBehavior(flowControlSettings.getLimitExceededBehavior())

.setInitialOutstandingElementCount(

flowControlSettings.getMaxOutstandingElementCount())

.setMaxOutstandingElementCount(flowControlSettings.getMaxOutstandingElementCount())

.setMinOutstandingElementCount(flowControlSettings.getMaxOutstandingElementCount())

.setInitialOutstandingRequestBytes(

flowControlSettings.getMaxOutstandingRequestBytes())

.setMinOutstandingRequestBytes(flowControlSettings.getMaxOutstandingRequestBytes())

.setMaxOutstandingRequestBytes(flowControlSettings.getMaxOutstandingRequestBytes())

.build();

}

return new BigtableBatchingCallSettings(this);

}

}