Fix concurrency issues in SockJS session impls

Cherry-picked from:
fcf6ae8328

Issue: SPR-11916

spring-websocket/src/test/java/org/springframework/web/socket/sockjs/transport/session/HttpSockJsSessionTests.java

@@ -94,10 +94,8 @@ public class HttpSockJsSessionTests extends AbstractSockJsSessionTests<TestAbstr
 		this.session.handleSuccessiveRequest(this.request, this.response, this.frameFormat);
 
 		assertTrue(this.servletRequest.isAsyncStarted());
-
 		assertTrue(this.session.wasHeartbeatScheduled());
 		assertTrue(this.session.wasCacheFlushed());
-
 		assertEquals("hhh\n", this.servletResponse.getContentAsString());
 
 		verifyNoMoreInteractions(this.webSocketHandler);
@@ -119,6 +117,11 @@ public class HttpSockJsSessionTests extends AbstractSockJsSessionTests<TestAbstr
 			super("1", config, handler, attributes);
 		}
 
+		@Override
+		protected boolean isStreaming() {
+			return false;
+		}
+
 		@Override
 		protected void writePrelude(ServerHttpRequest request, ServerHttpResponse response) throws IOException {
 			response.getBody().write("hhh\n".getBytes());
@@ -139,6 +142,7 @@ public class HttpSockJsSessionTests extends AbstractSockJsSessionTests<TestAbstr
 		@Override
 		protected void flushCache() {
 			this.cacheFlushed = true;
+			scheduleHeartbeat();
 		}
 
 		@Override

spring-websocket/src/test/java/org/springframework/web/socket/sockjs/transport/session/TestHttpSockJsSession.java

@@ -53,6 +53,11 @@ public class TestHttpSockJsSession extends AbstractHttpSockJsSession {
 		super(sessionId, config, wsHandler, attributes);
 	}
 
+	@Override
+	protected boolean isStreaming() {
+		return true;
+	}
+
 	@Override
 	public String getAcceptedProtocol() {
 		return this.subProtocol;

spring-websocket/src/test/java/org/springframework/web/socket/sockjs/transport/session/WebSocketServerSockJsSessionTests.java

@@ -28,6 +28,8 @@ import org.junit.Test;
 import org.springframework.web.socket.CloseStatus;
 import org.springframework.web.socket.TextMessage;
 import org.springframework.web.socket.WebSocketHandler;
+import org.springframework.web.socket.WebSocketSession;
+import org.springframework.web.socket.handler.TextWebSocketHandler;
 import org.springframework.web.socket.sockjs.transport.SockJsServiceConfig;
 import org.springframework.web.socket.sockjs.transport.session.WebSocketServerSockJsSessionTests.TestWebSocketServerSockJsSession;
 import org.springframework.web.socket.handler.TestWebSocketSession;
@@ -72,17 +74,27 @@ public class WebSocketServerSockJsSessionTests extends AbstractSockJsSessionTest
 
 	@Test
 	public void afterSessionInitialized() throws Exception {
-
 		this.session.initializeDelegateSession(this.webSocketSession);
-
-		assertEquals("Open frame not sent",
-				Collections.singletonList(new TextMessage("o")), this.webSocketSession.getSentMessages());
-
+		assertEquals(Collections.singletonList(new TextMessage("o")), this.webSocketSession.getSentMessages());
 		assertEquals(Arrays.asList("schedule"), this.session.heartbeatSchedulingEvents);
 		verify(this.webSocketHandler).afterConnectionEstablished(this.session);
 		verifyNoMoreInteractions(this.taskScheduler, this.webSocketHandler);
 	}
 
+	@Test
+	public void afterSessionInitializedOpenFrameFirst() throws Exception {
+		TextWebSocketHandler handler = new TextWebSocketHandler() {
+			@Override
+			public void afterConnectionEstablished(WebSocketSession session) throws Exception {
+				session.sendMessage(new TextMessage("go go"));
+			}
+		};
+		TestWebSocketServerSockJsSession session = new TestWebSocketServerSockJsSession(this.sockJsConfig, handler, null);
+		session.initializeDelegateSession(this.webSocketSession);
+		List<TextMessage> expected = Arrays.asList(new TextMessage("o"), new TextMessage("a[\"go go\"]"));
+		assertEquals(expected, this.webSocketSession.getSentMessages());
+	}
+
 	@Test
 	public void handleMessageEmptyPayload() throws Exception {
 		this.session.handleMessage(new TextMessage(""), this.webSocketSession);