Review comments

Author: dagnir

http-clients/apache5-client/src/main/java/software/amazon/awssdk/http/apache5/internal/utils/CancelOnInterruptWrapper.java

@@ -48,17 +48,36 @@ public boolean isDone() {
         return f.isDone();
     }
 
+
+    // This method attempts to cancel the wrapped future if the thread is interrupted while blocked on get(). This is done by
+    // attempting to cancel() the future when InterruptedException is thrown. If the the cancel() is unsuccessful (i.e.
+    // the future is completed either successfully or exceptionally), then get the result if present and return it.
     @Override
     public ResultT get() throws InterruptedException, ExecutionException {
-        return f.get();
+        try {
+            return f.get();
+        } catch (InterruptedException ie) {
+            if (!cancel(true)) {
+                try {
+                    // We couldn't cancel so the result will be available or it failed
+                    ResultT entry = f.get();
+                    Thread.currentThread().interrupt();
+                    return entry;
+                } catch (CancellationException | InterruptedException | ExecutionException e) {
+                    // no-op, let it fall through to throwing the original interrupted exception
+                }
+            }
+            throw ie;
+        }
     }
 
+
     // This method attempts to cancel the wrapped future if the thread is interrupted while blocked on get(). This is done by
     // attempting to cancel() the future when InterruptedException is thrown. If the the cancel() is unsuccessful (i.e.
     // the future is completed either successfully or exceptionally), then get the result if present and return it.
     @Override
     public ResultT get(long timeout, TimeUnit unit)
-            throws InterruptedException, ExecutionException, TimeoutException {
+        throws InterruptedException, ExecutionException, TimeoutException {
         try {
             return f.get(timeout, unit);
         } catch (InterruptedException ie) {
@@ -75,4 +94,4 @@ public ResultT get(long timeout, TimeUnit unit)
             throw ie;
         }
     }
-}
+}

…l/utils/CancelOnInterruptFutureTest.java …/utils/CancelOnInterruptWrapperTest.javahttp-clients/apache5-client/src/test/java/software/amazon/awssdk/http/apache5/internal/utils/CancelOnInterruptFutureTest.java renamed to http-clients/apache5-client/src/test/java/software/amazon/awssdk/http/apache5/internal/utils/CancelOnInterruptWrapperTest.java http-clients/apache5-client/src/test/java/software/amazon/awssdk/http/apache5/internal/utils/CancelOnInterruptFutureTest.java renamed to http-clients/apache5-client/src/test/java/software/amazon/awssdk/http/apache5/internal/utils/CancelOnInterruptWrapperTest.java

