Ensure SocketStream recalculates the timeout on each read (#1299)

JAVA-5298

Author: rozza

driver-core/src/main/com/mongodb/internal/TimeoutContext.java

@@ -213,6 +213,9 @@ public TimeoutContext withAdditionalReadTimeout(final int additionalReadTimeout)
 
     private long timeoutRemainingMS() {
         assertNotNull(timeout);
+        if (timeout.hasExpired()) {
+            throw createMongoTimeoutException("The operation timeout has expired.");
+        }
         return timeout.isInfinite() ? 0 : timeout.remaining(MILLISECONDS);
     }
 

driver-core/src/main/com/mongodb/internal/connection/SocketStream.java

@@ -171,16 +171,16 @@ public void write(final List<ByteBuf> buffers, final OperationContext operationC
 
     @Override
     public ByteBuf read(final int numBytes, final OperationContext operationContext) throws IOException {
-        int readTimeoutMS = (int) operationContext.getTimeoutContext().getReadTimeoutMS();
-        if (readTimeoutMS > 0) {
-            socket.setSoTimeout(readTimeoutMS);
-        }
         try {
             ByteBuf buffer = bufferProvider.getBuffer(numBytes);
             try {
                 int totalBytesRead = 0;
                 byte[] bytes = buffer.array();
                 while (totalBytesRead < buffer.limit()) {
+                    int readTimeoutMS = (int) operationContext.getTimeoutContext().getReadTimeoutMS();
+                    if (readTimeoutMS > 0) {
+                        socket.setSoTimeout(readTimeoutMS);
+                    }
                     int bytesRead = inputStream.read(bytes, totalBytesRead, buffer.limit() - totalBytesRead);
                     if (bytesRead == -1) {
                         throw new MongoSocketReadException("Prematurely reached end of stream", getAddress());

driver-core/src/test/unit/com/mongodb/internal/TimeoutContextTest.java

@@ -32,6 +32,7 @@
 import static com.mongodb.ClusterFixture.sleep;
 import static java.util.Arrays.asList;
 import static org.junit.jupiter.api.Assertions.assertAll;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -197,6 +198,30 @@ Collection<DynamicTest> timeoutContextTest() {
                             () -> assertTrue(smallTimeout.hasExpired())
                     );
                 }),
+                dynamicTest("throws when calculating timeout if expired", () -> {
+                    TimeoutContext smallTimeout = new TimeoutContext(TIMEOUT_SETTINGS.withTimeoutMS(1));
+                    TimeoutContext longTimeout =
+                            new TimeoutContext(TIMEOUT_SETTINGS.withTimeoutMS(9999999));
+                    TimeoutContext noTimeout = new TimeoutContext(TIMEOUT_SETTINGS);
+                    sleep(100);
+                    assertAll(
+                            () -> assertThrows(MongoOperationTimeoutException.class, smallTimeout::getReadTimeoutMS),
+                            () -> assertThrows(MongoOperationTimeoutException.class, smallTimeout::getWriteTimeoutMS),
+                            () -> assertThrows(MongoOperationTimeoutException.class, smallTimeout::getMaxTimeMS),
+                            () -> assertThrows(MongoOperationTimeoutException.class, smallTimeout::getMaxCommitTimeMS),
+                            () -> assertThrows(MongoOperationTimeoutException.class, () -> smallTimeout.timeoutOrAlternative(1)),
+                            () -> assertDoesNotThrow(longTimeout::getReadTimeoutMS),
+                            () -> assertDoesNotThrow(longTimeout::getWriteTimeoutMS),
+                            () -> assertDoesNotThrow(longTimeout::getMaxTimeMS),
+                            () -> assertDoesNotThrow(longTimeout::getMaxCommitTimeMS),
+                            () -> assertDoesNotThrow(() -> longTimeout.timeoutOrAlternative(1)),
+                            () -> assertDoesNotThrow(noTimeout::getReadTimeoutMS),
+                            () -> assertDoesNotThrow(noTimeout::getWriteTimeoutMS),
+                            () -> assertDoesNotThrow(noTimeout::getMaxTimeMS),
+                            () -> assertDoesNotThrow(noTimeout::getMaxCommitTimeMS),
+                            () -> assertDoesNotThrow(() -> noTimeout.timeoutOrAlternative(1))
+                    );
+                }),
                 dynamicTest("validates minRoundTripTime for maxTimeMS", () -> {
                     Supplier<TimeoutContext> supplier = () -> new TimeoutContext(TIMEOUT_SETTINGS.withTimeoutMS(100));
                     assertAll(