BenWhitehead commented on a change in pull request #14261: URL: https://github.com/apache/beam/pull/14261#discussion_r627767434
########## File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/RpcQosImpl.java ########## @@ -0,0 +1,909 @@ +/* + * 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 org.apache.beam.sdk.io.gcp.firestore; + +import com.google.api.gax.grpc.GrpcStatusCode; +import com.google.api.gax.rpc.ApiException; +import com.google.api.gax.rpc.StatusCode; +import com.google.rpc.Code; +import java.io.IOException; +import java.util.Iterator; +import java.util.Optional; +import java.util.Random; +import java.util.Set; +import java.util.WeakHashMap; +import java.util.function.Function; +import org.apache.beam.sdk.io.gcp.firestore.RpcQos.RpcAttempt.Context; +import org.apache.beam.sdk.io.gcp.firestore.RpcQos.RpcWriteAttempt.Element; +import org.apache.beam.sdk.io.gcp.firestore.RpcQos.RpcWriteAttempt.FlushBuffer; +import org.apache.beam.sdk.metrics.Counter; +import org.apache.beam.sdk.metrics.Distribution; +import org.apache.beam.sdk.metrics.MetricName; +import org.apache.beam.sdk.transforms.Sum; +import org.apache.beam.sdk.util.BackOff; +import org.apache.beam.sdk.util.BackOffUtils; +import org.apache.beam.sdk.util.FluentBackoff; +import org.apache.beam.sdk.util.MovingFunction; +import org.apache.beam.sdk.util.Sleeper; +import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting; +import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet; +import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.primitives.Ints; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.joda.time.Interval; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +final class RpcQosImpl implements RpcQos { + + /** Non-retryable errors. See https://cloud.google.com/apis/design/errors#handling_errors. */ + private static final Set<Integer> NON_RETRYABLE_ERROR_NUMBERS = + ImmutableSet.of( + Code.ALREADY_EXISTS, + Code.DATA_LOSS, + Code.FAILED_PRECONDITION, + Code.INVALID_ARGUMENT, + Code.OUT_OF_RANGE, + Code.NOT_FOUND, + Code.PERMISSION_DENIED, + Code.UNIMPLEMENTED) + .stream() + .map(Code::getNumber) + .collect(ImmutableSet.toImmutableSet()); + /** + * The target minimum number of requests per samplePeriodMs, even if no requests succeed. Must be + * greater than 0, else we could throttle to zero. Because every decision is probabilistic, there + * is no guarantee that the request rate in any given interval will not be zero. (This is the +1 + * from the formula in https://landing.google.com/sre/book/chapters/handling-overload.html) + */ + private static final double MIN_REQUESTS = 1; + + private final RpcQosOptions options; + + private final AdaptiveThrottler at; + private final WriteBatcher wb; + private final WriteRampUp writeRampUp; + private final FluentBackoff fb; + + private final WeakHashMap<Context, O11y> counters; + private final Random random; + private final Sleeper sleeper; + private final Function<Context, O11y> computeCounters; + private final DistributionFactory distributionFactory; + + RpcQosImpl( + RpcQosOptions options, + Random random, + Sleeper sleeper, + CounterFactory counterFactory, + DistributionFactory distributionFactory) { + this.options = options; + this.random = random; + this.sleeper = sleeper; + DistributionFactory filteringDistributionFactory = + new DiagnosticOnlyFilteringDistributionFactory( + !options.isShouldReportDiagnosticMetrics(), distributionFactory); + this.distributionFactory = filteringDistributionFactory; + at = + new AdaptiveThrottler( + options.getSamplePeriod(), + options.getSamplePeriodBucketSize(), + options.getThrottleDuration(), + options.getOverloadRatio()); + wb = + new WriteBatcher( + options.getSamplePeriod(), + options.getSamplePeriodBucketSize(), + options.getBatchInitialCount(), + options.getBatchTargetLatency(), + filteringDistributionFactory); + writeRampUp = + new WriteRampUp( + Math.max(1, 500 / options.getHintMaxNumWorkers()), filteringDistributionFactory); + // maxRetries is an inclusive value, we want exclusive since we are tracking all attempts + fb = + FluentBackoff.DEFAULT + .withMaxRetries(options.getMaxAttempts() - 1) + .withInitialBackoff(options.getInitialBackoff()); + counters = new WeakHashMap<>(); + computeCounters = (Context c) -> O11y.create(c, counterFactory, filteringDistributionFactory); + } + + @Override + public RpcWriteAttemptImpl newWriteAttempt(Context context) { + return new RpcWriteAttemptImpl( + context, counters.computeIfAbsent(context, computeCounters), fb.backoff(), sleeper); + } + + @Override + public RpcReadAttemptImpl newReadAttempt(Context context) { + return new RpcReadAttemptImpl( + context, counters.computeIfAbsent(context, computeCounters), fb.backoff(), sleeper); + } + + @Override + public boolean bytesOverLimit(long bytes) { + return bytes > options.getBatchMaxBytes(); + } + + private static MovingFunction createMovingFunction(Duration samplePeriod, Duration sampleUpdate) { + return new MovingFunction( + samplePeriod.getMillis(), + sampleUpdate.getMillis(), + 1 /* numSignificantBuckets */, + 1 /* numSignificantSamples */, + Sum.ofLongs()); + } + + private enum AttemptState { + PENDING, + STARTED, + COMPLETE_SUCCESS, + COMPLETE_ERROR; + + public void checkActive() { + switch (this) { + case PENDING: + case STARTED: + return; + case COMPLETE_SUCCESS: + throw new IllegalStateException( + "Expected state to be PENDING or STARTED, but was COMPLETE_SUCCESS"); + case COMPLETE_ERROR: + throw new IllegalStateException( + "Expected state to be PENDING or STARTED, but was COMPLETE_ERROR"); + } + } + + public void checkStarted() { + switch (this) { + case STARTED: + return; + case PENDING: + throw new IllegalStateException("Expected state to be STARTED, but was PENDING"); + case COMPLETE_SUCCESS: + throw new IllegalStateException("Expected state to be STARTED, but was COMPLETE_SUCCESS"); + case COMPLETE_ERROR: + throw new IllegalStateException("Expected state to be STARTED, but was COMPLETE_ERROR"); + } + } + } + + private abstract class BaseRpcAttempt implements RpcAttempt { + private final Logger logger; + final O11y o11y; + final BackOff backoff; + final Sleeper sleeper; + + AttemptState state; + Instant start; + + @SuppressWarnings( + "initialization.fields.uninitialized") // allow transient fields to be managed by component + // lifecycle + BaseRpcAttempt(Context context, O11y o11y, BackOff backoff, Sleeper sleeper) { + this.logger = LoggerFactory.getLogger(String.format("%s.RpcQos", context.getNamespace())); + this.o11y = o11y; + this.backoff = backoff; + this.sleeper = sleeper; + this.state = AttemptState.PENDING; + } + + @Override + public boolean awaitSafeToProceed(Instant instant) throws InterruptedException { + state.checkActive(); + Duration shouldThrottleRequest = at.shouldThrottleRequest(instant); + if (shouldThrottleRequest.compareTo(Duration.ZERO) > 0) { + long throttleRequestMillis = shouldThrottleRequest.getMillis(); + logger.debug("Delaying request by {}ms", throttleRequestMillis); + throttleRequest(shouldThrottleRequest); + return false; + } + + return true; + } + + @Override + public void checkCanRetry(RuntimeException exception) throws InterruptedException { + state.checkActive(); + + Optional<ApiException> findApiException = findApiException(exception); + + if (findApiException.isPresent()) { + ApiException apiException = findApiException.get(); + // order here is semi-important + // First we always want to test if the error code is one of the codes we have deemed + // non-retryable before delegating to the exceptions default set. + if (maxAttemptsExhausted() + || getStatusCodeNumber(apiException) + .map(NON_RETRYABLE_ERROR_NUMBERS::contains) + .orElse(false) + || !apiException.isRetryable()) { + state = AttemptState.COMPLETE_ERROR; + throw apiException; + } + } else { + state = AttemptState.COMPLETE_ERROR; + throw exception; + } + } + + @Override + public void completeSuccess() { + state.checkActive(); + state = AttemptState.COMPLETE_SUCCESS; + } + + @Override + public boolean isCodeRetryable(Code code) { + return !NON_RETRYABLE_ERROR_NUMBERS.contains(code.getNumber()); + } + + @Override + public void recordRequestSuccessful(Instant end) { + state.checkStarted(); + o11y.rpcSuccesses.inc(); + o11y.rpcDurationMs.update(durationMs(end)); + at.recordSuccessfulRequest(start); + } + + @Override + public void recordRequestFailed(Instant end) { + state.checkStarted(); + o11y.rpcFailures.inc(); + o11y.rpcDurationMs.update(durationMs(end)); + at.recordFailedRequest(start); + } + + private boolean maxAttemptsExhausted() throws InterruptedException { + try { + boolean exhausted = !BackOffUtils.next(sleeper, backoff); + if (exhausted) { + logger.error("Max attempts exhausted after {} attempts.", options.getMaxAttempts()); + } + return exhausted; + } catch (IOException e) { + // We are using FluentBackoff which does not ever throw an IOException from its methods + // Catch and wrap any potential IOException as a RuntimeException since it won't ever + // happen unless the implementation of FluentBackoff changes. + throw new RuntimeException(e); + } + } + + Logger getLogger() { + return logger; + } + + final void throttleRequest(Duration shouldThrottleRequest) throws InterruptedException { + o11y.throttlingMs.inc(shouldThrottleRequest.getMillis()); + sleeper.sleep(shouldThrottleRequest.getMillis()); + } + + final long durationMs(Instant end) { + return end.minus(start.getMillis()).getMillis(); + } + + private Optional<Integer> getStatusCodeNumber(ApiException apiException) { + StatusCode statusCode = apiException.getStatusCode(); + if (statusCode instanceof GrpcStatusCode) { + GrpcStatusCode grpcStatusCode = (GrpcStatusCode) statusCode; + return Optional.of(grpcStatusCode.getTransportCode().value()); + } + return Optional.empty(); + } + + private Optional<ApiException> findApiException(Throwable throwable) { + if (throwable instanceof ApiException) { + ApiException apiException = (ApiException) throwable; + return Optional.of(apiException); + } else { + Throwable cause = throwable.getCause(); + if (cause != null) { + return findApiException(cause); + } else { + return Optional.empty(); + } + } + } + } + + private final class RpcReadAttemptImpl extends BaseRpcAttempt implements RpcReadAttempt { + private RpcReadAttemptImpl(Context context, O11y o11y, BackOff backoff, Sleeper sleeper) { + super(context, o11y, backoff, sleeper); + } + + @Override + public void recordRequestStart(Instant start) { + at.recordStartRequest(start); + this.start = start; + state = AttemptState.STARTED; + } + + @Override + public void recordStreamValue(Instant now) { + state.checkActive(); + o11y.rpcStreamValueReceived.inc(); + } + } + + final class RpcWriteAttemptImpl extends BaseRpcAttempt implements RpcWriteAttempt { + + private RpcWriteAttemptImpl(Context context, O11y o11y, BackOff backoff, Sleeper sleeper) { + super(context, o11y, backoff, sleeper); + } + + @Override + public boolean awaitSafeToProceed(Instant instant) throws InterruptedException { + state.checkActive(); + Optional<Duration> shouldThrottle = writeRampUp.shouldThrottle(instant); + if (shouldThrottle.isPresent()) { + Duration throttleDuration = shouldThrottle.get(); + long throttleDurationMillis = throttleDuration.getMillis(); + getLogger().debug("Still ramping up, Delaying request by {}ms", throttleDurationMillis); + throttleRequest(throttleDuration); + return false; + } else { + return super.awaitSafeToProceed(instant); + } + } + + @Override + public <T, ElementT extends Element<T>> FlushBufferImpl<T, ElementT> newFlushBuffer( + Instant instantSinceEpoch) { + state.checkActive(); + int availableWriteCountBudget = writeRampUp.getAvailableWriteCountBudget(instantSinceEpoch); + int nextBatchMaxCount = wb.nextBatchMaxCount(instantSinceEpoch); + int batchMaxCount = + Ints.min( + Math.max(0, availableWriteCountBudget), + Math.max(0, nextBatchMaxCount), + options.getBatchMaxCount()); + o11y.batchCapacityCount.update(batchMaxCount); + return new FlushBufferImpl<>(batchMaxCount, options.getBatchMaxBytes()); + } + + @Override + public void recordRequestStart(Instant start, int numWrites) { + at.recordStartRequest(start, numWrites); + writeRampUp.recordWriteCount(start, numWrites); + this.start = start; + state = AttemptState.STARTED; + } + + @Override + public void recordWriteCounts(Instant end, int successfulWrites, int failedWrites) { + int totalWrites = successfulWrites + failedWrites; + state.checkStarted(); + wb.recordRequestLatency(start, end, totalWrites, o11y.latencyPerDocumentMs); + if (successfulWrites > 0) { + at.recordSuccessfulRequest(start, successfulWrites); + } + if (failedWrites > 0) { + at.recordFailedRequest(start, failedWrites); + } + } + } + + /** + * Determines batch sizes based on past performance. + * + * <p>It aims for a target response time per RPC: it uses the response times for previous RPCs and + * the number of documents contained in them, calculates a rolling average time-per-document, and Review comment: Correct, right now it is only tracking number of documents not the size of those documents. It could be extended if we felt so inclined to do so at some point. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: [email protected]
