From 7a1bbf08b364bfd9176aedc87785c1c7e3495fd4 Mon Sep 17 00:00:00 2001
From: Spencer Gibb <spencer@gibb.us>
Date: Wed, 16 May 2018 17:15:34 -0400
Subject: [PATCH] Uses Mono.zip().then() rather than Mono.when().

fixes gh-260
---
 .../filter/WebsocketRoutingFilter.java        |  2 +-
 .../websocket/WebSocketIntegrationTests.java  | 25 +++++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/WebsocketRoutingFilter.java b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/WebsocketRoutingFilter.java
index 4437f333..f8c4c18c 100644
--- a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/WebsocketRoutingFilter.java
+++ b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/WebsocketRoutingFilter.java
@@ -157,7 +157,7 @@ public class WebsocketRoutingFilter implements GlobalFilter, Ordered {
 					Mono<Void> serverSessionSend = session
 							.send(proxySession.receive().doOnNext(WebSocketMessage::retain));
                             // .log("sessionSend", Level.FINE);
-					return Mono.when(proxySessionSend, serverSessionSend);
+					return Mono.zip(proxySessionSend, serverSessionSend).then();
 				}
 
 				/**
diff --git a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/test/websocket/WebSocketIntegrationTests.java b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/test/websocket/WebSocketIntegrationTests.java
index 953537af..9815a04e 100644
--- a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/test/websocket/WebSocketIntegrationTests.java
+++ b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/test/websocket/WebSocketIntegrationTests.java
@@ -34,6 +34,7 @@ import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.reactivestreams.Publisher;
+import org.springframework.web.reactive.socket.CloseStatus;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.MonoProcessor;
@@ -177,6 +178,7 @@ public class WebSocketIntegrationTests {
 			map.put("/echoForHttp", new EchoWebSocketHandler());
 			map.put("/sub-protocol", new SubProtocolWebSocketHandler());
 			map.put("/custom-header", new CustomHeaderHandler());
+			map.put("/close", new SessionClosingHandler());
 
 			SimpleUrlHandlerMapping mapping = new SimpleUrlHandlerMapping();
 			mapping.setUrlMap(map);
@@ -287,6 +289,21 @@ public class WebSocketIntegrationTests {
 		assertEquals("my-header:my-value", output.block(Duration.ofMillis(5000)));
 	}
 
+	@Test
+	public void sessionClosing() throws Exception {
+		this.client.execute(getUrl("/close"),
+				session -> {
+					logger.debug("Starting..");
+					return session.receive()
+							.doOnNext(s -> logger.debug("inbound " + s))
+							.then()
+							.doFinally(signalType -> {
+								logger.debug("Completed with: " + signalType);
+							});
+				})
+				.block(Duration.ofMillis(5000));
+	}
+
 	private static class EchoWebSocketHandler implements WebSocketHandler {
 
 		@Override
@@ -330,6 +347,14 @@ public class WebSocketIntegrationTests {
 		}
 	}
 
+	private static class SessionClosingHandler implements WebSocketHandler {
+
+		@Override
+		public Mono<Void> handle(WebSocketSession session) {
+			return Flux.never().mergeWith(session.close(CloseStatus.GOING_AWAY)).then();
+		}
+	}
+
 	private static Mono<Void> doSend(WebSocketSession session, Publisher<WebSocketMessage> output) {
 		return session.send(output);
 		// workaround for suspected RxNetty WebSocket client issue