Expose request attributes for circuit breaker async invocation (#510)

spring-cloud-openfeign-core/src/main/java/org/springframework/cloud/openfeign/FeignCircuitBreakerInvocationHandler.java

@@ -30,6 +30,8 @@ import feign.Target;
 
 import org.springframework.cloud.client.circuitbreaker.CircuitBreaker;
 import org.springframework.cloud.client.circuitbreaker.CircuitBreakerFactory;
+import org.springframework.web.context.request.RequestAttributes;
+import org.springframework.web.context.request.RequestContextHolder;
 
 import static feign.Util.checkNotNull;
 
@@ -92,8 +94,10 @@ class FeignCircuitBreakerInvocationHandler implements InvocationHandler {
 	}
 
 	private Supplier<Object> asSupplier(final Method method, final Object[] args) {
+		final RequestAttributes requestAttributes = RequestContextHolder.getRequestAttributes();
 		return () -> {
 			try {
+				RequestContextHolder.setRequestAttributes(requestAttributes);
 				return this.dispatch.get(method).invoke(args);
 			}
 			catch (RuntimeException throwable) {
@@ -102,6 +106,9 @@ class FeignCircuitBreakerInvocationHandler implements InvocationHandler {
 			catch (Throwable throwable) {
 				throw new RuntimeException(throwable);
 			}
+			finally {
+				RequestContextHolder.resetRequestAttributes();
+			}
 		};
 	}
 

spring-cloud-openfeign-core/src/test/java/org/springframework/cloud/openfeign/circuitbreaker/AsyncCircuitBreaker.java

@@ -0,0 +1,59 @@
+/*
+ * Copyright 2013-2021 the original author or authors.
+ *
+ * Licensed 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
+ *
+ *      https://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.springframework.cloud.openfeign.circuitbreaker;
+
+import java.time.Duration;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+import org.springframework.cloud.client.circuitbreaker.CircuitBreaker;
+
+/**
+ * Asynchronous circuit breaker.
+ *
+ * @author John Niang
+ */
+class AsyncCircuitBreaker implements CircuitBreaker {
+
+	final Duration timeout;
+
+	final ExecutorService executorService;
+
+	AsyncCircuitBreaker(Duration timeout) {
+		this(timeout, Executors.newCachedThreadPool());
+	}
+
+	AsyncCircuitBreaker(Duration timeout, ExecutorService executorService) {
+		this.timeout = timeout;
+		this.executorService = executorService;
+	}
+
+	@Override
+	public <T> T run(Supplier<T> toRun, Function<Throwable, T> fallback) {
+		Future<T> future = executorService.submit(toRun::get);
+		try {
+			return future.get(timeout.toMillis(), TimeUnit.MILLISECONDS);
+		}
+		catch (Throwable t) {
+			return fallback.apply(t);
+		}
+	}
+}

spring-cloud-openfeign-core/src/test/java/org/springframework/cloud/openfeign/circuitbreaker/AsyncCircuitBreakerTest.java

@@ -0,0 +1,182 @@
+/*
+ * Copyright 2013-2021 the original author or authors.
+ *
+ * Licensed 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
+ *
+ *      https://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.springframework.cloud.openfeign.circuitbreaker;
+
+import java.time.Duration;
+import java.util.Objects;
+import java.util.UUID;
+import java.util.function.Function;
+
+import javax.servlet.http.HttpServletRequest;
+
+import feign.RequestInterceptor;
+import org.junit.jupiter.api.Test;
+
+import org.springframework.beans.factory.ObjectProvider;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
+import org.springframework.boot.test.autoconfigure.web.servlet.AutoConfigureMockMvc;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.cloud.client.circuitbreaker.CircuitBreaker;
+import org.springframework.cloud.client.circuitbreaker.CircuitBreakerFactory;
+import org.springframework.cloud.client.circuitbreaker.ConfigBuilder;
+import org.springframework.cloud.openfeign.EnableFeignClients;
+import org.springframework.cloud.openfeign.FeignClient;
+import org.springframework.cloud.openfeign.test.NoSecurityConfiguration;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.context.annotation.Import;
+import org.springframework.http.HttpHeaders;
+import org.springframework.test.web.servlet.MockMvc;
+import org.springframework.web.bind.annotation.GetMapping;
+import org.springframework.web.bind.annotation.PathVariable;
+import org.springframework.web.bind.annotation.RestController;
+import org.springframework.web.context.request.RequestContextHolder;
+import org.springframework.web.context.request.ServletRequestAttributes;
+
+import static org.springframework.boot.test.context.SpringBootTest.WebEnvironment.RANDOM_PORT;
+import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get;
+import static org.springframework.test.web.servlet.result.MockMvcResultHandlers.print;
+import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.content;
+import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status;
+
+/**
+ * Tests for asynchronous circuit breaker.
+ *
+ * @author John Niang
+ */
+@SpringBootTest(classes = AsyncCircuitBreakerTest.Application.class,
+		webEnvironment = RANDOM_PORT,
+		properties = "feign.circuitbreaker.enabled=true")
+@AutoConfigureMockMvc
+class AsyncCircuitBreakerTest {
+
+	@Autowired
+	MockMvc mvc;
+
+	@Test
+	void shouldWorkNormally() throws Exception {
+		mvc.perform(get("/hello/proxy"))
+				.andDo(print())
+				.andExpect(status().isOk())
+				.andExpect(content().string("openfeign"));
+	}
+
+	@Test
+	void shouldNotProxyAnyHeadersWithoutHeaderSet() throws Exception {
+		mvc.perform(get("/headers/" + HttpHeaders.AUTHORIZATION + "/proxy"))
+				.andDo(print())
+				.andExpect(status().isOk())
+				.andExpect(content().string(""));
+	}
+
+	@Test
+	void shouldProxyHeaderWhenHeaderSet() throws Exception {
+		String authorization = UUID.randomUUID().toString();
+		mvc.perform(get("/headers/" + HttpHeaders.AUTHORIZATION + "/proxy")
+				.header(HttpHeaders.AUTHORIZATION, authorization))
+				.andDo(print())
+				.andExpect(status().isOk())
+				.andExpect(content().string(authorization));
+	}
+
+	@EnableAutoConfiguration
+	@Configuration(proxyBeanMethods = false)
+	@EnableFeignClients(clients = {TestClient.class})
+	@Import({NoSecurityConfiguration.class, TestController.class})
+	static class Application {
+
+		@Bean
+		CircuitBreakerFactory<Duration, ConfigBuilder<Duration>> circuitBreakerFactory() {
+			return new CircuitBreakerFactory<Duration, ConfigBuilder<Duration>>() {
+
+				Function<String, Duration> defaultConfiguration = id -> Duration.ofMillis(1000);
+
+				@Override
+				public CircuitBreaker create(String id) {
+					Duration timeout = super.getConfigurations().computeIfAbsent(id, defaultConfiguration);
+					return new AsyncCircuitBreaker(timeout);
+				}
+
+				@Override
+				protected ConfigBuilder<Duration> configBuilder(String id) {
+					return () -> Duration.ofMillis(100);
+				}
+
+				@Override
+				public void configureDefault(Function<String, Duration> defaultConfiguration) {
+					this.defaultConfiguration = defaultConfiguration;
+				}
+			};
+		}
+
+		@Bean
+		RequestInterceptor proxyHeaderRequestInterceptor() {
+			return template -> {
+				ServletRequestAttributes requestAttributes =
+						(ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
+				String authorization = Objects.requireNonNull(requestAttributes).getRequest()
+						.getHeader(HttpHeaders.AUTHORIZATION);
+				if (authorization != null) {
+					// proxy authorization header
+					template.header(HttpHeaders.AUTHORIZATION, authorization);
+				}
+			};
+		}
+	}
+
+	@RestController
+	static class TestController {
+
+		final ObjectProvider<TestClient> testClient;
+
+		TestController(ObjectProvider<TestClient> testClient) {
+			this.testClient = testClient;
+		}
+
+		@GetMapping("/hello")
+		String hello() {
+			return "openfeign";
+		}
+
+		@GetMapping("/hello/proxy")
+		String helloProxy() {
+			return testClient.getObject().hello();
+		}
+
+		@GetMapping("/headers/{headerName}")
+		String header(HttpServletRequest request, @PathVariable String headerName) {
+			return request.getHeader(headerName);
+		}
+
+		@GetMapping("/headers/{headerName}/proxy")
+		String headerProxy(@PathVariable String headerName) {
+			return testClient.getObject().header(headerName);
+		}
+	}
+
+	@FeignClient(name = "async-circuit-breaker-test", url = "http://localhost:${local.server.port}")
+	interface TestClient {
+
+		@GetMapping("/hello")
+		String hello();
+
+		@GetMapping("/headers/{headerName}")
+		String header(@PathVariable String headerName);
+
+	}
+}