Refactor ServerWebExchange#getAttribute options

Issue: SPR-15718

spring-test/src/test/java/org/springframework/test/web/reactive/server/MockServerSpecTests.java

@@ -99,7 +99,7 @@ public class MockServerSpecTests {
 		@Override
 		public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
 			String name = "test-attribute";
-			String value = (String) exchange.getAttribute(name).orElse("");
+			String value = exchange.getAttributeOrDefault(name, "");
 			exchange.getAttributes().put(name, value + ":" + this.name);
 			return chain.filter(exchange);
 		}

spring-web/src/main/java/org/springframework/web/server/ServerWebExchange.java

@@ -19,7 +19,6 @@ package org.springframework.web.server;
 import java.security.Principal;
 import java.time.Instant;
 import java.util.Map;
-import java.util.Optional;
 import java.util.function.Consumer;
 
 import reactor.core.publisher.Mono;
@@ -29,6 +28,7 @@ import org.springframework.http.codec.multipart.Part;
 import org.springframework.http.server.reactive.ServerHttpRequest;
 import org.springframework.http.server.reactive.ServerHttpResponse;
 import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
 import org.springframework.util.MultiValueMap;
 import org.springframework.web.server.i18n.LocaleContextResolver;
 
@@ -63,7 +63,37 @@ public interface ServerWebExchange {
 	 * @param <T> the attribute type
 	 * @return the attribute value
 	 */
-	<T> Optional<T> getAttribute(String name);
+	@SuppressWarnings("unchecked")
+	@Nullable
+	default <T> T getAttribute(String name) {
+		return (T) getAttributes().get(name);
+	}
+
+	/**
+	 * Return the request attribute value or if not present raise an
+	 * {@link IllegalArgumentException}.
+	 * @param name the attribute name
+	 * @param <T> the attribute type
+	 * @return the attribute value
+	 */
+	@SuppressWarnings("unchecked")
+	default <T> T getRequiredAttribute(String name) {
+		T value = getAttribute(name);
+		Assert.notNull(value, "Required attribute '" + name + "' is missing.");
+		return value;
+	}
+
+	/**
+	 * Return the request attribute value, or a default, fallback value.
+	 * @param name the attribute name
+	 * @param defaultValue a default value to return instead
+	 * @param <T> the attribute type
+	 * @return the attribute value
+	 */
+	@SuppressWarnings("unchecked")
+	default <T> T getAttributeOrDefault(String name, T defaultValue) {
+		return (T) getAttributes().getOrDefault(name, defaultValue);
+	}
 
 	/**
 	 * Return the web session for the current request. Always guaranteed  to

spring-web/src/main/java/org/springframework/web/server/ServerWebExchangeDecorator.java

@@ -18,7 +18,6 @@ package org.springframework.web.server;
 import java.security.Principal;
 import java.time.Instant;
 import java.util.Map;
-import java.util.Optional;
 
 import reactor.core.publisher.Mono;
 
@@ -76,11 +75,6 @@ public class ServerWebExchangeDecorator implements ServerWebExchange {
 		return getDelegate().getAttributes();
 	}
 
-	@Override
-	public <T> Optional<T> getAttribute(String name) {
-		return getDelegate().getAttribute(name);
-	}
-
 	@Override
 	public Mono<WebSession> getSession() {
 		return getDelegate().getSession();

spring-web/src/main/java/org/springframework/web/server/adapter/DefaultServerWebExchange.java

@@ -23,7 +23,6 @@ import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 
 import reactor.core.publisher.Mono;
@@ -46,9 +45,9 @@ import org.springframework.util.CollectionUtils;
 import org.springframework.util.LinkedMultiValueMap;
 import org.springframework.util.MultiValueMap;
 import org.springframework.util.StringUtils;
-import org.springframework.web.server.i18n.LocaleContextResolver;
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.server.WebSession;
+import org.springframework.web.server.i18n.LocaleContextResolver;
 import org.springframework.web.server.session.WebSessionManager;
 
 /**
@@ -178,11 +177,6 @@ public class DefaultServerWebExchange implements ServerWebExchange {
 		return this.attributes;
 	}
 
-	@Override @SuppressWarnings("unchecked")
-	public <T> Optional<T> getAttribute(String name) {
-		return Optional.ofNullable((T) this.attributes.get(name));
-	}
-
 	@Override
 	public Mono<WebSession> getSession() {
 		return this.sessionMono;

spring-web/src/test/java/org/springframework/web/server/adapter/WebHttpHandlerBuilderTests.java

@@ -117,7 +117,8 @@ public class WebHttpHandlerBuilderTests {
 
 		private WebFilter createFilter(String name) {
 			return (exchange, chain) -> {
-				String value = exchange.getAttribute(ATTRIBUTE).map(v -> v + "::" + name).orElse(name);
+				String value = exchange.getAttribute(ATTRIBUTE);
+				value = (value != null ? value + "::" + name : name);
 				exchange.getAttributes().put(ATTRIBUTE, value);
 				return chain.filter(exchange);
 			};
@@ -126,7 +127,7 @@ public class WebHttpHandlerBuilderTests {
 		@Bean
 		public WebHandler webHandler() {
 			return exchange -> {
-				String value = exchange.getAttribute(ATTRIBUTE).map(v -> (String) v).orElse("none");
+				String value = exchange.getAttributeOrDefault(ATTRIBUTE, "none");
 				return writeToResponse(exchange, value);
 			};
 		}

spring-webflux/src/main/java/org/springframework/web/reactive/function/server/DefaultServerRequest.java

@@ -134,7 +134,7 @@ class DefaultServerRequest implements ServerRequest {
 
 	@Override
 	public <T> Optional<T> attribute(String name) {
-		return this.exchange.getAttribute(name);
+		return Optional.ofNullable(this.exchange.getAttribute(name));
 	}
 
 	@Override
@@ -150,8 +150,8 @@ class DefaultServerRequest implements ServerRequest {
 
 	@Override
 	public Map<String, String> pathVariables() {
-		return this.exchange.<Map<String, String>>getAttribute(RouterFunctions.URI_TEMPLATE_VARIABLES_ATTRIBUTE).
-				orElseGet(Collections::emptyMap);
+		return this.exchange.getAttributeOrDefault(
+				RouterFunctions.URI_TEMPLATE_VARIABLES_ATTRIBUTE, Collections.emptyMap());
 	}
 
 	@Override

spring-webflux/src/main/java/org/springframework/web/reactive/function/server/support/HandlerFunctionAdapter.java

@@ -57,11 +57,7 @@ public class HandlerFunctionAdapter implements HandlerAdapter {
 	@Override
 	public Mono<HandlerResult> handle(ServerWebExchange exchange, Object handler) {
 		HandlerFunction<?> handlerFunction = (HandlerFunction<?>) handler;
-		ServerRequest request =
-				exchange.<ServerRequest>getAttribute(RouterFunctions.REQUEST_ATTRIBUTE)
-						.orElseThrow(() -> new IllegalStateException(
-								"Could not find ServerRequest in exchange attributes"));
-
+		ServerRequest request = exchange.getRequiredAttribute(RouterFunctions.REQUEST_ATTRIBUTE);
 		return handlerFunction.handle(request)
 				.map(response -> new HandlerResult(handlerFunction, response, HANDLER_FUNCTION_RETURN_TYPE));
 	}

spring-webflux/src/main/java/org/springframework/web/reactive/resource/ResourceWebHandler.java

@@ -24,7 +24,6 @@ import java.util.ArrayList;
 import java.util.Collections;
 import java.util.EnumSet;
 import java.util.List;
-import java.util.Optional;
 import java.util.Set;
 
 import org.apache.commons.logging.Log;
@@ -311,14 +310,9 @@ public class ResourceWebHandler implements WebHandler, InitializingBean {
 
 	protected Mono<Resource> getResource(ServerWebExchange exchange) {
 
-		String attributeName = HandlerMapping.PATH_WITHIN_HANDLER_MAPPING_ATTRIBUTE;
-		Optional<String> optional = exchange.getAttribute(attributeName);
-		if (!optional.isPresent()) {
-			return Mono.error(new IllegalStateException(
-					"Required request attribute '" + attributeName + "' is not set"));
-		}
-
-		String path = processPath(optional.get());
+		String name = HandlerMapping.PATH_WITHIN_HANDLER_MAPPING_ATTRIBUTE;
+		String pathWithinHandler = exchange.getRequiredAttribute(name);
+		String path = processPath(pathWithinHandler);
 		if (!StringUtils.hasText(path) || isInvalidPath(path)) {
 			if (logger.isTraceEnabled()) {
 				logger.trace("Ignoring invalid resource path [" + path + "]");

spring-webflux/src/main/java/org/springframework/web/reactive/result/HandlerResultHandlerSupport.java

@@ -152,9 +152,8 @@ public abstract class HandlerResultHandlerSupport implements Ordered {
 	private List<MediaType> getProducibleTypes(ServerWebExchange exchange,
 			Supplier<List<MediaType>> producibleTypesSupplier) {
 
-		return exchange.getAttribute(HandlerMapping.PRODUCIBLE_MEDIA_TYPES_ATTRIBUTE)
-				.map(attribute -> (List<MediaType>) new ArrayList<>((Set<MediaType>) attribute))
-				.orElseGet(producibleTypesSupplier);
+		Set<MediaType> mediaTypes = exchange.getAttribute(HandlerMapping.PRODUCIBLE_MEDIA_TYPES_ATTRIBUTE);
+		return (mediaTypes != null ? new ArrayList<>(mediaTypes) : producibleTypesSupplier.get());
 	}
 
 	private MediaType selectMoreSpecificMediaType(MediaType acceptable, MediaType producible) {

spring-webflux/src/main/java/org/springframework/web/reactive/result/method/annotation/AbstractNamedValueSyncArgumentResolver.java

@@ -81,6 +81,7 @@ public abstract class AbstractNamedValueSyncArgumentResolver extends AbstractNam
 	/**
 	 * Actually resolve the value synchronously.
 	 */
-	protected abstract Optional<Object> resolveNamedValue(String name, MethodParameter param, ServerWebExchange exchange);
+	@Nullable
+	protected abstract Object resolveNamedValue(String name, MethodParameter param, ServerWebExchange exchange);
 
 }

