From 06012a1e86657707b5878b9cf93f5aff2ebd0a90 Mon Sep 17 00:00:00 2001
From: jrhee17 <guins_j@guins.org>
Date: Thu, 1 Aug 2024 14:32:47 +0900
Subject: [PATCH] address comments by @ikhoon

---
 .../armeria/internal/common/CancellationScheduler.java      | 3 ++-
 .../internal/common/DefaultCancellationScheduler.java       | 4 ++--
 .../com/linecorp/armeria/common/ContextPushHookTest.java    | 4 +---
 .../armeria/internal/common/CancellationSchedulerTest.java  | 6 +++---
 4 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/core/src/main/java/com/linecorp/armeria/internal/common/CancellationScheduler.java b/core/src/main/java/com/linecorp/armeria/internal/common/CancellationScheduler.java
index 2145907e50f..04f40a680bb 100644
--- a/core/src/main/java/com/linecorp/armeria/internal/common/CancellationScheduler.java
+++ b/core/src/main/java/com/linecorp/armeria/internal/common/CancellationScheduler.java
@@ -121,13 +121,14 @@ default void finishNow() {
 
     enum State {
         INIT,
-        PENDING,
+        SCHEDULED,
         FINISHED,
     }
 
     /**
      * A cancellation task invoked by the scheduler when its timeout exceeds or invoke by the user.
      */
+    @FunctionalInterface
     interface CancellationTask {
         /**
          * Returns {@code true} if the cancellation task can be scheduled.
diff --git a/core/src/main/java/com/linecorp/armeria/internal/common/DefaultCancellationScheduler.java b/core/src/main/java/com/linecorp/armeria/internal/common/DefaultCancellationScheduler.java
index e5792084926..39dda7ce3e4 100644
--- a/core/src/main/java/com/linecorp/armeria/internal/common/DefaultCancellationScheduler.java
+++ b/core/src/main/java/com/linecorp/armeria/internal/common/DefaultCancellationScheduler.java
@@ -123,11 +123,11 @@ public void start() {
             if (state != State.INIT) {
                 return;
             }
-            state = State.PENDING;
+            state = State.SCHEDULED;
             startTimeNanos = ticker.read();
             if (timeoutMode == TimeoutMode.SET_FROM_NOW) {
                 final long elapsedTimeNanos = startTimeNanos - setFromNowStartNanos;
-                timeoutNanos = LongMath.saturatedSubtract(timeoutNanos, elapsedTimeNanos);
+                timeoutNanos = Long.max(LongMath.saturatedSubtract(timeoutNanos, elapsedTimeNanos), 0);
             }
             if (timeoutNanos != Long.MAX_VALUE) {
                 scheduledFuture = eventLoop().schedule(() -> invokeTask(null), timeoutNanos, NANOSECONDS);
diff --git a/core/src/test/java/com/linecorp/armeria/common/ContextPushHookTest.java b/core/src/test/java/com/linecorp/armeria/common/ContextPushHookTest.java
index 6d76d7b4dc7..acd776f8d88 100644
--- a/core/src/test/java/com/linecorp/armeria/common/ContextPushHookTest.java
+++ b/core/src/test/java/com/linecorp/armeria/common/ContextPushHookTest.java
@@ -135,9 +135,7 @@ void shouldRunHooksWhenContextIsPushed() {
             hookEvents.clear();
             response = client.get("http://foo.com:" + server.httpPort() + "/virtualhost");
             assertThat(response.status()).isEqualTo(HttpStatus.OK);
-            // we don't do containsExactly here because there is no easy way to guarantee that
-            // all context hooks from the previous request have been completed
-            assertThat(hookEvents).contains(
+            assertThat(hookEvents).containsExactly(
                     "ClientBuilder/push",
                     "ClientContext/push",
                     "ServerBuilder/push",
diff --git a/core/src/test/java/com/linecorp/armeria/internal/common/CancellationSchedulerTest.java b/core/src/test/java/com/linecorp/armeria/internal/common/CancellationSchedulerTest.java
index 2739d9362b1..2b51fa69e88 100644
--- a/core/src/test/java/com/linecorp/armeria/internal/common/CancellationSchedulerTest.java
+++ b/core/src/test/java/com/linecorp/armeria/internal/common/CancellationSchedulerTest.java
@@ -139,7 +139,7 @@ void cancelTimeoutAfterDeadline() {
     void cancelTimeoutBySettingTimeoutZero() {
         executeInEventLoop(1000, scheduler -> {
             scheduler.setTimeoutNanos(SET_FROM_START, Long.MAX_VALUE);
-            assertThat(scheduler.state()).isEqualTo(State.PENDING);
+            assertThat(scheduler.state()).isEqualTo(State.SCHEDULED);
         });
     }
 
@@ -272,7 +272,7 @@ public void run(Throwable cause) {
             assertThat(scheduler.isFinished()).isFalse();
 
             scheduler.setTimeoutNanos(SET_FROM_NOW, MILLISECONDS.toNanos(1000));
-            assertThat(scheduler.state()).isEqualTo(State.PENDING);
+            assertThat(scheduler.state()).isEqualTo(State.SCHEDULED);
 
             schedulerRef.set(scheduler);
             whenTimedOutRef.set(scheduler.whenCancelled());
@@ -339,7 +339,7 @@ public void run(Throwable cause) {
             assertThat(scheduler.isFinished()).isFalse();
 
             scheduler.setTimeoutNanos(SET_FROM_NOW, MILLISECONDS.toNanos(1000));
-            assertThat(scheduler.state()).isEqualTo(State.PENDING);
+            assertThat(scheduler.state()).isEqualTo(State.SCHEDULED);
 
             schedulerRef.set(scheduler);
             whenCancellingRef.set(scheduler.whenCancelling());