chore: Improve task cancellation.

nstdio · Dec 29, 2022 · b7b8223 · b7b8223
1 parent 80117ab
commit b7b8223
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 14 deletions.
diff --git a/src/main/java/io/github/nstdio/http/ext/ByteArraySubscription.java b/src/main/java/io/github/nstdio/http/ext/ByteArraySubscription.java
@@ -19,14 +19,15 @@
 import java.nio.ByteBuffer;
 import java.util.List;
 import java.util.concurrent.Executor;
-import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Flow.Subscriber;
 import java.util.concurrent.Flow.Subscription;
 import java.util.concurrent.Future;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
 import java.util.function.Supplier;
 
+import static java.util.concurrent.CompletableFuture.runAsync;
+
 class ByteArraySubscription<T> implements Subscription {
   private final Subscriber<T> subscriber;
   private final Executor executor;
@@ -35,7 +36,7 @@ class ByteArraySubscription<T> implements Subscription {
   private final Supplier<byte[]> bytes;
 
   private final AtomicBoolean completed = new AtomicBoolean(false);
-  private Future<?> result;
+  Future<?> result;
 
   ByteArraySubscription(Subscriber<T> subscriber, Executor executor, Supplier<byte[]> bytes, Function<byte[], T> mapper) {
     this.subscriber = subscriber;
@@ -83,11 +84,7 @@ public void cancel() {
   }
 
   private void submit(Runnable r) {
-    if (executor instanceof ExecutorService) {
-      result = ((ExecutorService) executor).submit(r);
-    } else {
-      executor.execute(r);
-    }
+    result = runAsync(r, executor);
   }
 
   private enum DirectExecutor implements Executor {

diff --git a/src/test/kotlin/io/github/nstdio/http/ext/ByteArraySubscriptionTest.kt b/src/test/kotlin/io/github/nstdio/http/ext/ByteArraySubscriptionTest.kt
@@ -16,6 +16,7 @@
 package io.github.nstdio.http.ext
 
 import io.kotest.assertions.timing.eventually
+import io.kotest.matchers.shouldBe
 import io.kotest.property.Arb
 import io.kotest.property.arbitrary.next
 import io.kotest.property.arbitrary.string
@@ -28,7 +29,6 @@ import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.extension.ExtendWith
 import org.mockito.ArgumentMatchers
 import org.mockito.ArgumentMatchers.any
-import org.mockito.BDDMockito.given
 import org.mockito.Mock
 import org.mockito.Mockito.inOrder
 import org.mockito.Mockito.mock
@@ -42,7 +42,6 @@ import java.nio.ByteBuffer
 import java.nio.charset.StandardCharsets
 import java.util.concurrent.ExecutorService
 import java.util.concurrent.ForkJoinPool
-import java.util.concurrent.Future
 import kotlin.time.Duration.Companion.seconds
 
 @ExtendWith(MockitoExtension::class)
@@ -142,10 +141,6 @@ internal class ByteArraySubscriptionTest {
   fun `Should cancel result`() {
     //given
     val mockExecutor = mock(ExecutorService::class.java)
-    val mockFuture = mock(Future::class.java)
-
-    given(mockExecutor.submit(any(Runnable::class.java)))
-      .willReturn(mockFuture)
 
     val bytes = Arb.byteArray(8).next()
     val item = listOf(bytes.toBuffer())
@@ -156,7 +151,7 @@ internal class ByteArraySubscriptionTest {
     subscription.cancel()
 
     //then
-    verify(mockFuture).cancel(false)
+    subscription.result.isCancelled shouldBe true
   }
 
   private fun runAsyncAwait(block: suspend CoroutineScope.() -> Unit) {