This is an automated email from the ASF dual-hosted git repository.

markt-asf pushed a commit to branch 10.1.x
in repository https://gitbox.apache.org/repos/asf/tomcat.git


The following commit(s) were added to refs/heads/10.1.x by this push:
     new df2f47fe32 Refactor WebSocket writes to use a consistent per message 
timeout
df2f47fe32 is described below

commit df2f47fe32c4bd528af8b1db099c6710f34d9d5f
Author: Mark Thomas <[email protected]>
AuthorDate: Thu Sep 3 12:06:20 2026 +0100

    Refactor WebSocket writes to use a consistent per message timeout
    
    - Creates separate fields for blocking flag and timeout in MessagePart
    - Tracks the message level timeout
    - Write methods use the message level timeout
    
    Assisted-by: GPT-5.6-sol
---
 java/org/apache/tomcat/websocket/MessagePart.java  |  31 +++--
 .../apache/tomcat/websocket/PerMessageDeflate.java |  23 ++--
 .../tomcat/websocket/WsRemoteEndpointImplBase.java |  83 +++++++++-----
 .../websocket/WsRemoteEndpointImplClient.java      |  31 ++---
 .../server/WsRemoteEndpointImplServer.java         | 125 +++++++++++++--------
 .../tomcat/websocket/server/WsWriteTimeout.java    |   2 +-
 .../tomcat/websocket/TestPerMessageDeflate.java    |  20 +++-
 webapps/docs/changelog.xml                         |   4 +-
 8 files changed, 202 insertions(+), 117 deletions(-)

