From e5784f843e9634dc57b6735aa45deb6030ce274d Mon Sep 17 00:00:00 2001
From: Bryce Anderson <bryce_anderson@apple.com>
Date: Thu, 21 Dec 2023 11:24:32 -0700
Subject: [PATCH] loadbalancer: fix NormalizedTimeSourceExecutor to work with
 units other than nanos (#2791)

Motivation:

The time source normalizes itself using nanoseconds regardless
of what units the user is requesting.

Modifications:

- Always get the current time in nanos, normalize, then convert
  to the desired units.

Result:

Better behavior.
---
 .../loadbalancer/NormalizedTimeSourceExecutor.java |  7 ++++---
 .../NormalizedTimeSourceExecutorTest.java          | 14 ++++++++++++++
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/servicetalk-loadbalancer/src/main/java/io/servicetalk/loadbalancer/NormalizedTimeSourceExecutor.java b/servicetalk-loadbalancer/src/main/java/io/servicetalk/loadbalancer/NormalizedTimeSourceExecutor.java
index cee6fef96a..aebfb74333 100644
--- a/servicetalk-loadbalancer/src/main/java/io/servicetalk/loadbalancer/NormalizedTimeSourceExecutor.java
+++ b/servicetalk-loadbalancer/src/main/java/io/servicetalk/loadbalancer/NormalizedTimeSourceExecutor.java
@@ -27,15 +27,16 @@
  */
 final class NormalizedTimeSourceExecutor extends DelegatingExecutor {
 
-    private final long offset;
+    private final long offsetNanos;
 
     NormalizedTimeSourceExecutor(final Executor delegate) {
         super(delegate);
-        offset = delegate.currentTime(NANOSECONDS);
+        offsetNanos = delegate.currentTime(NANOSECONDS);
     }
 
     @Override
     public long currentTime(final TimeUnit unit) {
-        return delegate().currentTime(unit) - offset;
+        final long elapsedNanos = delegate().currentTime(NANOSECONDS) - offsetNanos;
+        return unit.convert(elapsedNanos, NANOSECONDS);
     }
 }
diff --git a/servicetalk-loadbalancer/src/test/java/io/servicetalk/loadbalancer/NormalizedTimeSourceExecutorTest.java b/servicetalk-loadbalancer/src/test/java/io/servicetalk/loadbalancer/NormalizedTimeSourceExecutorTest.java
index 1c13d5efae..9bbe34e4f9 100644
--- a/servicetalk-loadbalancer/src/test/java/io/servicetalk/loadbalancer/NormalizedTimeSourceExecutorTest.java
+++ b/servicetalk-loadbalancer/src/test/java/io/servicetalk/loadbalancer/NormalizedTimeSourceExecutorTest.java
@@ -24,7 +24,10 @@
 
 import static java.lang.Long.MAX_VALUE;
 import static java.lang.Long.MIN_VALUE;
+import static java.util.concurrent.TimeUnit.MICROSECONDS;
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static java.util.concurrent.TimeUnit.NANOSECONDS;
+import static java.util.concurrent.TimeUnit.SECONDS;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.is;
 
@@ -60,4 +63,15 @@ void minValue(long initialValue) {
         advanceAndVerify(10, MAX_VALUE + 10);
         advanceAndVerify(MAX_VALUE - 8, 0);
     }
+
+    @ParameterizedTest(name = "{displayName} [{index}]: initialValue={0}")
+    @ValueSource(longs = {MIN_VALUE, -100, 0, 100, MAX_VALUE})
+    void otherUnits(long initialValue) {
+        setUp(initialValue);
+        testExecutor.advanceTimeByNoExecuteTasks(1, SECONDS);
+        assertThat(executor.currentTime(SECONDS), is(1L));
+        assertThat(executor.currentTime(MILLISECONDS), is(1_000L));
+        assertThat(executor.currentTime(MICROSECONDS), is(1_000_000L));
+        assertThat(executor.currentTime(NANOSECONDS), is(1_000_000_000L));
+    }
 }