Polishing

Author: poutsma

spring-core/src/main/java/org/springframework/core/io/buffer/JettyDataBuffer.java

@@ -305,6 +305,22 @@ public String toString(int index, int length, Charset charset) {
 return this.delegate.toString(index, length, charset);
 }
 
+@Override
+public int hashCode() {
+return this.delegate.hashCode();
+}
+
+@Override
+public boolean equals(Object o) {
+return this == o || (o instanceof JettyDataBuffer other &&
+this.delegate.equals(other.delegate));
+}
+
+@Override
+public String toString() {
+return String.format("JettyDataBuffer (r: %d, w: %d, c: %d)",
+readPosition(), writePosition(), capacity());
+}
 
 private static final class JettyByteBufferIterator implements ByteBufferIterator {
 

spring-webflux/src/main/java/org/springframework/web/reactive/socket/adapter/JettyWebSocketHandlerAdapter.java

@@ -19,15 +19,12 @@
 import java.nio.ByteBuffer;
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
-import java.util.Objects;
 import java.util.function.Function;
 import java.util.function.IntPredicate;
 
 import org.eclipse.jetty.util.BufferUtil;
 import org.eclipse.jetty.websocket.api.Callback;
 import org.eclipse.jetty.websocket.api.Session;
-import org.reactivestreams.Subscriber;
-import org.reactivestreams.Subscription;
 
 import org.springframework.core.io.buffer.CloseableDataBuffer;
 import org.springframework.core.io.buffer.DataBuffer;
@@ -47,15 +44,16 @@
  * @author Rossen Stoyanchev
  * @since 5.0
  */
-@SuppressWarnings("NullAway")
 public class JettyWebSocketHandlerAdapter implements Session.Listener {
+
 private final WebSocketHandler delegateHandler;
 
 private final Function<Session, JettyWebSocketSession> sessionFactory;
 
 @Nullable
 private JettyWebSocketSession delegateSession;
 
+
 public JettyWebSocketHandlerAdapter(WebSocketHandler handler,
 Function<Session, JettyWebSocketSession> sessionFactory) {
 
@@ -67,32 +65,16 @@ public JettyWebSocketHandlerAdapter(WebSocketHandler handler,
 
 @Override
 public void onWebSocketOpen(Session session) {
-this.delegateSession = Objects.requireNonNull(this.sessionFactory.apply(session));
-this.delegateHandler.handle(this.delegateSession)
-.subscribe(new Subscriber<>() {
-@Override
-public void onSubscribe(Subscription s) {
-s.request(Long.MAX_VALUE);
-}
-
-@Override
-public void onNext(Void unused) {
-}
-
-@Override
-public void onError(Throwable t) {
-delegateSession.onHandlerError(t);
-}
-
-@Override
-public void onComplete() {
-delegateSession.onHandleComplete();
-}
-});
+JettyWebSocketSession delegateSession = this.sessionFactory.apply(session);
+this.delegateSession = delegateSession;
+this.delegateHandler.handle(delegateSession)
+.checkpoint(session.getUpgradeRequest().getRequestURI() + " [JettyWebSocketHandlerAdapter]")
+.subscribe(unused -> {}, delegateSession::onHandlerError, delegateSession::onHandleComplete);
 }
 
 @Override
 public void onWebSocketText(String message) {
+Assert.state(this.delegateSession != null, "No delegate session available");
 byte[] bytes = message.getBytes(StandardCharsets.UTF_8);
 DataBuffer buffer = this.delegateSession.bufferFactory().wrap(bytes);
 WebSocketMessage webSocketMessage = new WebSocketMessage(Type.TEXT, buffer);
@@ -101,37 +83,42 @@ public void onWebSocketText(String message) {
 
 @Override
 public void onWebSocketBinary(ByteBuffer byteBuffer, Callback callback) {
+Assert.state(this.delegateSession != null, "No delegate session available");
 DataBuffer buffer = this.delegateSession.bufferFactory().wrap(byteBuffer);
-buffer = new JettyDataBuffer(buffer, callback);
+buffer = new JettyCallbackDataBuffer(buffer, callback);
 WebSocketMessage webSocketMessage = new WebSocketMessage(Type.BINARY, buffer);
 this.delegateSession.handleMessage(webSocketMessage);
 }
 
 @Override
 public void onWebSocketPong(ByteBuffer payload) {
+Assert.state(this.delegateSession != null, "No delegate session available");
 DataBuffer buffer = this.delegateSession.bufferFactory().wrap(BufferUtil.copy(payload));
 WebSocketMessage webSocketMessage = new WebSocketMessage(Type.PONG, buffer);
 this.delegateSession.handleMessage(webSocketMessage);
 }
 
 @Override
 public void onWebSocketClose(int statusCode, String reason) {
+Assert.state(this.delegateSession != null, "No delegate session available");
 this.delegateSession.handleClose(CloseStatus.create(statusCode, reason));
 }
 
 @Override
 public void onWebSocketError(Throwable cause) {
+Assert.state(this.delegateSession != null, "No delegate session available");
 this.delegateSession.handleError(cause);
 }
 
 
-private static final class JettyDataBuffer implements CloseableDataBuffer {
+private static final class JettyCallbackDataBuffer implements CloseableDataBuffer {
 
 private final DataBuffer delegate;
 
 private final Callback callback;
 
-public JettyDataBuffer(DataBuffer delegate, Callback callback) {
+
+public JettyCallbackDataBuffer(DataBuffer delegate, Callback callback) {
 Assert.notNull(delegate, "'delegate` must not be null");
 Assert.notNull(callback, "Callback must not be null");
 this.delegate = delegate;
@@ -266,13 +253,13 @@ public DataBuffer write(ByteBuffer... buffers) {
 @Deprecated
 public DataBuffer slice(int index, int length) {
 DataBuffer delegateSlice = this.delegate.slice(index, length);
-return new JettyDataBuffer(delegateSlice, this.callback);
+return new JettyCallbackDataBuffer(delegateSlice, this.callback);
 }
 
 @Override
 public DataBuffer split(int index) {
 DataBuffer delegateSplit = this.delegate.split(index);
-return new JettyDataBuffer(delegateSplit, this.callback);
+return new JettyCallbackDataBuffer(delegateSplit, this.callback);
 }
 
 @Override

spring-webflux/src/main/java/org/springframework/web/reactive/socket/adapter/JettyWebSocketSession.java

@@ -53,9 +53,13 @@
 public class JettyWebSocketSession extends AbstractWebSocketSession<Session> {
 
 private final Flux<WebSocketMessage> flux;
+
 private final Sinks.One<CloseStatus> closeStatusSink = Sinks.one();
+
 private final Lock lock = new ReentrantLock();
+
 private long requested = 0;
+
 private boolean awaitingMessage = false;
 
 @Nullable

spring-websocket/src/main/java/org/springframework/web/socket/adapter/jetty/JettyWebSocketHandlerAdapter.java

@@ -39,16 +39,18 @@
  * @author Rossen Stoyanchev
  * @since 4.0
  */
-@SuppressWarnings("NullAway")
 public class JettyWebSocketHandlerAdapter implements Session.Listener {
+
 private static final Log logger = LogFactory.getLog(JettyWebSocketHandlerAdapter.class);
 
 private final WebSocketHandler webSocketHandler;
+
 private final JettyWebSocketSession wsSession;
 
 @Nullable
 private Session nativeSession;
 
+
 public JettyWebSocketHandlerAdapter(WebSocketHandler webSocketHandler, JettyWebSocketSession wsSession) {
 Assert.notNull(webSocketHandler, "WebSocketHandler must not be null");
 Assert.notNull(wsSession, "WebSocketSession must not be null");
@@ -71,6 +73,7 @@ public void onWebSocketOpen(Session session) {
 
 @Override
 public void onWebSocketText(String payload) {
+Assert.state(this.nativeSession != null, "No native session available");
 TextMessage message = new TextMessage(payload);
 try {
 this.webSocketHandler.handleMessage(this.wsSession, message);
@@ -83,6 +86,7 @@ public void onWebSocketText(String payload) {
 
 @Override
 public void onWebSocketBinary(ByteBuffer payload, Callback callback) {
+Assert.state(this.nativeSession != null, "No native session available");
 BinaryMessage message = new BinaryMessage(BufferUtil.copy(payload), true);
 callback.succeed();
 try {
@@ -96,6 +100,7 @@ public void onWebSocketBinary(ByteBuffer payload, Callback callback) {
 
 @Override
 public void onWebSocketPong(ByteBuffer payload) {
+Assert.state(this.nativeSession != null, "No native session available");
 PongMessage message = new PongMessage(BufferUtil.copy(payload));
 try {
 this.webSocketHandler.handleMessage(this.wsSession, message);
@@ -132,13 +137,14 @@ public void onWebSocketError(Throwable cause) {
 }
 
 private void tryCloseWithError(Throwable t) {
-
-if (this.nativeSession.isOpen()) {
-ExceptionWebSocketHandlerDecorator.tryCloseWithError(this.wsSession, t, logger);
-}
-else {
-// Session might be O-SHUT waiting for response close frame, so abort to close the connection.
-this.nativeSession.disconnect();
+if (this.nativeSession != null) {
+if (this.nativeSession.isOpen()) {
+ExceptionWebSocketHandlerDecorator.tryCloseWithError(this.wsSession, t, logger);
+}
+else {
+// Session might be O-SHUT waiting for response close frame, so abort to close the connection.
+this.nativeSession.disconnect();
+}
 }
 }
 }