diff --git a/java/org/apache/tomcat/websocket/MessagePart.java 
b/java/org/apache/tomcat/websocket/MessagePart.java
index 49cf0a1b78..8f84d78864 100644
--- a/java/org/apache/tomcat/websocket/MessagePart.java
+++ b/java/org/apache/tomcat/websocket/MessagePart.java
@@ -36,8 +36,10 @@ public class MessagePart {
     private final SendHandler intermediateHandler;
     /** End send handler. */
     private volatile SendHandler endHandler;
-    /** Blocking write timeout expiry. */
-    private final long blockingWriteTimeoutExpiry;
+    /** Whether the write is blocking. */
+    private final boolean blocking;
+    /** Write timeout expiry. */
+    private final long writeTimeoutExpiry;
 
     /**
      * Constructor.
@@ -47,17 +49,19 @@ public class MessagePart {
      * @param payload payload data
      * @param intermediateHandler intermediate send handler
      * @param endHandler end send handler
-     * @param blockingWriteTimeoutExpiry blocking write timeout expiry
+     * @param blocking whether the write is blocking
+     * @param writeTimeoutExpiry write timeout expiry
      */
     MessagePart(boolean fin, int rsv, byte opCode, ByteBuffer payload, 
SendHandler intermediateHandler,
-            SendHandler endHandler, long blockingWriteTimeoutExpiry) {
+            SendHandler endHandler, boolean blocking, long writeTimeoutExpiry) 
{
         this.fin = fin;
         this.rsv = rsv;
         this.opCode = opCode;
         this.payload = payload;
         this.intermediateHandler = intermediateHandler;
         this.endHandler = endHandler;
-        this.blockingWriteTimeoutExpiry = blockingWriteTimeoutExpiry;
+        this.blocking = blocking;
+        this.writeTimeoutExpiry = writeTimeoutExpiry;
     }
 
     /**
@@ -117,11 +121,18 @@ public class MessagePart {
     }
 
     /**
-     * Get the blocking write timeout expiry.
-     * @return the blocking write timeout expiry
+     * Determine if the write is blocking.
+     * @return {@code true} if the write is blocking, otherwise {@code false}
      */
-    public long getBlockingWriteTimeoutExpiry() {
-        return blockingWriteTimeoutExpiry;
+    public boolean isBlocking() {
+        return blocking;
     }
-}
 
+    /**
+     * Get the write timeout expiry.
+     * @return the write timeout expiry
+     */
+    public long getWriteTimeoutExpiry() {
+        return writeTimeoutExpiry;
+    }
+}
diff --git a/java/org/apache/tomcat/websocket/PerMessageDeflate.java 
b/java/org/apache/tomcat/websocket/PerMessageDeflate.java
index c3627032c4..10ffa1dbbc 100644
--- a/java/org/apache/tomcat/websocket/PerMessageDeflate.java
+++ b/java/org/apache/tomcat/websocket/PerMessageDeflate.java
@@ -467,28 +467,29 @@ public class PerMessageDeflate implements Transformation {
                     boolean fin = uncompressedPart.isFin();
                     boolean full = compressedPayload.limit() == 
compressedPayload.capacity();
                     boolean needsInput = deflater.needsInput();
-                    long blockingWriteTimeoutExpiry = 
uncompressedPart.getBlockingWriteTimeoutExpiry();
+                    boolean blocking = uncompressedPart.isBlocking();
+                    long writeTimeoutExpiry = 
uncompressedPart.getWriteTimeoutExpiry();
 
                     if (fin && !full && needsInput) {
                         // End of compressed message. Drop EOM bytes and 
output.
                         compressedPayload.limit(compressedPayload.limit() - 
EOM_BYTES.length);
                         compressedPart = new MessagePart(true, 
getRsv(uncompressedPart), opCode, compressedPayload,
-                                uncompressedIntermediateHandler, 
uncompressedIntermediateHandler,
-                                blockingWriteTimeoutExpiry);
+                                uncompressedIntermediateHandler, 
uncompressedIntermediateHandler, blocking,
+                                writeTimeoutExpiry);
                         deflateRequired = false;
                         startNewMessage();
                     } else if (full && !needsInput) {
                         // Write buffer full and input message not fully read.
                         // Output and start new compressed part.
                         compressedPart = new MessagePart(false, 
getRsv(uncompressedPart), opCode, compressedPayload,
-                                uncompressedIntermediateHandler, 
uncompressedIntermediateHandler,
-                                blockingWriteTimeoutExpiry);
+                                uncompressedIntermediateHandler, 
uncompressedIntermediateHandler, blocking,
+                                writeTimeoutExpiry);
                     } else if (!fin && full/* note: needsInput is true here 
*/) {
                         // Write buffer full and input message not fully read.
                         // Output and get more data.
                         compressedPart = new MessagePart(false, 
getRsv(uncompressedPart), opCode, compressedPayload,
-                                uncompressedIntermediateHandler, 
uncompressedIntermediateHandler,
-                                blockingWriteTimeoutExpiry);
+                                uncompressedIntermediateHandler, 
uncompressedIntermediateHandler, blocking,
+                                writeTimeoutExpiry);
                         deflateRequired = false;
                     } else if (fin && full/* note: needsInput is true here */) 
{
                         // Write buffer full. Input fully read. Deflater may be
@@ -508,8 +509,8 @@ public class PerMessageDeflate implements Transformation {
                             // EOM has just been completed
                             compressedPayload.limit(compressedPayload.limit() 
- EOM_BYTES.length + eomBufferWritten);
                             compressedPart = new MessagePart(true, 
getRsv(uncompressedPart), opCode, compressedPayload,
-                                    uncompressedIntermediateHandler, 
uncompressedIntermediateHandler,
-                                    blockingWriteTimeoutExpiry);
+                                    uncompressedIntermediateHandler, 
uncompressedIntermediateHandler, blocking,
+                                    writeTimeoutExpiry);
                             deflateRequired = false;
                             startNewMessage();
                         } else {
@@ -517,8 +518,8 @@ public class PerMessageDeflate implements Transformation {
                             // Copy bytes to new write buffer
                             writeBuffer.put(EOM_BUFFER, 0, eomBufferWritten);
                             compressedPart = new MessagePart(false, 
getRsv(uncompressedPart), opCode, compressedPayload,
-                                    uncompressedIntermediateHandler, 
uncompressedIntermediateHandler,
-                                    blockingWriteTimeoutExpiry);
+                                    uncompressedIntermediateHandler, 
uncompressedIntermediateHandler, blocking,
+                                    writeTimeoutExpiry);
                         }
                     } else {
                         throw new 
IllegalStateException(sm.getString("perMessageDeflate.invalidState"));
diff --git a/java/org/apache/tomcat/websocket/WsRemoteEndpointImplBase.java 
b/java/org/apache/tomcat/websocket/WsRemoteEndpointImplBase.java
index 9dd87a6ea2..f1a18ece0b 100644
--- a/java/org/apache/tomcat/websocket/WsRemoteEndpointImplBase.java
+++ b/java/org/apache/tomcat/websocket/WsRemoteEndpointImplBase.java
@@ -271,8 +271,8 @@ public abstract class WsRemoteEndpointImplBase implements 
RemoteEndpoint {
             throw new 
IllegalArgumentException(sm.getString("wsRemoteEndpoint.nullHandler"));
         }
         stateMachine.textStart();
-        TextMessageSendHandler tmsh =
-                new TextMessageSendHandler(handler, CharBuffer.wrap(text), 
true, encoder, encoderBuffer, this);
+        TextMessageSendHandler tmsh = new TextMessageSendHandler(handler, 
CharBuffer.wrap(text), true, encoder,
+                encoderBuffer, this, getAsyncSendTimeoutExpiry());
         tmsh.write();
         // TextMessageSendHandler will update stateMachine when it completes
     }