spring-webflux/src/main/java/org/springframework/web/reactive/result/method/annotation/CookieValueMethodArgumentResolver.java

@@ -16,8 +16,6 @@
 
 package org.springframework.web.reactive.result.method.annotation;
 
-import java.util.Optional;
-
 import org.springframework.beans.factory.config.ConfigurableBeanFactory;
 import org.springframework.core.MethodParameter;
 import org.springframework.core.ReactiveAdapterRegistry;
@@ -64,18 +62,13 @@ public class CookieValueMethodArgumentResolver extends AbstractNamedValueSyncArg
 	}
 
 	@Override
-	protected Optional<Object> resolveNamedValue(String name, MethodParameter parameter, ServerWebExchange exchange) {
+	protected Object resolveNamedValue(String name, MethodParameter parameter, ServerWebExchange exchange) {
 		HttpCookie cookie = exchange.getRequest().getCookies().getFirst(name);
 		Class<?> paramType = parameter.getNestedParameterType();
 		if (HttpCookie.class.isAssignableFrom(paramType)) {
-			return Optional.ofNullable(cookie);
-		}
-		else if (cookie != null) {
-			return Optional.of(cookie.getValue());
-		}
-		else {
-			return Optional.empty();
+			return cookie;
 		}
+		return (cookie != null ? cookie.getValue() : null);
 	}
 
 	@Override

spring-webflux/src/main/java/org/springframework/web/reactive/result/method/annotation/ExpressionValueMethodArgumentResolver.java

@@ -16,8 +16,6 @@
 
 package org.springframework.web.reactive.result.method.annotation;
 
-import java.util.Optional;
-
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.beans.factory.config.ConfigurableBeanFactory;
 import org.springframework.core.MethodParameter;
@@ -62,9 +60,9 @@ public class ExpressionValueMethodArgumentResolver extends AbstractNamedValueSyn
 	}
 
 	@Override
