Issue #5368 - ensure onMessage exits before next frame is read

Signed-off-by: Lachlan Roberts <lachlan@webtide.com>
This commit is contained in:
Lachlan Roberts 2020-10-01 15:42:50 +10:00
parent e3ed05fc1c
commit 941ffcead7
5 changed files with 98 additions and 67 deletions

View File

@ -118,10 +118,8 @@ public class JsrAnnotatedEventDriver extends AbstractJsrEventDriver
if (LOG.isDebugEnabled())
LOG.debug("Binary Message InputStream");
final MessageInputStream stream = new MessageInputStream(session);
MessageInputStream stream = new MessageInputStream(session);
activeMessage = stream;
// Always dispatch streaming read to another thread.
dispatch(() ->
{
try
@ -329,11 +327,8 @@ public class JsrAnnotatedEventDriver extends AbstractJsrEventDriver
if (LOG.isDebugEnabled())
LOG.debug("Text Message Writer");
MessageInputStream inputStream = new MessageInputStream(session);
final MessageReader reader = new MessageReader(inputStream);
activeMessage = inputStream;
// Always dispatch streaming read to another thread.
MessageReader reader = new MessageReader(session);
activeMessage = reader;
dispatch(() ->
{
try
@ -343,9 +338,10 @@ public class JsrAnnotatedEventDriver extends AbstractJsrEventDriver
catch (Throwable e)
{
session.close(e);
return;
}
inputStream.close();
reader.handlerComplete();
});
}
}

View File

@ -100,9 +100,10 @@ public class JsrEndpointEventDriver extends AbstractJsrEventDriver
catch (Throwable t)
{
session.close(t);
return;
}
inputStream.close();
inputStream.handlerComplete();
});
}
else
@ -197,8 +198,7 @@ public class JsrEndpointEventDriver extends AbstractJsrEventDriver
{
@SuppressWarnings("unchecked")
MessageHandler.Whole<Reader> handler = (Whole<Reader>)wrapper.getHandler();
MessageInputStream inputStream = new MessageInputStream(session);
MessageReader reader = new MessageReader(inputStream);
MessageReader reader = new MessageReader(session);
activeMessage = reader;
dispatch(() ->
{
@ -209,9 +209,10 @@ public class JsrEndpointEventDriver extends AbstractJsrEventDriver
catch (Throwable t)
{
session.close(t);
return;
}
inputStream.close();
reader.handlerComplete();
});
}
else

View File

@ -32,7 +32,6 @@ import org.eclipse.jetty.websocket.api.WebSocketPolicy;
import org.eclipse.jetty.websocket.api.annotations.WebSocket;
import org.eclipse.jetty.websocket.api.extensions.Frame;
import org.eclipse.jetty.websocket.common.CloseInfo;
import org.eclipse.jetty.websocket.common.message.MessageAppender;
import org.eclipse.jetty.websocket.common.message.MessageInputStream;
import org.eclipse.jetty.websocket.common.message.MessageReader;
import org.eclipse.jetty.websocket.common.message.NullMessage;
@ -105,7 +104,7 @@ public class JettyAnnotatedEventDriver extends AbstractEventDriver
}
else if (events.onBinary.isStreaming())
{
final MessageInputStream inputStream = new MessageInputStream(session);
MessageInputStream inputStream = new MessageInputStream(session);
activeMessage = inputStream;
dispatch(() ->
{
@ -115,11 +114,11 @@ public class JettyAnnotatedEventDriver extends AbstractEventDriver
}
catch (Throwable t)
{
// dispatched calls need to be reported
session.close(t);
return;
}
inputStream.close();
inputStream.handlerComplete();
});
}
else
@ -262,22 +261,21 @@ public class JettyAnnotatedEventDriver extends AbstractEventDriver
}
else if (events.onText.isStreaming())
{
MessageInputStream inputStream = new MessageInputStream(session);
activeMessage = new MessageReader(inputStream);
final MessageAppender msg = activeMessage;
MessageReader reader = new MessageReader(session);
activeMessage = reader;
dispatch(() ->
{
try
{
events.onText.call(websocket, session, msg);
events.onText.call(websocket, session, reader);
}
catch (Throwable t)
{
// dispatched calls need to be reported
session.close(t);
return;
}
inputStream.close();
reader.handlerComplete();
});
}
else

View File

