Backport support for custom HTTP status in a backwards compatible way. (#1087)

Author: OlgaMaciaszek

docs/src/main/asciidoc/_configprops.adoc

@@ -57,6 +57,7 @@
 |spring.cloud.loadbalancer.sticky-session |  | Properties for LoadBalancer sticky-session.
 |spring.cloud.loadbalancer.sticky-session.add-service-instance-cookie | `false` | Indicates whether a cookie with the newly selected instance should be added by LoadBalancer.
 |spring.cloud.loadbalancer.sticky-session.instance-id-cookie-name | `sc-lb-instance-id` | The name of the cookie holding the preferred instance id.
+|spring.cloud.loadbalancer.use-raw-status-code-in-response-data | `false` | Indicates that raw status codes should be used in {@link ResponseData}.
 |spring.cloud.loadbalancer.x-forwarded |  | Enabling X-Forwarded Host and Proto Headers.
 |spring.cloud.loadbalancer.x-forwarded.enabled | `false` | To Enable X-Forwarded Headers.
 |spring.cloud.loadbalancer.zone |  | Spring Cloud LoadBalancer zone.

docs/src/main/asciidoc/spring-cloud-commons.adoc

@@ -1218,20 +1218,28 @@ One type of bean that it may be useful to register using <<custom-loadbalancer-c
 
 The `LoadBalancerLifecycle` beans provide callback methods, named `onStart(Request<RC> request)`, `onStartRequest(Request<RC> request, Response<T> lbResponse)` and `onComplete(CompletionContext<RES, T, RC> completionContext)`, that you should implement to specify what actions should take place before and after load-balancing.
 
-`onStart(Request<RC> request)` takes a `Request` object as a parameter. It contains data that is used to select an appropriate instance, including the downstream client request and <<spring-cloud-loadbalancer-hints,hint>>. `onStartRequest` also takes the `Request` object and, additionally, the `Response<T>` object as parameters.  On the other hand, a `CompletionContext` object is provided to the `onComplete(CompletionContext<RES, T, RC> completionContext)` method. It contains the LoadBalancer `Response`, including the selected service instance, the `Status` of the request executed against that service instance and (if available) the response returned to the downstream client, and (if an exception has occurred) the corresponding `Throwable`.
+`onStart(Request<RC> request)` takes a `Request` object as a parameter.
+It contains data that is used to select an appropriate instance, including the downstream client request and <<spring-cloud-loadbalancer-hints,hint>>. `onStartRequest` also takes the `Request` object and, additionally, the `Response<T>` object as parameters.
+On the other hand, a `CompletionContext` object is provided to the `onComplete(CompletionContext<RES, T, RC> completionContext)` method.
+It contains the LoadBalancer `Response`, including the selected service instance, the `Status` of the request executed against that service instance and (if available) the response returned to the downstream client, and (if an exception has occurred) the corresponding `Throwable`.
 
 The `supports(Class requestContextClass, Class responseClass,
-Class serverTypeClass)` method can be used to determine whether the processor in question handles objects of provided types. If not overridden by the user, it returns `true`.
+Class serverTypeClass)` method can be used to determine whether the processor in question handles objects of provided types.
+If not overridden by the user, it returns `true`.
 
 NOTE: In the preceding method calls, `RC` means `RequestContext` type, `RES` means client response type, and `T` means returned server type.
 
+WARNING: If you are using custom HTTP status codes, you will be getting exceptions.
+In order to prevent this, you can set the value of `spring.cloud.loadbalancer.use-raw-status-code-in-response-data`.
+It will cause raw status codes to be used instead of `HttpStatus` enums.
+The `httpStatus` field in `ResponseData` will then be used, but you'll be able to get the raw status code from the `rawHttpStatus` field.
+
 [[loadbalancer-micrometer-stats-lifecycle]]
 === Spring Cloud LoadBalancer Statistics
 
 We provide a `LoadBalancerLifecycle` bean called `MicrometerStatsLoadBalancerLifecycle`, which uses Micrometer to provide statistics for load-balanced calls.
 
-In order to get this bean added to your application context,
-set the value of the `spring.cloud.loadbalancer.stats.micrometer.enabled` to `true` and have a `MeterRegistry` available (for example, by adding https://docs.spring.io/spring-boot/docs/current/reference/html/production-ready-features.html[Spring Boot Actuator] to your project).
+In order to get this bean added to your application context, set the value of the `spring.cloud.loadbalancer.stats.micrometer.enabled` to `true` and have a `MeterRegistry` available (for example, by adding https://docs.spring.io/spring-boot/docs/current/reference/html/production-ready-features.html[Spring Boot Actuator] to your project).
 
 `MicrometerStatsLoadBalancerLifecycle` registers the following meters in `MeterRegistry`:
 

spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/LoadBalancerProperties.java

@@ -68,6 +68,11 @@ public class LoadBalancerProperties {
 	 */
 	private StickySession stickySession = new StickySession();
 
+	/**
+	 * Indicates that raw status codes should be used in {@link ResponseData}.
+	 */
+	private boolean useRawStatusCodeInResponseData;
+
 	public HealthCheck getHealthCheck() {
 		return healthCheck;
 	}
@@ -121,6 +126,14 @@ public XForwarded getXForwarded() {
 		return xForwarded;
 	}
 
+	public boolean isUseRawStatusCodeInResponseData() {
+		return useRawStatusCodeInResponseData;
+	}
+
+	public void setUseRawStatusCodeInResponseData(boolean useRawStatusCodeInResponseData) {
+		this.useRawStatusCodeInResponseData = useRawStatusCodeInResponseData;
+	}
+
 	public static class StickySession {
 
 		/**

spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/ResponseData.java

@@ -48,27 +48,93 @@ public class ResponseData {
 
 	private final RequestData requestData;
 
+	private final Integer rawHttpStatus;
+
+	/**
+	 * @deprecated for removal; new constructors will be added in 4.x
+	 */
+	@Deprecated
+	public ResponseData(HttpHeaders headers, MultiValueMap<String, ResponseCookie> cookies, RequestData requestData,
+			Integer rawHttpStatus) {
+		this.httpStatus = null;
+		this.rawHttpStatus = rawHttpStatus;
+		this.headers = headers;
+		this.cookies = cookies;
+		this.requestData = requestData;
+	}
+
+	/**
+	 * @deprecated for removal; new constructors will be added in 4.x
+	 */
+	@Deprecated
 	public ResponseData(HttpStatus httpStatus, HttpHeaders headers, MultiValueMap<String, ResponseCookie> cookies,
 			RequestData requestData) {
 		this.httpStatus = httpStatus;
+		this.rawHttpStatus = httpStatus != null ? httpStatus.value() : null;
 		this.headers = headers;
 		this.cookies = cookies;
 		this.requestData = requestData;
 	}
 
+	/**
+	 * @deprecated for removal; new constructors will be added in 4.x
+	 */
+	@Deprecated
 	public ResponseData(ClientResponse response, RequestData requestData) {
 		this(response.statusCode(), response.headers().asHttpHeaders(), response.cookies(), requestData);
 	}
 
+	// Done this way to maintain backwards compatibility while allowing switching to raw
+	// HTTPStatus
+	// Will be removed in `4.x`
+	/**
+	 * @deprecated for removal; new constructors will be added in 4.x
+	 */
+	@Deprecated
+	public ResponseData(RequestData requestData, ClientResponse response) {
+		this(response.headers().asHttpHeaders(), response.cookies(), requestData, response.rawStatusCode());
+	}
+
+	/**
+	 * @deprecated for removal; new constructors will be added in 4.x
+	 */
+	@Deprecated
 	public ResponseData(ServerHttpResponse response, RequestData requestData) {
 		this(response.getStatusCode(), response.getHeaders(), response.getCookies(), requestData);
 	}
 
+	// Done this way to maintain backwards compatibility while allowing switching to raw
+	// HTTPStatus
+	// Will be removed in `4.x`
+	/**
+	 * @deprecated for removal; new constructors will be added in 4.x
+	 */
+	@Deprecated
+	public ResponseData(RequestData requestData, ServerHttpResponse response) {
+		this(response.getHeaders(), response.getCookies(), requestData, response.getRawStatusCode());
+	}
+
+	/**
+	 * @deprecated for removal; new constructors will be added in 4.x
+	 */
+	@Deprecated
 	public ResponseData(ClientHttpResponse clientHttpResponse, RequestData requestData) throws IOException {
 		this(clientHttpResponse.getStatusCode(), clientHttpResponse.getHeaders(),
 				buildCookiesFromHeaders(clientHttpResponse.getHeaders()), requestData);
 	}
 
+	// Done this way to maintain backwards compatibility while allowing switching to raw
+	// HTTPStatus
+	// Will be removed in `4.x`
+	/**
+	 * @deprecated for removal; new constructors will be added in 4.x
+	 */
+	@Deprecated
+	public ResponseData(RequestData requestData, ClientHttpResponse clientHttpResponse) throws IOException {
+		this(clientHttpResponse.getHeaders(), buildCookiesFromHeaders(clientHttpResponse.getHeaders()), requestData,
+				clientHttpResponse.getRawStatusCode());
+	}
+
 	public HttpStatus getHttpStatus() {
 		return httpStatus;
 	}
@@ -85,6 +151,10 @@ public RequestData getRequestData() {
 		return requestData;
 	}
 
+	public Integer getRawHttpStatus() {
+		return rawHttpStatus;
+	}
+
 	@Override
 	public String toString() {
 		ToStringCreator to = new ToStringCreator(this);
@@ -113,7 +183,7 @@ static MultiValueMap<String, ResponseCookie> buildCookiesFromHeaders(HttpHeaders
 
 	@Override
 	public int hashCode() {
-		return Objects.hash(httpStatus, headers, cookies, requestData);
+		return Objects.hash(httpStatus, headers, cookies, requestData, rawHttpStatus);
 	}
 
 	@Override
@@ -126,7 +196,8 @@ public boolean equals(Object o) {
 		}
 		ResponseData that = (ResponseData) o;
 		return httpStatus == that.httpStatus && Objects.equals(headers, that.headers)
-				&& Objects.equals(cookies, that.cookies) && Objects.equals(requestData, that.requestData);
+				&& Objects.equals(cookies, that.cookies) && Objects.equals(requestData, that.requestData)
+				&& Objects.equals(rawHttpStatus, that.rawHttpStatus);
 	}
 
 }

spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/ReactorLoadBalancerExchangeFilterFunction.java

@@ -128,8 +128,8 @@ public Mono<ClientResponse> filter(ClientRequest clientRequest, ExchangeFunction
 				LOG.debug(String.format("LoadBalancer has retrieved the instance for service %s: %s", serviceId,
 						instance.getUri()));
 			}
-			LoadBalancerProperties.StickySession stickySessionProperties = loadBalancerFactory.getProperties(serviceId)
-					.getStickySession();
+			LoadBalancerProperties properties = loadBalancerFactory.getProperties(serviceId);
+			LoadBalancerProperties.StickySession stickySessionProperties = properties.getStickySession();
 			ClientRequest newRequest = buildClientRequest(clientRequest, instance,
 					stickySessionProperties.getInstanceIdCookieName(),
 					stickySessionProperties.isAddServiceInstanceCookie(), transformers);
@@ -140,10 +140,19 @@ public Mono<ClientResponse> filter(ClientRequest clientRequest, ExchangeFunction
 									CompletionContext.Status.FAILED, throwable, lbRequest, lbResponse))))
 					.doOnSuccess(clientResponse -> supportedLifecycleProcessors.forEach(
 							lifecycle -> lifecycle.onComplete(new CompletionContext<>(CompletionContext.Status.SUCCESS,
-									lbRequest, lbResponse, new ResponseData(clientResponse, requestData)))));
+									lbRequest, lbResponse, buildResponseData(requestData, clientResponse,
+											properties.isUseRawStatusCodeInResponseData())))));
 		});
 	}
 
+	private ResponseData buildResponseData(RequestData requestData, ClientResponse clientResponse,
+			boolean useRawStatusCodes) {
+		if (useRawStatusCodes) {
+			return new ResponseData(requestData, clientResponse);
+		}
+		return new ResponseData(clientResponse, requestData);
+	}
+
 	protected Mono<Response<ServiceInstance>> choose(String serviceId, Request<RequestDataContext> request) {
 		ReactiveLoadBalancer<ServiceInstance> loadBalancer = loadBalancerFactory.getInstance(serviceId);
 		if (loadBalancer == null) {

spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/RetryableLoadBalancerExchangeFilterFunction.java

@@ -164,9 +164,11 @@ public Mono<ClientResponse> filter(ClientRequest clientRequest, ExchangeFunction
 					.doOnError(throwable -> supportedLifecycleProcessors.forEach(lifecycle -> lifecycle
 							.onComplete(new CompletionContext<ResponseData, ServiceInstance, RetryableRequestContext>(
 									CompletionContext.Status.FAILED, throwable, lbRequest, lbResponse))))
-					.doOnSuccess(clientResponse -> supportedLifecycleProcessors.forEach(
-							lifecycle -> lifecycle.onComplete(new CompletionContext<>(CompletionContext.Status.SUCCESS,
-									lbRequest, lbResponse, new ResponseData(clientResponse, requestData)))))
+					.doOnSuccess(
+							clientResponse -> supportedLifecycleProcessors.forEach(lifecycle -> lifecycle.onComplete(
+									new CompletionContext<>(CompletionContext.Status.SUCCESS, lbRequest, lbResponse,
+											buildResponseData(requestData, clientResponse,
+													properties.isUseRawStatusCodeInResponseData())))))
 					.map(clientResponse -> {
 						loadBalancerRetryContext.setClientResponse(clientResponse);
 						if (shouldRetrySameServiceInstance(retryPolicy, loadBalancerRetryContext)) {
@@ -192,6 +194,14 @@ lbRequest, lbResponse, new ResponseData(clientResponse, requestData)))))
 		}).retryWhen(exchangeRetry)).retryWhen(filterRetry);
 	}
 
+	private ResponseData buildResponseData(RequestData requestData, ClientResponse clientResponse,
+			boolean useRawStatusCodes) {
+		if (useRawStatusCodes) {
+			return new ResponseData(requestData, clientResponse);
+		}
+		return new ResponseData(clientResponse, requestData);
+	}
+
 	private Retry buildRetrySpec(int max, boolean transientErrors, LoadBalancerProperties.Retry retry) {
 		if (!retry.isEnabled()) {
 			return Retry.max(0).filter(this::isRetryException).transientErrors(transientErrors);

spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/blocking/client/BlockingLoadBalancerClient.java

@@ -96,7 +96,8 @@ public <T> T execute(String serviceId, ServiceInstance serviceInstance, LoadBala
 				.forEach(lifecycle -> lifecycle.onStartRequest(lbRequest, new DefaultResponse(serviceInstance)));
 		try {
 			T response = request.apply(serviceInstance);
-			Object clientResponse = getClientResponse(response);
+			LoadBalancerProperties properties = loadBalancerClientFactory.getProperties(serviceId);
+			Object clientResponse = getClientResponse(response, properties.isUseRawStatusCodeInResponseData());
 			supportedLifecycleProcessors
 					.forEach(lifecycle -> lifecycle.onComplete(new CompletionContext<>(CompletionContext.Status.SUCCESS,
 							lbRequest, defaultResponse, clientResponse)));
@@ -115,13 +116,16 @@ public <T> T execute(String serviceId, ServiceInstance serviceInstance, LoadBala
 		return null;
 	}
 
-	private <T> Object getClientResponse(T response) {
+	private <T> Object getClientResponse(T response, boolean useRawStatusCodes) {
 		ClientHttpResponse clientHttpResponse = null;
 		if (response instanceof ClientHttpResponse) {
 			clientHttpResponse = (ClientHttpResponse) response;
 		}
 		if (clientHttpResponse != null) {
 			try {
+				if (useRawStatusCodes) {
+					return new ResponseData(null, clientHttpResponse);
+				}
 				return new ResponseData(clientHttpResponse, null);
 			}
 			catch (IOException ignored) {