Issue #9682 - fix RetainableByteBuffer release bug in WebSocket

Signed-off-by: Lachlan Roberts <lachlan@webtide.com>
This commit is contained in:
Lachlan Roberts 2023-05-11 10:58:42 +10:00
parent 413a644a44
commit c8b8ef6bd5
3 changed files with 86 additions and 13 deletions

View File

@ -220,6 +220,15 @@ public class WebSocketConnection extends AbstractConnection implements Connectio
if (!coreSession.isClosed())
coreSession.onEof();
flusher.onClose(cause);
try (AutoLock l = lock.lock())
{
if (networkBuffer != null)
{
networkBuffer.clear();
releaseNetworkBuffer();
}
}
super.onClose(cause);
}

View File

@ -81,8 +81,8 @@ public class JettyWebSocketFrameHandler implements FrameHandler
private MessageSink activeMessageSink;
private WebSocketSession session;
private SuspendState state = SuspendState.DEMANDING;
private Runnable delayedOnFrame;
private CoreSession coreSession;
private Frame delayedFrame;
private Callback delayedCallback;
public JettyWebSocketFrameHandler(WebSocketContainer container,
Object endpointInstance,
@ -151,7 +151,6 @@ public class JettyWebSocketFrameHandler implements FrameHandler
try
{
customizer.customize(coreSession);
this.coreSession = coreSession;
session = new WebSocketSession(container, coreSession, this);
if (!session.isOpen())
throw new IllegalStateException("Session is not open");
@ -199,7 +198,8 @@ public class JettyWebSocketFrameHandler implements FrameHandler
break;
case SUSPENDING:
delayedOnFrame = () -> onFrame(frame, callback);
delayedFrame = frame;
delayedCallback = callback;
state = SuspendState.SUSPENDED;
return;
@ -283,12 +283,19 @@ public class JettyWebSocketFrameHandler implements FrameHandler
@Override
public void onClosed(CloseStatus closeStatus, Callback callback)
{
Callback suspendedCallback;
try (AutoLock l = lock.lock())
{
// We are now closed and cannot suspend or resume.
state = SuspendState.CLOSED;
delayedFrame = null;
suspendedCallback = delayedCallback;
delayedCallback = null;
}
if (suspendedCallback != null)
suspendedCallback.failed(new CloseException(closeStatus.getCode(), closeStatus.getCause()));
notifyOnClose(closeStatus, callback);
container.notifySessionListeners((listener) -> listener.onWebSocketSessionClosed(session));
}
@ -447,7 +454,8 @@ public class JettyWebSocketFrameHandler implements FrameHandler
public void resume()
{
boolean needDemand = false;
Runnable delayedFrame = null;
Frame frame = null;
Callback callback = null;
try (AutoLock l = lock.lock())
{
switch (state)
@ -457,13 +465,13 @@ public class JettyWebSocketFrameHandler implements FrameHandler
case SUSPENDED:
needDemand = true;
delayedFrame = delayedOnFrame;
delayedOnFrame = null;
frame = delayedFrame;
callback = delayedCallback;
state = SuspendState.DEMANDING;
break;
case SUSPENDING:
if (delayedOnFrame != null)
if (delayedFrame != null)
throw new IllegalStateException();
state = SuspendState.DEMANDING;
break;
@ -475,8 +483,8 @@ public class JettyWebSocketFrameHandler implements FrameHandler
if (needDemand)
{
if (delayedFrame != null)
delayedFrame.run();
if (frame != null)
onFrame(frame, callback);
else
session.getCoreSession().demand(1);
}

View File

@ -15,16 +15,21 @@ package org.eclipse.jetty.websocket.tests;
import java.io.IOException;
import java.net.URI;
import java.time.Duration;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import org.eclipse.jetty.io.ArrayRetainableByteBufferPool;
import org.eclipse.jetty.server.Server;
import org.eclipse.jetty.server.ServerConnector;
import org.eclipse.jetty.servlet.ServletContextHandler;
import org.eclipse.jetty.servlet.ServletHolder;
import org.eclipse.jetty.websocket.api.BatchMode;
import org.eclipse.jetty.websocket.api.Session;
import org.eclipse.jetty.websocket.api.StatusCode;
import org.eclipse.jetty.websocket.api.SuspendToken;
import org.eclipse.jetty.websocket.api.annotations.WebSocket;
import org.eclipse.jetty.websocket.api.exceptions.WebSocketTimeoutException;
import org.eclipse.jetty.websocket.client.WebSocketClient;
import org.eclipse.jetty.websocket.server.JettyWebSocketServlet;
import org.eclipse.jetty.websocket.server.JettyWebSocketServletFactory;
@ -34,7 +39,10 @@ import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.instanceOf;
import static org.hamcrest.Matchers.is;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
@ -64,14 +72,15 @@ public class SuspendResumeTest
}
}
private Server server = new Server();
private WebSocketClient client = new WebSocketClient();
private SuspendSocket serverSocket = new SuspendSocket();
private Server server;
private WebSocketClient client;
private SuspendSocket serverSocket;
private ServerConnector connector;
@BeforeEach
public void start() throws Exception
{
server = new Server();
connector = new ServerConnector(server);
server.addConnector(connector);
@ -79,10 +88,12 @@ public class SuspendResumeTest
contextHandler.setContextPath("/");
server.setHandler(contextHandler);
contextHandler.addServlet(new ServletHolder(new UpgradeServlet()), "/suspend");
serverSocket = new SuspendSocket();
JettyWebSocketServletContainerInitializer.configure(contextHandler, null);
server.start();
client = new WebSocketClient();
client.start();
}
@ -189,4 +200,49 @@ public class SuspendResumeTest
// suspend after closed throws ISE
assertThrows(IllegalStateException.class, () -> clientSocket.session.suspend());
}
@Test
public void timeoutWhileSuspended() throws Exception
{
URI uri = new URI("ws://localhost:" + connector.getLocalPort() + "/suspend");
EventSocket clientSocket = new EventSocket();
Future<Session> connect = client.connect(clientSocket, uri);
connect.get(5, TimeUnit.SECONDS);
assertTrue(serverSocket.openLatch.await(5, TimeUnit.SECONDS));
// Set short idleTimeout on server.
int idleTimeout = 1000;
serverSocket.session.setIdleTimeout(Duration.ofMillis(idleTimeout));
// Suspend on the server.
clientSocket.session.getRemote().sendString("suspend");
assertThat(serverSocket.textMessages.poll(5, TimeUnit.SECONDS), is("suspend"));
// Send two messages, with batching on, so they are read into same network buffer on the server.
// First frame is read and delayed inside the JettyWebSocketFrameHandler suspendState, second frame remains in the network buffer.
clientSocket.session.getRemote().setBatchMode(BatchMode.ON);
clientSocket.session.getRemote().sendString("no demand");
clientSocket.session.getRemote().sendString("this should sit in network buffer");
clientSocket.session.getRemote().flush();
assertNotNull(serverSocket.suspendToken);
// Make sure both sides are closed.
assertTrue(serverSocket.closeLatch.await(5, TimeUnit.SECONDS));
assertTrue(clientSocket.closeLatch.await(5, TimeUnit.SECONDS));
// We received no additional messages.
assertNull(serverSocket.textMessages.poll());
assertNull(serverSocket.binaryMessages.poll());
// Check the idleTimeout occurred.
assertThat(serverSocket.error, instanceOf(WebSocketTimeoutException.class));
assertNull(clientSocket.error);
assertThat(clientSocket.closeCode, equalTo(StatusCode.SHUTDOWN));
assertThat(clientSocket.closeReason, equalTo("Connection Idle Timeout"));
// We should have no used buffers in the pool.
ArrayRetainableByteBufferPool pool = (ArrayRetainableByteBufferPool)connector.getByteBufferPool().asRetainableByteBufferPool();
assertThat(pool.getHeapByteBufferCount(), equalTo(pool.getAvailableHeapByteBufferCount()));
assertThat(pool.getDirectByteBufferCount(), equalTo(pool.getAvailableDirectByteBufferCount()));
}
}