@ -55,6 +55,7 @@ public class MessageInputStream extends InputStream implements MessageAppender
{
RESUMED,
SUSPENDED,
COMPLETE,
CLOSED
}
@ -76,23 +77,11 @@ public class MessageInputStream extends InputStream implements MessageAppender
if (LOG.isDebugEnabled())
LOG.debug("Appending {} chunk: {}", fin ? "final" : "non-final", BufferUtil.toDetailString(framePayload));
// Early non atomic test that we aren't closed to avoid an unnecessary copy (will be checked again later).
if (state == State.CLOSED)
return;
// Put the payload into the queue, by copying it.
// Copying is necessary because the payload will
// be processed after this method returns.
try
{
if (framePayload == null || !framePayload.hasRemaining())
if (BufferUtil.isEmpty(framePayload))
return;
ByteBuffer copy = acquire(framePayload.remaining(), framePayload.isDirect());
BufferUtil.clearToFill(copy);
copy.put(framePayload);
BufferUtil.flipToFlush(copy, 0);
synchronized (this)
{
switch (state)
@ -105,11 +94,14 @@ public class MessageInputStream extends InputStream implements MessageAppender
state = State.SUSPENDED;
break;
case SUSPENDED:
default:
throw new IllegalStateException();
}
buffers.put(copy);
// Put the payload into the queue, by copying it.
// Copying is necessary because the payload will
// be processed after this method returns.
buffers.put(copy(framePayload));
}
}
catch (InterruptedException e)
@ -121,7 +113,23 @@ public class MessageInputStream extends InputStream implements MessageAppender
@Override
public void close()
{
SuspendToken resume = null;
synchronized (this)
{
if (state == State.CLOSED)
return;
state = State.CLOSED;
buffers.clear();
buffers.offer(EOF);
}
}
@Override
public void messageComplete()
{
if (LOG.isDebugEnabled())
LOG.debug("Message completed");
synchronized (this)
{
switch (state)
@ -130,45 +138,35 @@ public class MessageInputStream extends InputStream implements MessageAppender
return;
case SUSPENDED:
resume = suspendToken;
suspendToken = null;
state = State.CLOSED;
case RESUMED:
state = State.COMPLETE;
break;
case RESUMED:
state = State.CLOSED;
break;
default:
throw new IllegalStateException();
}
buffers.offer(EOF);
}
}
public void handlerComplete()
{
// May need to resume to resume and read to the next message.
SuspendToken resume;
synchronized (this)
{
state = State.CLOSED;
resume = suspendToken;
suspendToken = null;
buffers.clear();
buffers.offer(EOF);
}
// May need to resume to discard until we reach next message.
if (resume != null)
resume.resume();
}
@Override
public void mark(int readlimit)
{
// Not supported.
}
@Override
public boolean markSupported()
{
return false;
}
@Override
public void messageComplete()
{
if (LOG.isDebugEnabled())
LOG.debug("Message completed");
buffers.offer(EOF);
}
@Override
public int read() throws IOException
{
@ -186,6 +184,7 @@ public class MessageInputStream extends InputStream implements MessageAppender
{
if (LOG.isDebugEnabled())
LOG.debug("Waiting {} ms to read", timeoutMs);
if (timeoutMs < 0)
{
// Wait forever until a buffer is available.
@ -212,7 +211,6 @@ public class MessageInputStream extends InputStream implements MessageAppender
int result = activeBuffer.get() & 0xFF;
if (!activeBuffer.hasRemaining())
{
SuspendToken resume = null;
synchronized (this)
{
@ -221,6 +219,11 @@ public class MessageInputStream extends InputStream implements MessageAppender
case CLOSED:
return -1;
case COMPLETE:
// If we are complete we have read the last frame but
// don't want to resume reading until onMessage() exits.
break;
case SUSPENDED:
resume = suspendToken;
suspendToken = null;
@ -254,6 +257,27 @@ public class MessageInputStream extends InputStream implements MessageAppender
throw new IOException("reset() not supported");
}
@Override
public void mark(int readlimit)
{
// Not supported.
}
@Override
public boolean markSupported()
{
return false;
}
private ByteBuffer copy(ByteBuffer buffer)
{
ByteBuffer copy = acquire(buffer.remaining(), buffer.isDirect());
BufferUtil.clearToFill(copy);
copy.put(buffer);
BufferUtil.flipToFlush(copy, 0);
return copy;
}
private ByteBuffer acquire(int capacity, boolean direct)
{
ByteBuffer buffer;

View File

@ -24,6 +24,8 @@ import java.nio.ByteBuffer;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import org.eclipse.jetty.websocket.api.Session;
/**
* Support class for reading a (single) WebSocket TEXT message via a Reader.
* <p>
@ -33,6 +35,11 @@ public class MessageReader extends InputStreamReader implements MessageAppender
{
private final MessageInputStream stream;
public MessageReader(Session session)
{
this(new MessageInputStream(session));
}
public MessageReader(MessageInputStream stream)
{
super(stream, StandardCharsets.UTF_8);
@ -50,4 +57,9 @@ public class MessageReader extends InputStreamReader implements MessageAppender
{
this.stream.messageComplete();
}
public void handlerComplete()
{
this.stream.handlerComplete();
}
}