-	protected Optional<Object> resolveNamedValue(String name, MethodParameter parameter, ServerWebExchange exchange) {
+	protected Object resolveNamedValue(String name, MethodParameter parameter, ServerWebExchange exchange) {
 		// No name to resolve
-		return Optional.empty();
+		return null;
 	}
 
 	@Override

spring-webflux/src/main/java/org/springframework/web/reactive/result/method/annotation/PathVariableMapMethodArgumentResolver.java

@@ -63,7 +63,7 @@ public class PathVariableMapMethodArgumentResolver extends HandlerMethodArgument
 			BindingContext context, ServerWebExchange exchange) {
 
 		String name = HandlerMapping.URI_TEMPLATE_VARIABLES_ATTRIBUTE;
-		Object value = exchange.getAttribute(name).orElse(Collections.emptyMap());
+		Object value = exchange.getAttributeOrDefault(name, Collections.emptyMap());
 		return Optional.of(value);
 	}
 

spring-webflux/src/main/java/org/springframework/web/reactive/result/method/annotation/PathVariableMethodArgumentResolver.java

@@ -16,8 +16,8 @@
 
 package org.springframework.web.reactive.result.method.annotation;
 
+import java.util.Collections;
 import java.util.Map;
-import java.util.Optional;
 
 import org.springframework.beans.factory.config.ConfigurableBeanFactory;
 import org.springframework.core.MethodParameter;