@@ -30,17 +30,24 @@
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
-public class CancelOnInterruptFutureTest {
+public class CancelOnInterruptWrapperTest {
     private Future<String> mockDelegate;
 
     @BeforeEach
     void setup() {
         mockDelegate = mock(Future.class);
     }
 
+    @AfterEach
+    void teardown() {
+        Thread.interrupted(); // clear the flag if it was set by the last test
+    }
+
     @Test
     void cancel_callsDelegate() {
         CancelOnInterruptWrapper<String> wrapper = new CancelOnInterruptWrapper<>(mockDelegate);
@@ -77,7 +84,7 @@ void getTimeout_callsDelegate() throws ExecutionException, InterruptedException,
     }
 
     @Test
-    void get_interrupted_cancelSuccessful_throws() throws ExecutionException, InterruptedException, TimeoutException {
+    void getTimeout_interrupted_cancelSuccessful_throws() throws ExecutionException, InterruptedException, TimeoutException {
         when(mockDelegate.get(anyLong(), any(TimeUnit.class))).thenThrow(new InterruptedException("interrupt"));
         when(mockDelegate.cancel(eq(true))).thenReturn(true);
 
@@ -94,7 +101,8 @@ void get_interrupted_cancelSuccessful_throws() throws ExecutionException, Interr
     }
 
     @Test
-    void get_interrupted_cancelUnsuccessful_returnsEntry() throws ExecutionException, InterruptedException, TimeoutException {
+    void getTimeout_interrupted_cancelUnsuccessful_returnsEntry() throws ExecutionException, InterruptedException,
+                                                                         TimeoutException {
         String result = "hello there";
 
         when(mockDelegate.get(anyLong(), any(TimeUnit.class))).thenThrow(new InterruptedException("interrupt"));
@@ -107,7 +115,8 @@ void get_interrupted_cancelUnsuccessful_returnsEntry() throws ExecutionException
     }
 
     @Test
-    void get_interrupted_cancelUnsuccessful_getUnsuccessful_rethrowsOriginalIe() throws ExecutionException, InterruptedException, TimeoutException {
+    void getTimeout_interrupted_cancelUnsuccessful_getUnsuccessful_rethrowsOriginalIe() throws ExecutionException,
+                                                                                                  InterruptedException, TimeoutException {
         InterruptedException interrupt = new InterruptedException("interrupt");
 
         when(mockDelegate.get(anyLong(), any(TimeUnit.class))).thenThrow(interrupt);
@@ -118,4 +127,93 @@ void get_interrupted_cancelUnsuccessful_getUnsuccessful_rethrowsOriginalIe() thr
 
         assertThatThrownBy(() -> wrapper.get(1, TimeUnit.SECONDS)).isSameAs(interrupt);
     }
+
+    @Test
+    void get_interrupted_cancelSuccessful_throws() throws ExecutionException, InterruptedException, TimeoutException {
+        when(mockDelegate.get()).thenThrow(new InterruptedException("interrupt"));
+        when(mockDelegate.cancel(eq(true))).thenReturn(true);
+
+        CancelOnInterruptWrapper<String> wrapper = new CancelOnInterruptWrapper<>(mockDelegate);
+
+        assertThatThrownBy(wrapper::get)
+            .isInstanceOf(InterruptedException.class)
+            .hasMessage("interrupt");
+
+        verify(mockDelegate).get();
+        verify(mockDelegate).cancel(eq(true));
+
+        verifyNoMoreInteractions(mockDelegate);
+    }
+
+    @Test
+    void get_interrupted_cancelUnsuccessful_returnsEntry() throws ExecutionException, InterruptedException, TimeoutException {
+        String result = "hello there";
+
+        AtomicBoolean first = new AtomicBoolean(true);
+        when(mockDelegate.get()).thenAnswer(i -> {
+           if (first.compareAndSet(true, false)) {
+               throw new InterruptedException("interrupt");
+           }
+           return result;
+        });
+        when(mockDelegate.cancel(eq(true))).thenReturn(false);
+
+        CancelOnInterruptWrapper<String> wrapper = new CancelOnInterruptWrapper<>(mockDelegate);
+
+        assertThat(wrapper.get()).isEqualTo(result);
+    }
+
+    @Test
+    void get_interrupted_cancelUnsuccessful_getUnsuccessful_rethrowsOriginalIe() throws ExecutionException, InterruptedException, TimeoutException {
+        InterruptedException interrupt = new InterruptedException("interrupt");
+
+        AtomicBoolean first = new AtomicBoolean(true);
+        when(mockDelegate.get()).thenAnswer(i -> {
+            if (first.compareAndSet(true, false)) {
+                throw interrupt;
+            }
+            throw new CancellationException("cancelled");
+        });
+        when(mockDelegate.cancel(eq(true))).thenReturn(false);
+
+        CancelOnInterruptWrapper<String> wrapper = new CancelOnInterruptWrapper<>(mockDelegate);
+
+        assertThatThrownBy(wrapper::get).isSameAs(interrupt);
+    }
+
+    @Test
+    void get_interrupted_cancelUnsuccessful_cancelUnsuccessful_preservesInterruptedFlag() throws ExecutionException, InterruptedException {
+        String result = "hello there";
+
+        AtomicBoolean first = new AtomicBoolean(true);
+        when(mockDelegate.get()).thenAnswer(i -> {
+            if (first.compareAndSet(true, false)) {
+                throw new InterruptedException("interrupt");
+            }
+            return result;
+        });
+        when(mockDelegate.cancel(eq(true))).thenReturn(false);
+
+        CancelOnInterruptWrapper<String> wrapper = new CancelOnInterruptWrapper<>(mockDelegate);
+
+        wrapper.get();
+
+        assertThat(Thread.interrupted()).isTrue();
+    }
+
+    @Test
+    void getTimeout_interrupted_cancelUnsuccessful_preservesInterruptedFlag() throws ExecutionException, InterruptedException,
+                                                                         TimeoutException {
+        String result = "hello there";
+
+        when(mockDelegate.get(anyLong(), any(TimeUnit.class))).thenThrow(new InterruptedException("interrupt"));
+        when(mockDelegate.cancel(eq(true))).thenReturn(false);
+        when(mockDelegate.get()).thenReturn(result);
+
+        CancelOnInterruptWrapper<String> wrapper = new CancelOnInterruptWrapper<>(mockDelegate);
+
+        wrapper.get(1, TimeUnit.SECONDS);
+
+        assertThat(Thread.interrupted()).isTrue();
+    }
 }