@@ -314,7 +314,11 @@ public abstract class WsRemoteEndpointImplBase implements 
RemoteEndpoint {
 
 
     void sendMessageBlock(CharBuffer part, boolean last) throws IOException {
-        long timeout = getBlockingSendTimeout();
+        sendMessageBlock(part, last, 
getTimeoutExpiry(getBlockingSendTimeout()));
+    }
+
+
+    private void sendMessageBlock(CharBuffer part, boolean last, long 
timeoutExpiry) throws IOException {
         boolean isDone = false;
         while (!isDone) {
             encoderBuffer.clear();
@@ -324,7 +328,7 @@ public abstract class WsRemoteEndpointImplBase implements 
RemoteEndpoint {
             }
             isDone = !cr.isOverflow();
             encoderBuffer.flip();
-            sendMessageBlock(Constants.OPCODE_TEXT, encoderBuffer, last && 
isDone, timeout);
+            sendMessageBlockInternal(Constants.OPCODE_TEXT, encoderBuffer, 
last && isDone, timeoutExpiry);
         }
         stateMachine.complete(last);
     }
@@ -360,7 +364,7 @@ public abstract class WsRemoteEndpointImplBase implements 
RemoteEndpoint {
         BlockingSendHandler bsh = new BlockingSendHandler();
 
         List<MessagePart> messageParts = new ArrayList<>();
-        messageParts.add(new MessagePart(last, 0, opCode, payload, bsh, bsh, 
timeoutExpiry));
+        messageParts.add(new MessagePart(last, 0, opCode, payload, bsh, bsh, 
true, timeoutExpiry));
 
         messageParts = transformation.sendMessagePart(messageParts);
 
@@ -434,12 +438,17 @@ public abstract class WsRemoteEndpointImplBase implements 
RemoteEndpoint {
 
 
     void startMessage(byte opCode, ByteBuffer payload, boolean last, 
SendHandler handler) {
+        startMessage(opCode, payload, last, handler, 
getAsyncSendTimeoutExpiry());
+    }
+
+
+    void startMessage(byte opCode, ByteBuffer payload, boolean last, 
SendHandler handler, long timeoutExpiry) {
 
         wsSession.updateLastActiveWrite();
 
         List<MessagePart> messageParts = new ArrayList<>();
         messageParts.add(new MessagePart(last, 0, opCode, payload, 
intermediateMessageHandler,
-                new EndMessageHandler(this, handler), -1));
+                new EndMessageHandler(this, handler), false, timeoutExpiry));
 
         try {
             messageParts = transformation.sendMessagePart(messageParts);
@@ -539,7 +548,7 @@ public abstract class WsRemoteEndpointImplBase implements 
RemoteEndpoint {
             nextText = text;
             outputBuffer.flip();
             SendHandler flushHandler = new 
OutputBufferFlushSendHandler(outputBuffer, mp.getEndHandler());
-            doWrite(flushHandler, mp.getBlockingWriteTimeoutExpiry(), 
outputBuffer);
+            doWrite(flushHandler, mp.isBlocking(), mp.getWriteTimeoutExpiry(), 
outputBuffer);
             return;
         }
 
@@ -592,12 +601,12 @@ public abstract class WsRemoteEndpointImplBase implements 
RemoteEndpoint {
         if (getBatchingAllowed() || isMasked()) {
             // Need to write via output buffer
             OutputBufferSendHandler obsh =
-                    new OutputBufferSendHandler(mp.getEndHandler(), 
mp.getBlockingWriteTimeoutExpiry(), headerBuffer,
-                            mp.getPayload(), mask, outputBuffer, 
!getBatchingAllowed(), this);
+                    new OutputBufferSendHandler(mp.getEndHandler(), 
mp.isBlocking(), mp.getWriteTimeoutExpiry(),
+                            headerBuffer, mp.getPayload(), mask, outputBuffer, 
!getBatchingAllowed(), this);
             obsh.write();
         } else {
             // Can write directly
-            doWrite(mp.getEndHandler(), mp.getBlockingWriteTimeoutExpiry(), 
headerBuffer, mp.getPayload());
+            doWrite(mp.getEndHandler(), mp.isBlocking(), 
mp.getWriteTimeoutExpiry(), headerBuffer, mp.getPayload());
         }
 
         updateStats(payloadSize);
@@ -629,6 +638,15 @@ public abstract class WsRemoteEndpointImplBase implements 
RemoteEndpoint {
     }
 
 
+    private long getAsyncSendTimeoutExpiry() {
+        long timeout = getSendTimeout();
+        if (timeout <= 0) {
+            return Long.MAX_VALUE;
+        }
+        return System.currentTimeMillis() + timeout;
+    }
+
+
     /**
      * Wraps the user provided handler so that the end point is notified when 
the message is complete.
      */
@@ -859,11 +877,12 @@ public abstract class WsRemoteEndpointImplBase implements 
RemoteEndpoint {
 
     /**
      * Writes data to the underlying connection.
-     * @param handler the send handler
-     * @param blockingWriteTimeoutExpiry the timeout expiry time
-     * @param data the data buffers to write
+     * @param handler            the send handler
+     * @param blocking           whether the write is blocking
+     * @param writeTimeoutExpiry the timeout expiry time
+     * @param data               the data buffers to write
      */
-    protected abstract void doWrite(SendHandler handler, long 
blockingWriteTimeoutExpiry, ByteBuffer... data);
+    protected abstract void doWrite(SendHandler handler, boolean blocking, 
long writeTimeoutExpiry, ByteBuffer... data);
 
     /**
      * Checks if frames should be masked.
@@ -940,6 +959,7 @@ public abstract class WsRemoteEndpointImplBase implements 
RemoteEndpoint {
     private class TextMessageSendHandler implements SendHandler {
 
         private final SendHandler handler;
+        private final long writeTimeoutExpiry;
         private final CharBuffer message;
         private final boolean isLast;
         private final CharsetEncoder encoder;
@@ -948,8 +968,9 @@ public abstract class WsRemoteEndpointImplBase implements 
RemoteEndpoint {
         private volatile boolean isDone = false;
 
         TextMessageSendHandler(SendHandler handler, CharBuffer message, 
boolean isLast, CharsetEncoder encoder,
-                ByteBuffer encoderBuffer, WsRemoteEndpointImplBase endpoint) {
+                ByteBuffer encoderBuffer, WsRemoteEndpointImplBase endpoint, 
long writeTimeoutExpiry) {
             this.handler = handler;
+            this.writeTimeoutExpiry = writeTimeoutExpiry;
             this.message = message;
             this.isLast = isLast;
             this.encoder = encoder.reset();
@@ -965,7 +986,7 @@ public abstract class WsRemoteEndpointImplBase implements 
RemoteEndpoint {
             }
             isDone = !cr.isOverflow();
             buffer.flip();
-            endpoint.startMessage(Constants.OPCODE_TEXT, buffer, isDone && 
isLast, this);
+            endpoint.startMessage(Constants.OPCODE_TEXT, buffer, isDone && 
isLast, this, writeTimeoutExpiry);
         }
 
         @Override
@@ -991,7 +1012,8 @@ public abstract class WsRemoteEndpointImplBase implements 
RemoteEndpoint {
     private static class OutputBufferSendHandler implements SendHandler {
 
         private final SendHandler handler;
-        private final long blockingWriteTimeoutExpiry;
+        private final boolean blocking;
+        private final long writeTimeoutExpiry;
         private final ByteBuffer headerBuffer;
         private final ByteBuffer payload;
         private final byte[] mask;
@@ -1000,10 +1022,11 @@ public abstract class WsRemoteEndpointImplBase 
implements RemoteEndpoint {
         private final WsRemoteEndpointImplBase endpoint;
         private volatile int maskIndex = 0;
 
-        OutputBufferSendHandler(SendHandler completion, long 
blockingWriteTimeoutExpiry, ByteBuffer headerBuffer,
-                ByteBuffer payload, byte[] mask, ByteBuffer outputBuffer, 
boolean flushRequired,
+        OutputBufferSendHandler(SendHandler completion, boolean blocking, long 
writeTimeoutExpiry,
+                ByteBuffer headerBuffer, ByteBuffer payload, byte[] mask, 
ByteBuffer outputBuffer, boolean flushRequired,
                 WsRemoteEndpointImplBase endpoint) {
-            this.blockingWriteTimeoutExpiry = blockingWriteTimeoutExpiry;
+            this.blocking = blocking;
+            this.writeTimeoutExpiry = writeTimeoutExpiry;
             this.handler = completion;
             this.headerBuffer = headerBuffer;
             this.payload = payload;
@@ -1021,7 +1044,7 @@ public abstract class WsRemoteEndpointImplBase implements 
RemoteEndpoint {
             if (headerBuffer.hasRemaining()) {
                 // Still more headers to write, need to flush
                 outputBuffer.flip();
-                endpoint.doWrite(this, blockingWriteTimeoutExpiry, 
outputBuffer);
+                endpoint.doWrite(this, blocking, writeTimeoutExpiry, 
outputBuffer);
                 return;
             }
 
@@ -1054,7 +1077,7 @@ public abstract class WsRemoteEndpointImplBase implements 
RemoteEndpoint {
                 payload.limit(payloadLimit);
                 // Still more data to write, need to flush
                 outputBuffer.flip();
-                endpoint.doWrite(this, blockingWriteTimeoutExpiry, 
outputBuffer);
+                endpoint.doWrite(this, blocking, writeTimeoutExpiry, 
outputBuffer);
                 return;
             }
 
@@ -1063,7 +1086,7 @@ public abstract class WsRemoteEndpointImplBase implements 
RemoteEndpoint {
                 if (outputBuffer.remaining() == 0) {
                     handler.onResult(SENDRESULT_OK);
                 } else {
-                    endpoint.doWrite(this, blockingWriteTimeoutExpiry, 
outputBuffer);
+                    endpoint.doWrite(this, blocking, writeTimeoutExpiry, 
outputBuffer);
                 }
             } else {
                 handler.onResult(SENDRESULT_OK);
@@ -1075,7 +1098,7 @@ public abstract class WsRemoteEndpointImplBase implements 
RemoteEndpoint {
         public void onResult(SendResult result) {
             if (result.isOK()) {
                 if (outputBuffer.hasRemaining()) {
-                    endpoint.doWrite(this, blockingWriteTimeoutExpiry, 
outputBuffer);
+                    endpoint.doWrite(this, blocking, writeTimeoutExpiry, 
outputBuffer);
                 } else {
                     outputBuffer.clear();
                     write();
@@ -1117,6 +1140,7 @@ public abstract class WsRemoteEndpointImplBase implements 
RemoteEndpoint {
         private final Object closeLock = new Object();
         private volatile boolean closed = false;
         private volatile boolean used = false;
+        private long timeoutExpiry = 0;
 
         WsOutputStream(WsRemoteEndpointImplBase endpoint) {
             this.endpoint = endpoint;
@@ -1192,8 +1216,11 @@ public abstract class WsRemoteEndpointImplBase 
implements RemoteEndpoint {
 
         private void doWrite(boolean last) throws IOException {
             if (used) {
+                if (timeoutExpiry == 0) {
+                    timeoutExpiry = 
endpoint.getTimeoutExpiry(endpoint.getBlockingSendTimeout());
+                }
                 buffer.flip();
-                endpoint.sendMessageBlock(Constants.OPCODE_BINARY, buffer, 
last);
+                endpoint.sendMessageBlockInternal(Constants.OPCODE_BINARY, 
buffer, last, timeoutExpiry);
             }
             endpoint.stateMachine.complete(last);
             buffer.clear();
@@ -1208,6 +1235,7 @@ public abstract class WsRemoteEndpointImplBase implements 
RemoteEndpoint {
         private final Object closeLock = new Object();
         private volatile boolean closed = false;
         private volatile boolean used = false;
+        private long timeoutExpiry = 0;
 
         WsWriter(WsRemoteEndpointImplBase endpoint) {
             this.endpoint = endpoint;
@@ -1270,8 +1298,11 @@ public abstract class WsRemoteEndpointImplBase 
implements RemoteEndpoint {
 
         private void doWrite(boolean last) throws IOException {
             if (used) {
+                if (timeoutExpiry == 0) {
+                    timeoutExpiry = 
endpoint.getTimeoutExpiry(endpoint.getBlockingSendTimeout());
+                }
                 buffer.flip();
-                endpoint.sendMessageBlock(buffer, last);
+                endpoint.sendMessageBlock(buffer, last, timeoutExpiry);
                 buffer.clear();
             } else {
                 endpoint.stateMachine.complete(last);
diff --git a/java/org/apache/tomcat/websocket/WsRemoteEndpointImplClient.java 
b/java/org/apache/tomcat/websocket/WsRemoteEndpointImplClient.java
index 27cb15985a..efed761de5 100644
--- a/java/org/apache/tomcat/websocket/WsRemoteEndpointImplClient.java
+++ b/java/org/apache/tomcat/websocket/WsRemoteEndpointImplClient.java
@@ -26,11 +26,16 @@ import java.util.concurrent.locks.ReentrantLock;
 import jakarta.websocket.SendHandler;
 import jakarta.websocket.SendResult;
 
+import org.apache.juli.logging.Log;
+import org.apache.juli.logging.LogFactory;
+
 /**
  * Client-side implementation of a WebSocket remote endpoint.
  */
 public class WsRemoteEndpointImplClient extends WsRemoteEndpointImplBase {
 
+    private final Log log = 
LogFactory.getLog(WsRemoteEndpointImplClient.class); // must not be static
+
     private final AsyncChannelWrapper channel;
     private final ReentrantLock lock = new ReentrantLock();
 
@@ -50,26 +55,26 @@ public class WsRemoteEndpointImplClient extends 
WsRemoteEndpointImplBase {
 
 
     @Override
-    protected void doWrite(SendHandler handler, long 
blockingWriteTimeoutExpiry, ByteBuffer... data) {
-        long timeout;
+    protected void doWrite(SendHandler handler, boolean blocking, long 
writeTimeoutExpiry, ByteBuffer... data) {
         for (ByteBuffer byteBuffer : data) {
-            if (blockingWriteTimeoutExpiry == -1) {
-                timeout = getSendTimeout();
-                if (timeout < 1) {
-                    timeout = Long.MAX_VALUE;
-                }
+            long timeout;
+            if (writeTimeoutExpiry == Long.MAX_VALUE) {
+                timeout = Long.MAX_VALUE;
             } else {
-                timeout = blockingWriteTimeoutExpiry - 
System.currentTimeMillis();
-                if (timeout < 0) {
-                    SendResult sr = new SendResult(new 
IOException(sm.getString("wsRemoteEndpoint.writeTimeout")));
-                    handler.onResult(sr);
-                    return;
-                }
+                timeout = writeTimeoutExpiry - System.currentTimeMillis();
+            }
+            if (timeout <= 0) {
+                SendResult sr =
+                        new SendResult(new 
IOException(sm.getString("wsRemoteEndpoint.writeTimeout")));
+                handler.onResult(sr);
+                return;
             }
 
             try {
                 channel.write(byteBuffer).get(timeout, TimeUnit.MILLISECONDS);
             } catch (InterruptedException | ExecutionException | 
TimeoutException e) {
+                log.warn(sm.getString("wsRemoteEndpointClient.writeFailed", 
Long.valueOf(writeTimeoutExpiry),
+                        Long.valueOf(timeout)), e);
                 handler.onResult(new SendResult(e));
                 return;
             }
diff --git 
a/java/org/apache/tomcat/websocket/server/WsRemoteEndpointImplServer.java 
b/java/org/apache/tomcat/websocket/server/WsRemoteEndpointImplServer.java
index fdf1df0d08..56d4bcf00c 100644
--- a/java/org/apache/tomcat/websocket/server/WsRemoteEndpointImplServer.java
+++ b/java/org/apache/tomcat/websocket/server/WsRemoteEndpointImplServer.java
@@ -56,6 +56,8 @@ public class WsRemoteEndpointImplServer extends 
WsRemoteEndpointImplBase {
     private final ReentrantLock writeCompletionLock = new ReentrantLock();
     private volatile SendHandler handler = null;
     private volatile ByteBuffer[] buffers = null;
+    private boolean blockingWriteInProgress = false;
+    private boolean blockingWriteTimedOut = false;
 
     private volatile long timeoutExpiry = -1;
 
@@ -158,25 +160,21 @@ public class WsRemoteEndpointImplServer extends 
WsRemoteEndpointImplBase {
 
 
     @Override
-    protected void doWrite(SendHandler handler, long 
blockingWriteTimeoutExpiry, ByteBuffer... buffers) {
+    protected void doWrite(SendHandler handler, boolean block, long 
writeTimeoutExpiry, ByteBuffer... buffers) {
         if (socketWrapper.hasAsyncIO()) {
-            final boolean block = (blockingWriteTimeoutExpiry != -1);
-            long timeout;
-            if (block) {
-                timeout = blockingWriteTimeoutExpiry - 
System.currentTimeMillis();
-                if (timeout <= 0) {
-                    SendResult sr = new SendResult(new 
SocketTimeoutException());
-                    handler.onResult(sr);
-                    return;
-                }
-            } else {
+            long timeout = getTimeout(writeTimeoutExpiry);
+            if (timeout == 0) {
+                SendResult sr = new SendResult(new SocketTimeoutException());
+                handler.onResult(sr);
+                return;
+            }
+            if (!block) {
                 writeCompletionLock.lock();
                 try {
                     this.handler = handler;
-                    timeout = getSendTimeout();
-                    if (timeout > 0) {
+                    timeoutExpiry = writeTimeoutExpiry;
+                    if (writeTimeoutExpiry != Long.MAX_VALUE) {
                         // Register with timeout thread
-                        timeoutExpiry = timeout + System.currentTimeMillis();
                         wsWriteTimeout.register(this);
                     }
                 } finally {
@@ -188,8 +186,8 @@ public class WsRemoteEndpointImplServer extends 
WsRemoteEndpointImplBase {
                         @Override
                         public void completed(Long result, Void attachment) {
                             if (block) {
-                                long timeout = blockingWriteTimeoutExpiry - 
System.currentTimeMillis();
-                                if (timeout <= 0) {
+                                long timeout = getTimeout(writeTimeoutExpiry);
+                                if (timeout == 0) {
                                     failed(new SocketTimeoutException(), null);
                                 } else {
                                     handler.onResult(SENDRESULT_OK);
@@ -211,11 +209,15 @@ public class WsRemoteEndpointImplServer extends 
WsRemoteEndpointImplBase {
                         }
                     }, buffers);
         } else {
-            if (blockingWriteTimeoutExpiry == -1) {
+            if (!block) {
                 writeCompletionLock.lock();
                 try {
                     this.handler = handler;
                     this.buffers = buffers;
+                    timeoutExpiry = writeTimeoutExpiry;
+                    if (writeTimeoutExpiry != Long.MAX_VALUE) {
+                        wsWriteTimeout.register(this);
+                    }
                 } finally {
                     writeCompletionLock.unlock();
                 }
@@ -224,35 +226,72 @@ public class WsRemoteEndpointImplServer extends 
WsRemoteEndpointImplBase {
                 onWritePossible(true);
             } else {
                 // Blocking
+                writeCompletionLock.lock();
+                try {
+                    blockingWriteInProgress = true;
+                    blockingWriteTimedOut = false;
+                    timeoutExpiry = writeTimeoutExpiry;
+                    if (writeTimeoutExpiry != Long.MAX_VALUE) {
+                        wsWriteTimeout.register(this);
+                    }
+                } finally {
+                    writeCompletionLock.unlock();
+                }
+                SendResult sendResult;
                 try {
+                    boolean timedOut = false;
                     for (ByteBuffer buffer : buffers) {
-                        long timeout = blockingWriteTimeoutExpiry - 
System.currentTimeMillis();
-                        if (timeout <= 0) {
-                            SendResult sr = new SendResult(new 
SocketTimeoutException());
-                            handler.onResult(sr);
-                            return;
+                        long timeout = getTimeout(writeTimeoutExpiry);
+                        if (timeout == 0) {
+                            timedOut = true;
+                            break;
                         }
                         socketWrapper.setWriteTimeout(timeout);
                         socketWrapper.write(true, buffer);
                     }
-                    long timeout = blockingWriteTimeoutExpiry - 
System.currentTimeMillis();
-                    if (timeout <= 0) {
-                        SendResult sr = new SendResult(new 
SocketTimeoutException());
-                        handler.onResult(sr);
-                        return;
+                    if (!timedOut) {
+                        long timeout = getTimeout(writeTimeoutExpiry);
+                        if (timeout == 0) {
+                            timedOut = true;
+                        } else {
+                            socketWrapper.setWriteTimeout(timeout);
+                            socketWrapper.flush(true);
+                        }
+                    }
+                    if (timedOut) {
+                        sendResult = new SendResult(new 
SocketTimeoutException());
+                    } else {
+                        sendResult = new SendResult();
                     }
-                    socketWrapper.setWriteTimeout(timeout);
-                    socketWrapper.flush(true);
-                    handler.onResult(SENDRESULT_OK);
                 } catch (IOException ioe) {
-                    SendResult sr = new SendResult(ioe);
-                    handler.onResult(sr);
+                    sendResult = new SendResult(ioe);
+                } finally {
+                    writeCompletionLock.lock();
+                    try {
+                        if (blockingWriteTimedOut) {
+                            sendResult = new SendResult(new 
SocketTimeoutException());
+                        }
+                        blockingWriteInProgress = false;
+                        blockingWriteTimedOut = false;
+                        wsWriteTimeout.unregister(this);
+                    } finally {
+                        writeCompletionLock.unlock();
+                    }
                 }
+                handler.onResult(sendResult);
             }
         }
     }
 
 
+    private static long getTimeout(long writeTimeoutExpiry) {
+        if (writeTimeoutExpiry == Long.MAX_VALUE) {
+            return -1;
+        }
+        return Math.max(0, writeTimeoutExpiry - System.currentTimeMillis());
+    }
+
+
     @Override
     protected void updateStats(long payloadLength) {
         upgradeInfo.addMsgsSent(1);
@@ -299,21 +338,6 @@ public class WsRemoteEndpointImplServer extends 
WsRemoteEndpointImplBase {
             close();
         }
 
-        if (!complete) {
-            writeCompletionLock.lock();
-            try {
-                // The write may have completed or timed out while obtaining 
the lock
-                if (handler != null) {
-                    long timeout = getSendTimeout();
-                    if (timeout > 0) {
-                        timeoutExpiry = timeout + System.currentTimeMillis();
-                        wsWriteTimeout.register(this);
-                    }
-                }
-            } finally {
-                writeCompletionLock.unlock();
-            }
-        }
     }
 
 
@@ -359,7 +383,12 @@ public class WsRemoteEndpointImplServer extends 
WsRemoteEndpointImplBase {
         writeCompletionLock.lock();
         try {
             // Re-check timeout in case of concurrent completion and new write
-            if (handler != null && getTimeoutExpiry() < now) {
+            if (blockingWriteInProgress && getTimeoutExpiry() < now) {
+                blockingWriteInProgress = false;
+                blockingWriteTimedOut = true;
+                wsWriteTimeout.unregister(this);
+                close();
+            } else if (handler != null && getTimeoutExpiry() < now) {
                 wsWriteTimeout.unregister(this);
                 clearHandlerInternal(new SocketTimeoutException(), 
useDispatch);
                 close();
diff --git a/java/org/apache/tomcat/websocket/server/WsWriteTimeout.java 
b/java/org/apache/tomcat/websocket/server/WsWriteTimeout.java
index db6d9963a7..15d9949e1e 100644
--- a/java/org/apache/tomcat/websocket/server/WsWriteTimeout.java
+++ b/java/org/apache/tomcat/websocket/server/WsWriteTimeout.java
@@ -24,7 +24,7 @@ import org.apache.tomcat.websocket.BackgroundProcess;
 import org.apache.tomcat.websocket.BackgroundProcessManager;
 
 /**
- * Provides timeouts for asynchronous web socket writes. On the server side we 
only have access to
+ * Provides timeouts for WebSocket writes. On the server side we only have 
access to
  * {@link jakarta.servlet.ServletOutputStream} and {@link 
jakarta.servlet.ServletInputStream} so there is no way to set
  * a timeout for writes to the client.
  */
diff --git a/test/org/apache/tomcat/websocket/TestPerMessageDeflate.java 
b/test/org/apache/tomcat/websocket/TestPerMessageDeflate.java
index 2798cecbee..74a77d112e 100644
--- a/test/org/apache/tomcat/websocket/TestPerMessageDeflate.java
+++ b/test/org/apache/tomcat/websocket/TestPerMessageDeflate.java
@@ -47,14 +47,16 @@ public class TestPerMessageDeflate {
         perMessageDeflate.setNext(new TesterTransformation());
 
         ByteBuffer bb1 = ByteBuffer.wrap("A".getBytes(StandardCharsets.UTF_8));
-        MessagePart mp1 = new MessagePart(true, 0, Constants.OPCODE_TEXT, bb1, 
null, null, -1);
+        MessagePart mp1 =
+                new MessagePart(true, 0, Constants.OPCODE_TEXT, bb1, null, 
null, false, Long.MAX_VALUE);
 
         List<MessagePart> uncompressedParts1 = new ArrayList<>();
         uncompressedParts1.add(mp1);
         perMessageDeflate.sendMessagePart(uncompressedParts1);
 
         ByteBuffer bb2 = ByteBuffer.wrap("".getBytes(StandardCharsets.UTF_8));
-        MessagePart mp2 = new MessagePart(true, 0, Constants.OPCODE_TEXT, bb2, 
null, null, -1);
+        MessagePart mp2 =
+                new MessagePart(true, 0, Constants.OPCODE_TEXT, bb2, null, 
null, false, Long.MAX_VALUE);
 
         List<MessagePart> uncompressedParts2 = new ArrayList<>();
         uncompressedParts2.add(mp2);
@@ -80,13 +82,17 @@ public class TestPerMessageDeflate {
         byte[] data = new byte[8192];
 
         ByteBuffer bb = ByteBuffer.wrap(data);
-        MessagePart mp = new MessagePart(true, 0, Constants.OPCODE_BINARY, bb, 
null, null, -1);
+        long writeTimeoutExpiry = System.currentTimeMillis() + 10000;
+        MessagePart mp =
+                new MessagePart(true, 0, Constants.OPCODE_BINARY, bb, null, 
null, true, writeTimeoutExpiry);
 
         List<MessagePart> uncompressedParts = new ArrayList<>();
         uncompressedParts.add(mp);
         List<MessagePart> compressedParts = 
perMessageDeflateTx.sendMessagePart(uncompressedParts);
 
         MessagePart compressedPart = compressedParts.get(0);
+        Assert.assertTrue(compressedPart.isBlocking());
+        Assert.assertEquals(writeTimeoutExpiry, 
compressedPart.getWriteTimeoutExpiry());
 
         // Set up the decompression and process the received message
         PerMessageDeflate perMessageDeflateRx = 
PerMessageDeflate.build(preferences, true);
@@ -120,7 +126,7 @@ public class TestPerMessageDeflate {
             Arrays.fill(data, (byte) 0x80);
             List<MessagePart> uncompressedParts = new ArrayList<>();
             uncompressedParts.add(new MessagePart(true, 0, 
Constants.OPCODE_BINARY,
-                    ByteBuffer.wrap(data), null, null, -1));
+                    ByteBuffer.wrap(data), null, null, false, Long.MAX_VALUE));
             MessagePart compressedPart = 
perMessageDeflateTx.sendMessagePart(uncompressedParts).get(0);
 
             // Decompress the way WsFrameBase.processDataBinary does: fill an 
8192
@@ -166,11 +172,13 @@ public class TestPerMessageDeflate {
         // First message part
         byte[] data = new byte[1024];
         ByteBuffer bb = ByteBuffer.wrap(data);
-        MessagePart mp1 = new MessagePart(true, 0, Constants.OPCODE_BINARY, 
bb, null, null, -1);
+        MessagePart mp1 =
+                new MessagePart(true, 0, Constants.OPCODE_BINARY, bb, null, 
null, false, Long.MAX_VALUE);
         uncompressedParts.add(mp1);
 
         // Flush message (replicates result of calling flushBatch()
-        MessagePart mp2 = new MessagePart(true, 0, 
Constants.INTERNAL_OPCODE_FLUSH, null, null, null, -1);
+        MessagePart mp2 = new MessagePart(true, 0, 
Constants.INTERNAL_OPCODE_FLUSH, null, null, null, false,
+                Long.MAX_VALUE);
         uncompressedParts.add(mp2);
 
         List<MessagePart> compressedParts = 
perMessageDeflateTx.sendMessagePart(uncompressedParts);
diff --git a/webapps/docs/changelog.xml b/webapps/docs/changelog.xml
index f7de426c41..f3b4eb56b9 100644
--- a/webapps/docs/changelog.xml
+++ b/webapps/docs/changelog.xml
@@ -257,8 +257,8 @@
         Improve robustness of client handshakes. (remm)
       </fix>
       <fix>
-        Ensure that WebSocket non-blocking write timeouts are not lost if two
-        writes have the same timeout. (markt)
+        Ensure that WebSocket write timeouts apply to the complete message and
+        are not lost if two writes have the same timeout. (markt)
       </fix>
     </changelog>
   </subsection>


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to