@@ -83,10 +83,9 @@ public class PathVariableMethodArgumentResolver extends AbstractNamedValueSyncAr
 
 	@Override
 	@SuppressWarnings("unchecked")
-	protected Optional<Object> resolveNamedValue(String name, MethodParameter parameter, ServerWebExchange exchange) {
+	protected Object resolveNamedValue(String name, MethodParameter parameter, ServerWebExchange exchange) {
 		String attributeName = HandlerMapping.URI_TEMPLATE_VARIABLES_ATTRIBUTE;
-		return exchange.getAttribute(attributeName)
-				.map(value -> ((Map<String, String>) value).get(name));
+		return exchange.getAttributeOrDefault(attributeName, Collections.emptyMap()).get(name);
 	}
 
 	@Override

spring-webflux/src/main/java/org/springframework/web/reactive/result/method/annotation/RequestAttributeMethodArgumentResolver.java

@@ -16,8 +16,6 @@
 
 package org.springframework.web.reactive.result.method.annotation;
 
-import java.util.Optional;
-
 import org.springframework.beans.factory.config.ConfigurableBeanFactory;
 import org.springframework.core.MethodParameter;
 import org.springframework.core.ReactiveAdapterRegistry;
@@ -63,7 +61,7 @@ public class RequestAttributeMethodArgumentResolver extends AbstractNamedValueSy
 	}
 
 	@Override
-	protected Optional<Object> resolveNamedValue(String name, MethodParameter parameter, ServerWebExchange exchange) {
+	protected Object resolveNamedValue(String name, MethodParameter parameter, ServerWebExchange exchange) {
 		return exchange.getAttribute(name);
 	}
 

spring-webflux/src/main/java/org/springframework/web/reactive/result/method/annotation/RequestHeaderMethodArgumentResolver.java

@@ -76,15 +76,13 @@ public class RequestHeaderMethodArgumentResolver extends AbstractNamedValueSyncA
 	}
 
 	@Override
