sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/RetryHttpRequestInitializer.java - beam - Git at Google

 /*
  * 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.extensions.gcp.util;

 import com.google.api.client.http.HttpIOExceptionHandler;
 import com.google.api.client.http.HttpRequest;
 import com.google.api.client.http.HttpRequestInitializer;
 import com.google.api.client.http.HttpResponse;
 import com.google.api.client.http.HttpResponseInterceptor;
 import com.google.api.client.http.HttpUnsuccessfulResponseHandler;
 import com.google.api.client.util.BackOff;
 import com.google.api.client.util.ExponentialBackOff;
 import com.google.api.client.util.NanoClock;
 import com.google.api.client.util.Sleeper;
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.Set;
 import javax.annotation.Nullable;
 import org.apache.beam.sdk.metrics.Counter;
 import org.apache.beam.sdk.metrics.Metrics;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 /**
  * Implements a request initializer that adds retry handlers to all HttpRequests.
  *
  * <p>Also can take a HttpResponseInterceptor to be applied to the responses.
  */
 public class RetryHttpRequestInitializer implements HttpRequestInitializer {

   private static final Logger LOG = LoggerFactory.getLogger(RetryHttpRequestInitializer.class);

   /** Http response codes that should be silently ignored. */
   private static final Set<Integer> DEFAULT_IGNORED_RESPONSE_CODES =
       new HashSet<>(
           Arrays.asList(
               307 /* Redirect, handled by the client library */,
               308 /* Resume Incomplete, handled by the client library */));

   /** Http response timeout to use for hanging gets. */
   private static final int HANGING_GET_TIMEOUT_SEC = 80;

   private int writeTimeout;

   /** Handlers used to provide additional logging information on unsuccessful HTTP requests. */
   private static class LoggingHttpBackOffHandler
       implements HttpIOExceptionHandler, HttpUnsuccessfulResponseHandler {

     private final Sleeper sleeper;
     private final BackOff ioExceptionBackOff;
     private final BackOff unsuccessfulResponseBackOff;
     private final Set<Integer> ignoredResponseCodes;
     // aggregate the total time spent in exponential backoff
     private final Counter throttlingSeconds =
         Metrics.counter(LoggingHttpBackOffHandler.class, "cumulativeThrottlingSeconds");
     private int ioExceptionRetries;
     private int unsuccessfulResponseRetries;
     @Nullable private CustomHttpErrors customHttpErrors;

     private LoggingHttpBackOffHandler(
         Sleeper sleeper,
         BackOff ioExceptionBackOff,
         BackOff unsucessfulResponseBackOff,
         Set<Integer> ignoredResponseCodes,
         @Nullable CustomHttpErrors customHttpErrors) {
       this.sleeper = sleeper;
       this.ioExceptionBackOff = ioExceptionBackOff;
       this.unsuccessfulResponseBackOff = unsucessfulResponseBackOff;
       this.ignoredResponseCodes = ignoredResponseCodes;
       this.customHttpErrors = customHttpErrors;
     }

     @Override
     public boolean handleIOException(HttpRequest request, boolean supportsRetry)
         throws IOException {
       // We will retry if the request supports retry or the backoff was successful.
       // Note that the order of these checks is important since
       // backOffWasSuccessful will perform a sleep.
       boolean willRetry = supportsRetry && backOffWasSuccessful(ioExceptionBackOff);
       if (willRetry) {
         ioExceptionRetries += 1;
         LOG.debug("Request failed with IOException, will retry: {}", request.getUrl());
       } else {
         String message =
             "Request failed with IOException, "
                 + "performed {} retries due to IOExceptions, "
                 + "performed {} retries due to unsuccessful status codes, "
                 + "HTTP framework says request {} be retried, "
                 + "(caller responsible for retrying): {}";
         LOG.warn(
             message,
             ioExceptionRetries,
             unsuccessfulResponseRetries,
             supportsRetry ? "can" : "cannot",
             request.getUrl());
       }
       return willRetry;
     }

     @Override
     public boolean handleResponse(HttpRequest request, HttpResponse response, boolean supportsRetry)
         throws IOException {
       // We will retry if the request supports retry and the status code requires a backoff
       // and the backoff was successful. Note that the order of these checks is important since
       // backOffWasSuccessful will perform a sleep.
       boolean willRetry =
           supportsRetry
               && retryOnStatusCode(response.getStatusCode())
               && backOffWasSuccessful(unsuccessfulResponseBackOff);
       if (willRetry) {
         unsuccessfulResponseRetries += 1;
         LOG.debug(
             "Request failed with code {}, will retry: {}",
             response.getStatusCode(),
             request.getUrl());
       } else {

         String message =
             "Request failed with code {}, "
                 + "performed {} retries due to IOExceptions, "
                 + "performed {} retries due to unsuccessful status codes, "
                 + "HTTP framework says request {} be retried, "
                 + "(caller responsible for retrying): {}. {}";
         String customLogMessage = "";
         if (customHttpErrors != null) {
           String error =
               customHttpErrors.getCustomError(
                   new HttpRequestWrapper(request), new HttpResponseWrapper(response));
           if (error != null) {
             customLogMessage = error;
           }
         }
         if (ignoredResponseCodes.contains(response.getStatusCode())) {
           // Log ignored response codes at a lower level
           LOG.debug(
               message,
               response.getStatusCode(),
               ioExceptionRetries,
               unsuccessfulResponseRetries,
               supportsRetry ? "can" : "cannot",
               request.getUrl(),
               customLogMessage);
         } else {
           LOG.warn(
               message,
               response.getStatusCode(),
               ioExceptionRetries,
               unsuccessfulResponseRetries,
               supportsRetry ? "can" : "cannot",
               request.getUrl(),
               customLogMessage);
         }
       }
       return willRetry;
     }

     /** Returns true iff performing the backoff was successful. */
     private boolean backOffWasSuccessful(BackOff backOff) {
       try {
         long backOffTime = backOff.nextBackOffMillis();
         if (backOffTime == BackOff.STOP) {
           return false;
         }
         throttlingSeconds.inc(backOffTime / 1000);
         sleeper.sleep(backOffTime);
         return true;
       } catch (InterruptedException | IOException e) {
         return false;
       }
     }

     /** Returns true iff the {@code statusCode} represents an error that should be retried. */
     private boolean retryOnStatusCode(int statusCode) {
       return (statusCode == 0) // Code 0 usually means no response / network error
           || (statusCode / 100 == 5) // 5xx: server error
           || statusCode == 429; // 429: Too many requests
     }
   }

   private final HttpResponseInterceptor responseInterceptor; // response Interceptor to use

   private CustomHttpErrors customHttpErrors = null;

   private final NanoClock nanoClock; // used for testing

   private final Sleeper sleeper; // used for testing

   private Set<Integer> ignoredResponseCodes = new HashSet<>(DEFAULT_IGNORED_RESPONSE_CODES);

   public RetryHttpRequestInitializer() {
     this(Collections.emptyList());
   }

   /**
    * @param additionalIgnoredResponseCodes a list of HTTP status codes that should not be logged.
    */
   public RetryHttpRequestInitializer(Collection<Integer> additionalIgnoredResponseCodes) {
     this(additionalIgnoredResponseCodes, null);
   }

   /**
    * @param additionalIgnoredResponseCodes a list of HTTP status codes that should not be logged.
    * @param responseInterceptor HttpResponseInterceptor to be applied on all requests. May be null.
    */
   public RetryHttpRequestInitializer(
       Collection<Integer> additionalIgnoredResponseCodes,
       @Nullable HttpResponseInterceptor responseInterceptor) {
     this(NanoClock.SYSTEM, Sleeper.DEFAULT, additionalIgnoredResponseCodes, responseInterceptor);
   }

   /**
    * Visible for testing.
    *
    * @param nanoClock used as a timing source for knowing how much time has elapsed.
    * @param sleeper used to sleep between retries.
    * @param additionalIgnoredResponseCodes a list of HTTP status codes that should not be logged.
    */
   RetryHttpRequestInitializer(
       NanoClock nanoClock,
       Sleeper sleeper,
       Collection<Integer> additionalIgnoredResponseCodes,
       HttpResponseInterceptor responseInterceptor) {
     this.nanoClock = nanoClock;
     this.sleeper = sleeper;
     this.ignoredResponseCodes.addAll(additionalIgnoredResponseCodes);
     this.responseInterceptor = responseInterceptor;
     this.writeTimeout = 0;
   }

   @Override
   public void initialize(HttpRequest request) throws IOException {
     // Set a timeout for hanging-gets.
     // TODO: Do this exclusively for work requests.
     request.setReadTimeout(HANGING_GET_TIMEOUT_SEC * 1000);
     request.setWriteTimeout(this.writeTimeout);

     LoggingHttpBackOffHandler loggingHttpBackOffHandler =
         new LoggingHttpBackOffHandler(
             sleeper,
             // Back off on retryable http errors and IOExceptions.
             // A back-off multiplier of 2 raises the maximum request retrying time
             // to approximately 5 minutes (keeping other back-off parameters to
             // their default values).
             new ExponentialBackOff.Builder().setNanoClock(nanoClock).setMultiplier(2).build(),
             new ExponentialBackOff.Builder().setNanoClock(nanoClock).setMultiplier(2).build(),
             ignoredResponseCodes,
             this.customHttpErrors);

     request.setUnsuccessfulResponseHandler(loggingHttpBackOffHandler);
     request.setIOExceptionHandler(loggingHttpBackOffHandler);

     // Set response initializer
     if (responseInterceptor != null) {
       request.setResponseInterceptor(responseInterceptor);
     }
   }

   public void setCustomErrors(CustomHttpErrors customErrors) {
     this.customHttpErrors = customErrors;
   }

   public void setWriteTimeout(int writeTimeout) {
     this.writeTimeout = writeTimeout;
   }
 }
	/*
	* 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.extensions.gcp.util;

	import com.google.api.client.http.HttpIOExceptionHandler;
	import com.google.api.client.http.HttpRequest;
	import com.google.api.client.http.HttpRequestInitializer;
	import com.google.api.client.http.HttpResponse;
	import com.google.api.client.http.HttpResponseInterceptor;
	import com.google.api.client.http.HttpUnsuccessfulResponseHandler;
	import com.google.api.client.util.BackOff;
	import com.google.api.client.util.ExponentialBackOff;
	import com.google.api.client.util.NanoClock;
	import com.google.api.client.util.Sleeper;
	import java.io.IOException;
	import java.util.Arrays;
	import java.util.Collection;
	import java.util.Collections;
	import java.util.HashSet;
	import java.util.Set;
	import javax.annotation.Nullable;
	import org.apache.beam.sdk.metrics.Counter;
	import org.apache.beam.sdk.metrics.Metrics;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	/**
	* Implements a request initializer that adds retry handlers to all HttpRequests.
	*
	* <p>Also can take a HttpResponseInterceptor to be applied to the responses.
	*/
	public class RetryHttpRequestInitializer implements HttpRequestInitializer {

	private static final Logger LOG = LoggerFactory.getLogger(RetryHttpRequestInitializer.class);

	/** Http response codes that should be silently ignored. */
	private static final Set<Integer> DEFAULT_IGNORED_RESPONSE_CODES =
	new HashSet<>(
	Arrays.asList(
	307 /* Redirect, handled by the client library */,
	308 /* Resume Incomplete, handled by the client library */));

	/** Http response timeout to use for hanging gets. */
	private static final int HANGING_GET_TIMEOUT_SEC = 80;

	private int writeTimeout;

	/** Handlers used to provide additional logging information on unsuccessful HTTP requests. */
	private static class LoggingHttpBackOffHandler
	implements HttpIOExceptionHandler, HttpUnsuccessfulResponseHandler {

	private final Sleeper sleeper;
	private final BackOff ioExceptionBackOff;
	private final BackOff unsuccessfulResponseBackOff;
	private final Set<Integer> ignoredResponseCodes;
	// aggregate the total time spent in exponential backoff
	private final Counter throttlingSeconds =
	Metrics.counter(LoggingHttpBackOffHandler.class, "cumulativeThrottlingSeconds");
	private int ioExceptionRetries;
	private int unsuccessfulResponseRetries;
	@Nullable private CustomHttpErrors customHttpErrors;

	private LoggingHttpBackOffHandler(
	Sleeper sleeper,
	BackOff ioExceptionBackOff,
	BackOff unsucessfulResponseBackOff,
	Set<Integer> ignoredResponseCodes,
	@Nullable CustomHttpErrors customHttpErrors) {
	this.sleeper = sleeper;
	this.ioExceptionBackOff = ioExceptionBackOff;
	this.unsuccessfulResponseBackOff = unsucessfulResponseBackOff;
	this.ignoredResponseCodes = ignoredResponseCodes;
	this.customHttpErrors = customHttpErrors;
	}

	@Override
	public boolean handleIOException(HttpRequest request, boolean supportsRetry)
	throws IOException {
	// We will retry if the request supports retry or the backoff was successful.
	// Note that the order of these checks is important since
	// backOffWasSuccessful will perform a sleep.
	boolean willRetry = supportsRetry && backOffWasSuccessful(ioExceptionBackOff);
	if (willRetry) {
	ioExceptionRetries += 1;
	LOG.debug("Request failed with IOException, will retry: {}", request.getUrl());
	} else {
	String message =
	"Request failed with IOException, "
	+ "performed {} retries due to IOExceptions, "
	+ "performed {} retries due to unsuccessful status codes, "
	+ "HTTP framework says request {} be retried, "
	+ "(caller responsible for retrying): {}";
	LOG.warn(
	message,
	ioExceptionRetries,
	unsuccessfulResponseRetries,
	supportsRetry ? "can" : "cannot",
	request.getUrl());
	}
	return willRetry;
	}

	@Override
	public boolean handleResponse(HttpRequest request, HttpResponse response, boolean supportsRetry)
	throws IOException {
	// We will retry if the request supports retry and the status code requires a backoff
	// and the backoff was successful. Note that the order of these checks is important since
	// backOffWasSuccessful will perform a sleep.
	boolean willRetry =
	supportsRetry
	&& retryOnStatusCode(response.getStatusCode())
	&& backOffWasSuccessful(unsuccessfulResponseBackOff);
	if (willRetry) {
	unsuccessfulResponseRetries += 1;
	LOG.debug(
	"Request failed with code {}, will retry: {}",
	response.getStatusCode(),
	request.getUrl());
	} else {

	String message =
	"Request failed with code {}, "
	+ "performed {} retries due to IOExceptions, "
	+ "performed {} retries due to unsuccessful status codes, "
	+ "HTTP framework says request {} be retried, "
	+ "(caller responsible for retrying): {}. {}";
	String customLogMessage = "";
	if (customHttpErrors != null) {
	String error =
	customHttpErrors.getCustomError(
	new HttpRequestWrapper(request), new HttpResponseWrapper(response));
	if (error != null) {
	customLogMessage = error;
	}
	}
	if (ignoredResponseCodes.contains(response.getStatusCode())) {
	// Log ignored response codes at a lower level
	LOG.debug(
	message,
	response.getStatusCode(),
	ioExceptionRetries,
	unsuccessfulResponseRetries,
	supportsRetry ? "can" : "cannot",
	request.getUrl(),
	customLogMessage);
	} else {
	LOG.warn(
	message,
	response.getStatusCode(),
	ioExceptionRetries,
	unsuccessfulResponseRetries,
	supportsRetry ? "can" : "cannot",
	request.getUrl(),
	customLogMessage);
	}
	}
	return willRetry;
	}

	/** Returns true iff performing the backoff was successful. */
	private boolean backOffWasSuccessful(BackOff backOff) {
	try {
	long backOffTime = backOff.nextBackOffMillis();
	if (backOffTime == BackOff.STOP) {
	return false;
	}
	throttlingSeconds.inc(backOffTime / 1000);
	sleeper.sleep(backOffTime);
	return true;
	} catch (InterruptedException \| IOException e) {
	return false;
	}
	}

	/** Returns true iff the {@code statusCode} represents an error that should be retried. */
	private boolean retryOnStatusCode(int statusCode) {
	return (statusCode == 0) // Code 0 usually means no response / network error
	\|\| (statusCode / 100 == 5) // 5xx: server error
	\|\| statusCode == 429; // 429: Too many requests
	}
	}

	private final HttpResponseInterceptor responseInterceptor; // response Interceptor to use

	private CustomHttpErrors customHttpErrors = null;

	private final NanoClock nanoClock; // used for testing

	private final Sleeper sleeper; // used for testing

	private Set<Integer> ignoredResponseCodes = new HashSet<>(DEFAULT_IGNORED_RESPONSE_CODES);

	public RetryHttpRequestInitializer() {
	this(Collections.emptyList());
	}

	/**
	* @param additionalIgnoredResponseCodes a list of HTTP status codes that should not be logged.
	*/
	public RetryHttpRequestInitializer(Collection<Integer> additionalIgnoredResponseCodes) {
	this(additionalIgnoredResponseCodes, null);
	}

	/**
	* @param additionalIgnoredResponseCodes a list of HTTP status codes that should not be logged.
	* @param responseInterceptor HttpResponseInterceptor to be applied on all requests. May be null.
	*/
	public RetryHttpRequestInitializer(
	Collection<Integer> additionalIgnoredResponseCodes,
	@Nullable HttpResponseInterceptor responseInterceptor) {
	this(NanoClock.SYSTEM, Sleeper.DEFAULT, additionalIgnoredResponseCodes, responseInterceptor);
	}

	/**
	* Visible for testing.
	*
	* @param nanoClock used as a timing source for knowing how much time has elapsed.
	* @param sleeper used to sleep between retries.
	* @param additionalIgnoredResponseCodes a list of HTTP status codes that should not be logged.
	*/
	RetryHttpRequestInitializer(
	NanoClock nanoClock,
	Sleeper sleeper,
	Collection<Integer> additionalIgnoredResponseCodes,
	HttpResponseInterceptor responseInterceptor) {
	this.nanoClock = nanoClock;
	this.sleeper = sleeper;
	this.ignoredResponseCodes.addAll(additionalIgnoredResponseCodes);
	this.responseInterceptor = responseInterceptor;
	this.writeTimeout = 0;
	}

	@Override
	public void initialize(HttpRequest request) throws IOException {
	// Set a timeout for hanging-gets.
	// TODO: Do this exclusively for work requests.
	request.setReadTimeout(HANGING_GET_TIMEOUT_SEC * 1000);
	request.setWriteTimeout(this.writeTimeout);

	LoggingHttpBackOffHandler loggingHttpBackOffHandler =
	new LoggingHttpBackOffHandler(
	sleeper,
	// Back off on retryable http errors and IOExceptions.
	// A back-off multiplier of 2 raises the maximum request retrying time
	// to approximately 5 minutes (keeping other back-off parameters to
	// their default values).
	new ExponentialBackOff.Builder().setNanoClock(nanoClock).setMultiplier(2).build(),
	new ExponentialBackOff.Builder().setNanoClock(nanoClock).setMultiplier(2).build(),
	ignoredResponseCodes,
	this.customHttpErrors);

	request.setUnsuccessfulResponseHandler(loggingHttpBackOffHandler);
	request.setIOExceptionHandler(loggingHttpBackOffHandler);

	// Set response initializer
	if (responseInterceptor != null) {
	request.setResponseInterceptor(responseInterceptor);
	}
	}

	public void setCustomErrors(CustomHttpErrors customErrors) {
	this.customHttpErrors = customErrors;
	}

	public void setWriteTimeout(int writeTimeout) {
	this.writeTimeout = writeTimeout;
	}
	}