Spring/spring-graphql
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Propagate context in WebMvc GraphQlWebSocketHandler

Browse Source 
See gh-342
This commit is contained in:
rstoyanchev
2022-04-12 10:07:46 +01:00
parent e2948e9a36
commit 85dc2e5145
9 changed files with 196 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
spring-graphql/src/main/java/org/springframework/graphql/execution/CompositeThreadLocalAccessor.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2002-2021 the original author or authors.
* Copyright 2002-2022 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.
@@ -16,21 +16,24 @@
package org.springframework.graphql.execution;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
/**
* Default implementation of a composite accessor that is returned from
* {@link ThreadLocalAccessor#composite(List)}.
*
* @author Rossen Stoyanchev
* @since 1.0.0
*/
class CompositeThreadLocalAccessor implements ThreadLocalAccessor {
private final List<ThreadLocalAccessor> accessors;
CompositeThreadLocalAccessor(List<ThreadLocalAccessor> accessors) {
this.accessors = accessors;
this.accessors = new ArrayList<>(accessors);
}
@Override

24
spring-graphql/src/main/java/org/springframework/graphql/server/DefaultWebGraphQlHandlerBuilder.java
View File

@@ -96,26 +96,34 @@ class DefaultWebGraphQlHandlerBuilder implements WebGraphQlHandler.Builder {
.map(interceptor -> interceptor.apply(endOfChain))
.orElse(endOfChain);
ThreadLocalAccessor accessor = (CollectionUtils.isEmpty(this.accessors) ? null :
ThreadLocalAccessor.composite(this.accessors));
return new WebGraphQlHandler() {
@Override
public WebSocketGraphQlInterceptor getWebSocketInterceptor() {
return (webSocketInterceptor != null ?
webSocketInterceptor : new WebSocketGraphQlInterceptor() {});
}
@Nullable
@Override
public ThreadLocalAccessor getThreadLocalAccessor() {
return accessor;
}
@Override
public Mono<WebGraphQlResponse> handleRequest(WebGraphQlRequest request) {
return executionChain.next(request)
.contextWrite(context -> {
if (!CollectionUtils.isEmpty(accessors)) {
ThreadLocalAccessor accessor = ThreadLocalAccessor.composite(accessors);
if (accessor != null) {
return ReactorContextManager.extractThreadLocalValues(accessor, context);
}
return context;
});
}
@Override
public WebSocketGraphQlInterceptor webSocketInterceptor() {
return (webSocketInterceptor != null ?
webSocketInterceptor : new WebSocketGraphQlInterceptor() {});
}
};
}

21
spring-graphql/src/main/java/org/springframework/graphql/server/WebGraphQlHandler.java
View File

@@ -22,6 +22,7 @@ import reactor.core.publisher.Mono;
import org.springframework.graphql.ExecutionGraphQlService;
import org.springframework.graphql.execution.ThreadLocalAccessor;
import org.springframework.lang.Nullable;
/**
@@ -33,6 +34,19 @@ import org.springframework.graphql.execution.ThreadLocalAccessor;
*/
public interface WebGraphQlHandler {
/**
* Return the single interceptor of type
* {@link WebSocketGraphQlInterceptor} among all the configured
* interceptors.
*/
WebSocketGraphQlInterceptor getWebSocketInterceptor();
/**
* Return the composite {@link ThreadLocalAccessor} that the handler is
* configured with.
*/
@Nullable
ThreadLocalAccessor getThreadLocalAccessor();
/**
* Execute the given request and return the response.
@@ -41,13 +55,6 @@ public interface WebGraphQlHandler {
*/
Mono<WebGraphQlResponse> handleRequest(WebGraphQlRequest request);
/**
* Return the single interceptor of type
* {@link WebSocketGraphQlInterceptor} among all the configured
* interceptors.
*/
WebSocketGraphQlInterceptor webSocketInterceptor();
/**
* Provides access to a builder to create a {@link WebGraphQlHandler} instance.

2
spring-graphql/src/main/java/org/springframework/graphql/server/webflux/GraphQlWebSocketHandler.java
View File

@@ -83,7 +83,7 @@ public class GraphQlWebSocketHandler implements WebSocketHandler {
Assert.notNull(graphQlHandler, "WebGraphQlHandler is required");
this.graphQlHandler = graphQlHandler;
this.webSocketInterceptor = this.graphQlHandler.webSocketInterceptor();
this.webSocketInterceptor = this.graphQlHandler.getWebSocketInterceptor();
this.codecDelegate = new CodecDelegate(codecConfigurer);
this.initTimeoutDuration = connectionInitTimeout;
}

92
spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java
View File

@@ -18,6 +18,7 @@ package org.springframework.graphql.server.webmvc;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.Closeable;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
@@ -25,6 +26,7 @@ import java.net.URI;
import java.time.Duration;
import java.util.Arrays;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
@@ -43,6 +45,7 @@ import reactor.core.publisher.Mono;
import reactor.core.scheduler.Scheduler;
import reactor.core.scheduler.Schedulers;
import org.springframework.graphql.execution.ThreadLocalAccessor;
import org.springframework.graphql.server.WebGraphQlHandler;
import org.springframework.graphql.server.WebGraphQlRequest;
import org.springframework.graphql.server.WebGraphQlResponse;
@@ -53,15 +56,21 @@ import org.springframework.http.HttpInputMessage;
import org.springframework.http.HttpOutputMessage;
import org.springframework.http.converter.GenericHttpMessageConverter;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.SubProtocolCapable;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.ExceptionWebSocketHandlerDecorator;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.web.socket.server.HandshakeHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;
import org.springframework.web.socket.server.support.WebSocketHttpRequestHandler;
/**
* WebSocketHandler for GraphQL based on
@@ -81,7 +90,9 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
private final WebGraphQlHandler graphQlHandler;
private final WebSocketGraphQlInterceptor webSocketInterceptor;
private final ContextHandshakeInterceptor contextHandshakeInterceptor;
private final WebSocketGraphQlInterceptor webSocketGraphQlInterceptor;
private final Duration initTimeoutDuration;
@@ -103,7 +114,8 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
Assert.notNull(converter, "HttpMessageConverter for JSON is required");
this.graphQlHandler = graphQlHandler;
this.webSocketInterceptor = this.graphQlHandler.webSocketInterceptor();
this.contextHandshakeInterceptor = new ContextHandshakeInterceptor(graphQlHandler.getThreadLocalAccessor());
this.webSocketGraphQlInterceptor = this.graphQlHandler.getWebSocketInterceptor();
this.initTimeoutDuration = connectionInitTimeout;
this.converter = converter;
}
@@ -113,6 +125,18 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
return SUB_PROTOCOL_LIST;
}
/**
* Return a {@link WebSocketHttpRequestHandler} that uses this instance as
* its {@link WebGraphQlHandler} and adds a {@link HandshakeInterceptor} to
* propagate context.
*/
public WebSocketHttpRequestHandler asWebSocketHttpRequestHandler(HandshakeHandler handshakeHandler) {
WebSocketHttpRequestHandler handler = new WebSocketHttpRequestHandler(this, handshakeHandler);
handler.setHandshakeInterceptors(Collections.singletonList(this.contextHandshakeInterceptor));
return handler;
}
@Override
public void afterConnectionEstablished(WebSocketSession session) {
if ("graphql-ws".equalsIgnoreCase(session.getAcceptedProtocol())) {
@@ -137,8 +161,15 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
}
@SuppressWarnings({"unused", "try"})
@Override
protected void handleTextMessage(WebSocketSession session, TextMessage webSocketMessage) throws Exception {
try (Closeable closeable = this.contextHandshakeInterceptor.restoreThreadLocalValue(session)) {
handleInternal(session, webSocketMessage);
}
}
private void handleInternal(WebSocketSession session, TextMessage webSocketMessage) throws IOException {
GraphQlWebSocketMessage message = decode(webSocketMessage);
String id = message.getId();
Map<String, Object> payload = message.getPayload();
@@ -174,7 +205,7 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
if (subscription != null) {
subscription.cancel();
}
this.webSocketInterceptor.handleCancelledSubscription(session.getId(), id)
this.webSocketGraphQlInterceptor.handleCancelledSubscription(session.getId(), id)
.block(Duration.ofSeconds(10));
}
return;
@@ -183,7 +214,7 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
GraphQlStatus.closeSession(session, GraphQlStatus.TOO_MANY_INIT_REQUESTS_STATUS);
return;
}
this.webSocketInterceptor.handleConnectionInitialization(session.getId(), payload)
this.webSocketGraphQlInterceptor.handleConnectionInitialization(session.getId(), payload)
.defaultIfEmpty(Collections.emptyMap())
.publishOn(sessionState.getScheduler()) // Serial blocking send via single thread
.doOnNext(ackPayload -> {
@@ -285,7 +316,7 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
info.dispose();
Map<String, Object> connectionInitPayload = info.getConnectionInitPayload();
if (connectionInitPayload != null) {
this.webSocketInterceptor.handleConnectionClosed(id, closeStatus.getCode(), connectionInitPayload);
this.webSocketGraphQlInterceptor.handleConnectionClosed(id, closeStatus.getCode(), connectionInitPayload);
}
}
}
@@ -296,6 +327,57 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
}
/**
* {@code HandshakeInterceptor} that propagates ThreadLocal context through
* the attributes map in {@code WebSocketSession}.
*/
private static class ContextHandshakeInterceptor implements HandshakeInterceptor {
private static final String SAVED_CONTEXT_KEY = ContextHandshakeInterceptor.class.getName();
@Nullable
private final ThreadLocalAccessor accessor;
ContextHandshakeInterceptor(@Nullable ThreadLocalAccessor accessor) {
this.accessor = accessor;
}
@Override
public boolean beforeHandshake(
ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,
Map<String, Object> attributes) {
if (this.accessor != null) {
Map<String, Object> valuesMap = new LinkedHashMap<>();
this.accessor.extractValues(valuesMap);
attributes.put(SAVED_CONTEXT_KEY, valuesMap);
}
return true;
}
@Override
public void afterHandshake(
ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,
@Nullable Exception exception) {
}
@SuppressWarnings("unchecked")
public Closeable restoreThreadLocalValue(WebSocketSession session) {
if (this.accessor != null) {
Map<String, Object> valuesMap = (Map<String, Object>) session.getAttributes().get(SAVED_CONTEXT_KEY);
// Uncomment when Boot is updated to use HandshakeInterceptor
// Assert.state(valuesMap != null, "No context");
if (valuesMap != null) {
this.accessor.restoreValues(valuesMap);
return () -> this.accessor.resetValues(valuesMap);
}
}
return () -> {};
}
}
private static class GraphQlStatus {
private static final CloseStatus INVALID_MESSAGE_STATUS = new CloseStatus(4400, "Invalid message");

11
spring-graphql/src/test/java/org/springframework/graphql/server/WebSocketHandlerTestSupport.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2002-2021 the original author or authors.
* Copyright 2002-2022 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.
@@ -20,6 +20,8 @@ import reactor.core.publisher.Flux;
import org.springframework.graphql.BookSource;
import org.springframework.graphql.GraphQlSetup;
import org.springframework.graphql.execution.ThreadLocalAccessor;
import org.springframework.lang.Nullable;
public abstract class WebSocketHandlerTestSupport {
@@ -65,6 +67,12 @@ public abstract class WebSocketHandlerTestSupport {
protected WebGraphQlHandler initHandler(WebGraphQlInterceptor... interceptors) {
return initHandler(null, interceptors);
}
protected WebGraphQlHandler initHandler(
@Nullable ThreadLocalAccessor accessor, WebGraphQlInterceptor... interceptors) {
return GraphQlSetup.schemaResource(BookSource.schema)
.queryFetcher("bookById", environment -> {
Long id = Long.parseLong(environment.getArgument("id"));
@@ -75,6 +83,7 @@ public abstract class WebSocketHandlerTestSupport {
return Flux.fromIterable(BookSource.books())
.filter((book) -> book.getAuthor().getFullName().contains(author));
})
.threadLocalAccessor(accessor)
.interceptor(interceptors)
.toWebGraphQlHandler();
}

55
spring-graphql/src/test/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandlerTests.java
View File

@@ -35,18 +35,21 @@ import reactor.core.publisher.Mono;
import reactor.test.StepVerifier;
import org.springframework.graphql.GraphQlSetup;
import org.springframework.graphql.TestThreadLocalAccessor;
import org.springframework.graphql.execution.ThreadLocalAccessor;
import org.springframework.graphql.server.ConsumeOneAndNeverCompleteInterceptor;
import org.springframework.graphql.server.WebGraphQlHandler;
import org.springframework.graphql.server.WebGraphQlInterceptor;
import org.springframework.graphql.server.WebSocketGraphQlInterceptor;
import org.springframework.graphql.server.WebSocketHandlerTestSupport;
import org.springframework.graphql.server.support.GraphQlWebSocketMessage;
import org.springframework.graphql.server.support.GraphQlWebSocketMessageType;
import org.springframework.graphql.server.ConsumeOneAndNeverCompleteInterceptor;
import org.springframework.graphql.server.WebSocketHandlerTestSupport;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpInputMessage;
import org.springframework.http.converter.GenericHttpMessageConverter;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.converter.json.MappingJackson2HttpMessageConverter;
import org.springframework.lang.Nullable;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketMessage;
@@ -371,6 +374,45 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport {
.verify(TIMEOUT);
}
@Test
void contextPropagation() throws Exception {
ThreadLocal<String> threadLocal = new ThreadLocal<>();
threadLocal.set("foo");
WebGraphQlInterceptor threadLocalInterceptor = (request, chain) -> {
assertThat(threadLocal.get()).isEqualTo("foo");
return chain.next(request);
};
GraphQlWebSocketHandler handler = initWebSocketHandler(
new TestThreadLocalAccessor<>(threadLocal), threadLocalInterceptor);
// Use HandshakeInterceptor to capture ThreadLocal context
handler.asWebSocketHttpRequestHandler((request, response, wsHandler, attributes) -> false)
.getHandshakeInterceptors().get(0)
.beforeHandshake(null, null, null, this.session.getAttributes());
// Context should propagate, if message is handled on different thread
Thread thread = new Thread(() -> {
try {
handle(handler,
new TextMessage("{\"type\":\"connection_init\"}"),
new TextMessage(BOOK_QUERY));
}
catch (Exception ex) {
throw new IllegalStateException(ex);
}
});
thread.start();
StepVerifier.create(this.session.getOutput())
.expectNextCount(2)
.consumeNextWith((message) -> assertMessageType(message, GraphQlWebSocketMessageType.COMPLETE))
.then(this.session::close) // Complete output Flux
.expectComplete()
.verify(TIMEOUT);
}
private void handle(GraphQlWebSocketHandler handler, TextMessage... textMessages) throws Exception {
handler.afterConnectionEstablished(this.session);
for (TextMessage message : textMessages) {
@@ -379,8 +421,15 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport {
}
private GraphQlWebSocketHandler initWebSocketHandler(WebGraphQlInterceptor... interceptors) {
return initWebSocketHandler(null, interceptors);
}
private GraphQlWebSocketHandler initWebSocketHandler(
@Nullable ThreadLocalAccessor accessor, WebGraphQlInterceptor... interceptors) {
try {
return new GraphQlWebSocketHandler(initHandler(interceptors), converter, Duration.ofSeconds(60));
return new GraphQlWebSocketHandler(
initHandler(accessor, interceptors), converter, Duration.ofSeconds(60));
}
catch (Exception ex) {
throw new IllegalStateException(ex);

13
spring-graphql/src/testFixtures/java/org/springframework/graphql/GraphQlSetup.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2002-2021 the original author or authors.
* Copyright 2002-2022 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.
@@ -35,10 +35,11 @@ import org.springframework.graphql.execution.DefaultExecutionGraphQlService;
import org.springframework.graphql.execution.GraphQlSource;
import org.springframework.graphql.execution.RuntimeWiringConfigurer;
import org.springframework.graphql.execution.ThreadLocalAccessor;
import org.springframework.graphql.server.webflux.GraphQlHttpHandler;
import org.springframework.graphql.server.WebGraphQlHandler;
import org.springframework.graphql.server.WebGraphQlSetup;
import org.springframework.graphql.server.WebGraphQlInterceptor;
import org.springframework.graphql.server.WebGraphQlSetup;
import org.springframework.graphql.server.webflux.GraphQlHttpHandler;
import org.springframework.lang.Nullable;
/**
* Workflow for GraphQL tests setup that starts with {@link GraphQlSource.Builder}
@@ -141,8 +142,10 @@ public class GraphQlSetup implements GraphQlServiceSetup {
}
@Override
public WebGraphQlSetup threadLocalAccessor(ThreadLocalAccessor... accessors) {
this.accessors.addAll(Arrays.asList(accessors));
public WebGraphQlSetup threadLocalAccessor(@Nullable ThreadLocalAccessor accessor) {
if (accessor != null) {
this.accessors.add(accessor);
}
return this;
}

5
spring-graphql/src/testFixtures/java/org/springframework/graphql/server/WebGraphQlSetup.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2002-2021 the original author or authors.
* Copyright 2002-2022 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.
@@ -17,6 +17,7 @@ package org.springframework.graphql.server;
import org.springframework.graphql.execution.ThreadLocalAccessor;
import org.springframework.graphql.server.webflux.GraphQlHttpHandler;
import org.springframework.lang.Nullable;
/**
* Workflow that results in the creation of a {@link WebGraphQlHandler} or
@@ -28,7 +29,7 @@ public interface WebGraphQlSetup {
WebGraphQlSetup interceptor(WebGraphQlInterceptor... interceptors);
WebGraphQlSetup threadLocalAccessor(ThreadLocalAccessor... accessors);
WebGraphQlSetup threadLocalAccessor(@Nullable ThreadLocalAccessor accessor);
WebGraphQlHandler toWebGraphQlHandler();