-	protected Optional<Object> resolveNamedValue(String name, MethodParameter parameter,
-			ServerWebExchange exchange) {
-
+	protected Object resolveNamedValue(String name, MethodParameter parameter, ServerWebExchange exchange) {
 		List<String> headerValues = exchange.getRequest().getHeaders().get(name);
 		Object result = null;
 		if (headerValues != null) {
 			result = (headerValues.size() == 1 ? headerValues.get(0) : headerValues);
 		}
-		return Optional.ofNullable(result);
+		return result;
 	}
 
 	@Override

spring-webflux/src/main/java/org/springframework/web/reactive/result/method/annotation/RequestParamMethodArgumentResolver.java

@@ -18,7 +18,6 @@ package org.springframework.web.reactive.result.method.annotation;
 
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 
 import org.springframework.beans.BeanUtils;
 import org.springframework.beans.factory.config.ConfigurableBeanFactory;
@@ -99,15 +98,13 @@ public class RequestParamMethodArgumentResolver extends AbstractNamedValueSyncAr
 	}
 
 	@Override
-	protected Optional<Object> resolveNamedValue(String name, MethodParameter parameter,
-			ServerWebExchange exchange) {
-
+	protected Object resolveNamedValue(String name, MethodParameter parameter, ServerWebExchange exchange) {
 		List<String> paramValues = exchange.getRequest().getQueryParams().get(name);
 		Object result = null;
 		if (paramValues != null) {
 			result = (paramValues.size() == 1 ? paramValues.get(0) : paramValues);
 		}
-		return Optional.ofNullable(result);
+		return result;
 	}
 
 	@Override

spring-webflux/src/main/java/org/springframework/web/reactive/result/view/RedirectView.java

@@ -227,7 +227,7 @@ public class RedirectView extends AbstractUrlBasedView {
 	@SuppressWarnings("unchecked")
 	private Map<String, String> getCurrentUriVariables(ServerWebExchange exchange) {
 		String name = HandlerMapping.URI_TEMPLATE_VARIABLES_ATTRIBUTE;
-		return (Map<String, String>) exchange.getAttribute(name).orElse(Collections.emptyMap());
+		return exchange.getAttributeOrDefault(name, Collections.<String, String>emptyMap());
 	}
 
 	/**

spring-webflux/src/main/java/org/springframework/web/reactive/result/view/RequestContext.java

@@ -404,7 +404,7 @@ public class RequestContext {
 	protected <T> T getModelObject(String modelName) {
 		T modelObject = (T) this.model.get(modelName);
 		if (modelObject == null) {
-			modelObject = (T) this.exchange.getAttribute(modelName).orElse(null);
+			modelObject = this.exchange.getAttribute(modelName);
 		}
 		return modelObject;
 	}

spring-webflux/src/test/java/org/springframework/web/reactive/handler/SimpleUrlHandlerMappingTests.java

@@ -103,7 +103,7 @@ public class SimpleUrlHandlerMappingTests {
 			assertNotNull(actual);
 			assertSame(bean, actual);
 			//noinspection OptionalGetWithoutIsPresent
-			assertEquals(pathWithinMapping, exchange.getAttribute(PATH_WITHIN_HANDLER_MAPPING_ATTRIBUTE).get());
+			assertEquals(pathWithinMapping, exchange.getAttribute(PATH_WITHIN_HANDLER_MAPPING_ATTRIBUTE));
 		}
 		else {
 			assertNull(actual);

spring-webflux/src/test/java/org/springframework/web/reactive/resource/ResourceWebHandlerTests.java

@@ -388,7 +388,7 @@ public class ResourceWebHandlerTests {
 		assertEquals(HttpStatus.NOT_FOUND, exchange.getResponse().getStatusCode());
 	}
 
-	@Test(expected = IllegalStateException.class)
+	@Test(expected = IllegalArgumentException.class)
 	public void noPathWithinHandlerMappingAttribute() throws Exception {
 		MockServerWebExchange exchange = MockServerHttpRequest.get("").toExchange();
 		this.handler.handle(exchange).block(TIMEOUT);