CancellationScheduler finishes even if not started

core/src/main/java/com/linecorp/armeria/client/ClientRequestContextBuilder.java

@@ -20,8 +20,6 @@
 
 import java.net.InetSocketAddress;
 import java.net.URI;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
 
 import javax.net.ssl.SSLSession;
 
@@ -116,17 +114,8 @@ public ClientRequestContext build() {
         if (timedOut()) {
             responseCancellationScheduler = CancellationScheduler.finished(false);
         } else {
-            responseCancellationScheduler = CancellationScheduler.of(0);
-            final CountDownLatch latch = new CountDownLatch(1);
-            eventLoop().execute(() -> {
-                responseCancellationScheduler.init(eventLoop(), noopCancellationTask, 0, /* server */ false);
-                latch.countDown();
-            });
-
-            try {
-                latch.await(1000, TimeUnit.MILLISECONDS);
-            } catch (InterruptedException ignored) {
-            }
+            responseCancellationScheduler = CancellationScheduler.ofClient(0);
+            responseCancellationScheduler.initAndStart(eventLoop(), noopCancellationTask);
         }
 
         final DefaultClientRequestContext ctx = new DefaultClientRequestContext(

core/src/main/java/com/linecorp/armeria/client/HttpResponseWrapper.java

@@ -291,9 +291,7 @@ void initTimeout() {
         if (ctxExtension != null) {
             final CancellationScheduler responseCancellationScheduler =
                     ctxExtension.responseCancellationScheduler();
-            responseCancellationScheduler.init(
-                    ctx.eventLoop(), newCancellationTask(),
-                    TimeUnit.MILLISECONDS.toNanos(responseTimeoutMillis), /* server */ false);
+            responseCancellationScheduler.start(newCancellationTask());
         }
     }
 

core/src/main/java/com/linecorp/armeria/internal/client/DefaultClientRequestContext.java

@@ -180,14 +180,15 @@ private static SessionProtocol desiredSessionProtocol(SessionProtocol protocol,
      *                               e.g. {@code System.currentTimeMillis() * 1000}.
      */
     public DefaultClientRequestContext(
-            @Nullable EventLoop eventLoop, MeterRegistry meterRegistry, SessionProtocol sessionProtocol,
+            EventLoop eventLoop, MeterRegistry meterRegistry, SessionProtocol sessionProtocol,
             RequestId id, HttpMethod method, RequestTarget reqTarget,
             ClientOptions options, @Nullable HttpRequest req, @Nullable RpcRequest rpcReq,
             RequestOptions requestOptions, CancellationScheduler responseCancellationScheduler,
             long requestStartTimeNanos, long requestStartTimeMicros) {
-        this(eventLoop, meterRegistry, sessionProtocol,
+        this(requireNonNull(eventLoop, "eventLoop"), meterRegistry, sessionProtocol,
              id, method, reqTarget, options, req, rpcReq, requestOptions, serviceRequestContext(),
-             responseCancellationScheduler, requestStartTimeNanos, requestStartTimeMicros);
+             requireNonNull(responseCancellationScheduler, "responseCancellationScheduler"),
+             requestStartTimeNanos, requestStartTimeMicros);
     }
 
     /**
@@ -226,6 +227,9 @@ private DefaultClientRequestContext(
               firstNonNull(requestOptions.exchangeType(), ExchangeType.BIDI_STREAMING),
               requestAutoAbortDelayMillis(options, requestOptions), req, rpcReq,
               getAttributes(root), options.contextHook());
+        assert (eventLoop == null && responseCancellationScheduler == null) ||
+               (eventLoop != null && responseCancellationScheduler != null)
+                : "'eventLoop' and 'responseCancellationScheduler' should be both null or non-null";
 
         this.eventLoop = eventLoop;
         this.options = requireNonNull(options, "options");
@@ -240,7 +244,8 @@ private DefaultClientRequestContext(
                 responseTimeoutMillis = options().responseTimeoutMillis();
             }
             this.responseCancellationScheduler =
-                    CancellationScheduler.of(TimeUnit.MILLISECONDS.toNanos(responseTimeoutMillis));
+                    CancellationScheduler.ofClient(TimeUnit.MILLISECONDS.toNanos(responseTimeoutMillis));
+            // the cancellationScheduler is not initialized here since the eventLoop is guaranteed to be null
         } else {
             this.responseCancellationScheduler = responseCancellationScheduler;
         }
@@ -427,6 +432,7 @@ private void acquireEventLoop(EndpointGroup endpointGroup) {
                     options().factory().acquireEventLoop(sessionProtocol(), endpointGroup, endpoint);
             eventLoop = releasableEventLoop.get();
             log.whenComplete().thenAccept(unused -> releasableEventLoop.release());
+            responseCancellationScheduler.init(eventLoop());
         }
     }
 
@@ -517,7 +523,7 @@ private DefaultClientRequestContext(DefaultClientRequestContext ctx,
         log = RequestLog.builder(this);
         log.startRequest();
         responseCancellationScheduler =
-                CancellationScheduler.of(TimeUnit.MILLISECONDS.toNanos(ctx.responseTimeoutMillis()));
+                CancellationScheduler.ofClient(TimeUnit.MILLISECONDS.toNanos(ctx.responseTimeoutMillis()));
         writeTimeoutMillis = ctx.writeTimeoutMillis();
         maxResponseLength = ctx.maxResponseLength();
 
@@ -534,6 +540,7 @@ private DefaultClientRequestContext(DefaultClientRequestContext ctx,
         // the root context.
         if (endpoint == null || ctx.endpoint() == endpoint && ctx.log.children().isEmpty()) {
             eventLoop = ctx.eventLoop().withoutContext();
+            responseCancellationScheduler.init(eventLoop());
         } else {
             acquireEventLoop(endpoint);
         }

core/src/main/java/com/linecorp/armeria/internal/common/CancellationScheduler.java

@@ -18,17 +18,19 @@
 
 import java.util.concurrent.CompletableFuture;
 
-import com.google.common.annotations.VisibleForTesting;
-
 import com.linecorp.armeria.common.annotation.Nullable;
 import com.linecorp.armeria.common.util.TimeoutMode;
 
 import io.netty.util.concurrent.EventExecutor;
 
 public interface CancellationScheduler {
 
-    static CancellationScheduler of(long timeoutNanos) {
-        return new DefaultCancellationScheduler(timeoutNanos);
+    static CancellationScheduler ofClient(long timeoutNanos) {
+        return new DefaultCancellationScheduler(timeoutNanos, false);
+    }
+
+    static CancellationScheduler ofServer(long timeoutNanos) {
+        return new DefaultCancellationScheduler(timeoutNanos, true);
     }
 
     /**
@@ -59,7 +61,11 @@ public boolean canSchedule() {
         public void run(Throwable cause) { /* no-op */ }
     };
 
-    void init(EventExecutor eventLoop, CancellationTask task, long timeoutNanos, boolean server);
+    void initAndStart(EventExecutor eventLoop, CancellationTask task);
+
+    void init(EventExecutor eventLoop);
+
+    void start(CancellationTask task);
 
     void clearTimeout();
 
@@ -89,9 +95,6 @@ public void run(Throwable cause) { /* no-op */ }
     @Deprecated
     CompletableFuture<Void> whenTimedOut();
 
-    @VisibleForTesting
-    boolean isInitialized();
-
     enum State {
         INIT,
         INACTIVE,

core/src/main/java/com/linecorp/armeria/internal/common/DefaultCancellationScheduler.java

@@ -64,6 +64,10 @@ final class DefaultCancellationScheduler implements CancellationScheduler {
     private static final AtomicLongFieldUpdater<DefaultCancellationScheduler> pendingTimeoutNanosUpdater =
             AtomicLongFieldUpdater.newUpdater(DefaultCancellationScheduler.class, "pendingTimeoutNanos");
 
+    private static final AtomicReferenceFieldUpdater<DefaultCancellationScheduler, EventExecutor>
+            eventLoopUpdater = AtomicReferenceFieldUpdater.newUpdater(
+            DefaultCancellationScheduler.class, EventExecutor.class, "eventLoop");
+
     private static final Runnable noopPendingTask = () -> {
     };
 
@@ -74,7 +78,7 @@ final class DefaultCancellationScheduler implements CancellationScheduler {
     private long timeoutNanos;
     private long startTimeNanos;
     @Nullable
-    private EventExecutor eventLoop;
+    private volatile EventExecutor eventLoop;
     @Nullable
     private CancellationTask task;
     @Nullable
@@ -91,42 +95,61 @@ final class DefaultCancellationScheduler implements CancellationScheduler {
     private volatile TimeoutFuture whenTimedOut;
     @SuppressWarnings("FieldMayBeFinal")
     private volatile long pendingTimeoutNanos;
-    private boolean server;
+    private final boolean server;
     @Nullable
     private Throwable cause;
 
+    @VisibleForTesting
     DefaultCancellationScheduler(long timeoutNanos) {
+        this(timeoutNanos, true);
+    }
+
+    DefaultCancellationScheduler(long timeoutNanos, boolean server) {
         this.timeoutNanos = timeoutNanos;
         pendingTimeoutNanos = timeoutNanos;
+        this.server = server;
     }
 
     /**
      * Initializes this {@link DefaultCancellationScheduler}.
      */
     @Override
-    public void init(EventExecutor eventLoop, CancellationTask task, long timeoutNanos, boolean server) {
+    public void initAndStart(EventExecutor eventLoop, CancellationTask task) {
+        init(eventLoop);
         if (!eventLoop.inEventLoop()) {
-            eventLoop.execute(() -> init0(eventLoop, task, timeoutNanos, server));
+            eventLoop.execute(() -> start(task));
         } else {
-            init0(eventLoop, task, timeoutNanos, server);
+            start(task);
+        }
+    }
+
+    @Override
+    public void init(EventExecutor eventLoop) {
+        if (!eventLoopUpdater.compareAndSet(this, null, eventLoop)) {
+            throw new IllegalStateException("Can't init() more than once");
         }
     }
 
-    private void init0(EventExecutor eventLoop, CancellationTask task, long timeoutNanos, boolean server) {
-        if (state != State.INIT) {
+    @Override
+    public void start(CancellationTask task) {
+        assert eventLoop != null;
+        assert eventLoop.inEventLoop();
+        if (isFinished()) {
+            assert cause != null;
+            task.run(cause);
             return;
         }
-        this.eventLoop = eventLoop;
-        this.task = task;
-        if (timeoutNanos > 0) {
-            this.timeoutNanos = timeoutNanos;
+        if (this.task != null) {
+            // just replace the task
+            this.task = task;
+            return;
         }
-        this.server = server;
+        this.task = task;
         startTimeNanos = System.nanoTime();
-        if (this.timeoutNanos != 0) {
+        if (timeoutNanos != 0) {
             state = State.SCHEDULED;
             scheduledFuture =
-                    eventLoop.schedule(() -> invokeTask(null), this.timeoutNanos, NANOSECONDS);
+                    eventLoop.schedule(() -> invokeTask(null), timeoutNanos, NANOSECONDS);
         } else {
             state = State.INACTIVE;
         }
@@ -312,14 +335,16 @@ public void finishNow(@Nullable Throwable cause) {
         if (isFinishing()) {
             return;
         }
+        assert eventLoop != null;
+        if (!eventLoop.inEventLoop()) {
+            eventLoop.execute(() -> finishNow(cause));
+            return;
+        }
         if (isInitialized()) {
-            if (eventLoop.inEventLoop()) {
-                finishNow0(cause);
-            } else {
-                eventLoop.execute(() -> finishNow0(cause));
-            }
+            finishNow0(cause);
         } else {
-            addPendingTask(() -> finishNow0(cause));
+            start(noopCancellationTask);
+            finishNow0(cause);
         }
     }
 
@@ -430,9 +455,7 @@ public CompletableFuture<Void> whenTimedOut() {
         }
     }
 
-    @Override
-    @VisibleForTesting
-    public boolean isInitialized() {
+    private boolean isInitialized() {
         return pendingTask == noopPendingTask && eventLoop != null;
     }
 
@@ -531,9 +554,8 @@ void doComplete() {
     }
 
     private static CancellationScheduler finished0(boolean server) {
-        final CancellationScheduler cancellationScheduler = CancellationScheduler.of(0);
-        cancellationScheduler
-                .init(ImmediateEventExecutor.INSTANCE, noopCancellationTask, 0, server);
+        final CancellationScheduler cancellationScheduler = new DefaultCancellationScheduler(0, server);
+        cancellationScheduler.initAndStart(ImmediateEventExecutor.INSTANCE, noopCancellationTask);
         cancellationScheduler.finishNow();
         return cancellationScheduler;
     }

core/src/main/java/com/linecorp/armeria/internal/common/NoopCancellationScheduler.java

@@ -37,7 +37,15 @@ private NoopCancellationScheduler() {
     }
 
     @Override
-    public void init(EventExecutor eventLoop, CancellationTask task, long timeoutNanos, boolean server) {
+    public void initAndStart(EventExecutor eventLoop, CancellationTask task) {
+    }
+
+    @Override
+    public void init(EventExecutor eventLoop) {
+    }
+
+    @Override
+    public void start(CancellationTask task) {
     }
 
     @Override
@@ -100,9 +108,4 @@ public CompletableFuture<Void> whenTimingOut() {
     public CompletableFuture<Void> whenTimedOut() {
         return VOID_FUTURE;
     }
-
-    @Override
-    public boolean isInitialized() {
-        return false;
-    }
 }

core/src/main/java/com/linecorp/armeria/internal/server/DefaultServiceRequestContext.java

@@ -177,7 +177,8 @@ public DefaultServiceRequestContext(
             this.requestCancellationScheduler = requestCancellationScheduler;
         } else {
             this.requestCancellationScheduler =
-                    CancellationScheduler.of(TimeUnit.MILLISECONDS.toNanos(cfg.requestTimeoutMillis()));
+                    CancellationScheduler.ofServer(TimeUnit.MILLISECONDS.toNanos(cfg.requestTimeoutMillis()));
+            this.requestCancellationScheduler.init(eventLoop());
         }
         this.sslSession = sslSession;
         this.proxiedAddresses = requireNonNull(proxiedAddresses, "proxiedAddresses");

core/src/main/java/com/linecorp/armeria/server/AbstractHttpResponseHandler.java

@@ -243,8 +243,7 @@ final void maybeWriteAccessLog() {
      */
     final void scheduleTimeout() {
         // Schedule the initial request timeout with the timeoutNanos in the CancellationScheduler
-        reqCtx.requestCancellationScheduler().init(reqCtx.eventLoop(), newCancellationTask(),
-                                                   0, /* server */ true);
+        reqCtx.requestCancellationScheduler().start(newCancellationTask());
     }
 
     /**

core/src/main/java/com/linecorp/armeria/server/ServiceRequestContextBuilder.java

@@ -22,8 +22,6 @@
 import java.net.InetSocketAddress;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 
 import javax.net.ssl.SSLSession;
@@ -234,17 +232,8 @@ public ServiceRequestContext build() {
         if (timedOut()) {
             requestCancellationScheduler = CancellationScheduler.finished(true);
         } else {
-            requestCancellationScheduler = CancellationScheduler.of(0);
-            final CountDownLatch latch = new CountDownLatch(1);
-            eventLoop().execute(() -> {
-                requestCancellationScheduler.init(eventLoop(), noopCancellationTask, 0, /* server */ true);
-                latch.countDown();
-            });
-
-            try {
-                latch.await(1000, TimeUnit.MILLISECONDS);
-            } catch (InterruptedException ignored) {
-            }
+            requestCancellationScheduler = CancellationScheduler.ofServer(0);
+            requestCancellationScheduler.initAndStart(eventLoop(), noopCancellationTask);
         }
 
         // Build the context with the properties set by a user and the fake objects.

core/src/test/java/com/linecorp/armeria/internal/client/DefaultClientRequestContextTest.java

@@ -300,7 +300,7 @@ private static DefaultClientRequestContext newContext(ClientOptions clientOption
         return new DefaultClientRequestContext(
                 mock(EventLoop.class), NoopMeterRegistry.get(), SessionProtocol.H2C,
                 RequestId.random(), HttpMethod.POST, reqTarget, clientOptions, httpRequest,
-                null, RequestOptions.of(), CancellationScheduler.of(0), System.nanoTime(),
+                null, RequestOptions.of(), CancellationScheduler.ofClient(0), System.nanoTime(),
                 SystemInfo.currentTimeMicros());
     }