chore(spanner): address review comments

Author: olavloite

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/EwmaLatencyTracker.java

@@ -19,6 +19,8 @@
 import com.google.api.core.BetaApi;
 import com.google.api.core.InternalApi;
 import com.google.common.base.Preconditions;
+import java.time.Duration;
+import java.util.concurrent.TimeUnit;
 import javax.annotation.concurrent.GuardedBy;
 
 /**
@@ -66,20 +68,28 @@ public double getScore() {
   }
 
   @Override
-  public void update(long latencyMillis) {
+  public void update(Duration latency) {
+    long latencyMicros;
+    try {
+      latencyMicros = TimeUnit.MICROSECONDS.convert(latency.toNanos(), TimeUnit.NANOSECONDS);
+    } catch (ArithmeticException e) {
+      // Duration is too large to fit in nanoseconds (292+ years).
+      // Use Long.MAX_VALUE to give it the lowest possible priority.
+      latencyMicros = Long.MAX_VALUE;
+    }
     synchronized (lock) {
       if (!initialized) {
-        score = latencyMillis;
+        score = latencyMicros;
         initialized = true;
       } else {
-        score = alpha * latencyMillis + (1 - alpha) * score;
+        score = alpha * latencyMicros + (1 - alpha) * score;
       }
     }
   }
 
   @Override
-  public void recordError(long penaltyMillis) {
+  public void recordError(Duration penalty) {
     // Treat the error as a sample with high latency (penalty)
-    update(penaltyMillis);
+    update(penalty);
   }
 }

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/LatencyTracker.java

@@ -18,6 +18,7 @@
 
 import com.google.api.core.BetaApi;
 import com.google.api.core.InternalApi;
+import java.time.Duration;
 
 /**
  * Interface for tracking latency scores of Spanner servers.
@@ -39,14 +40,14 @@ public interface LatencyTracker {
   /**
    * Updates the latency score with a new observation.
    *
-   * @param latencyMillis the observed latency in milliseconds.
+   * @param latency the observed latency.
    */
-  void update(long latencyMillis);
+  void update(Duration latency);
 
   /**
    * Records an error and applies a latency penalty.
    *
-   * @param penaltyMillis the penalty in milliseconds to apply.
+   * @param penalty the penalty to apply.
    */
-  void recordError(long penaltyMillis);
+  void recordError(Duration penalty);
 }

java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/EwmaLatencyTrackerTest.java

@@ -19,6 +19,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertThrows;
 
+import java.time.Duration;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -29,7 +30,7 @@ public class EwmaLatencyTrackerTest {
   @Test
   public void testInitialization() {
     EwmaLatencyTracker tracker = new EwmaLatencyTracker();
-    tracker.update(100);
+    tracker.update(Duration.ofNanos(100 * 1000));
     assertEquals(100.0, tracker.getScore(), 0.001);
   }
 
@@ -39,26 +40,33 @@ public void testUninitializedScore() {
     assertEquals(Double.MAX_VALUE, tracker.getScore(), 0.001);
   }
 
+  @Test
+  public void testOverflowScore() {
+    EwmaLatencyTracker tracker = new EwmaLatencyTracker();
+    tracker.update(Duration.ofSeconds(Long.MAX_VALUE));
+    assertEquals((double) Long.MAX_VALUE, tracker.getScore(), 0.001);
+  }
+
   @Test
   public void testEwmaCalculation() {
     double alpha = 0.5;
     EwmaLatencyTracker tracker = new EwmaLatencyTracker(alpha);
 
-    tracker.update(100); // Initial score = 100
+    tracker.update(Duration.ofNanos(100 * 1000)); // Initial score = 100
     assertEquals(100.0, tracker.getScore(), 0.001);
 
-    tracker.update(200); // Score = 0.5 * 200 + 0.5 * 100 = 150
+    tracker.update(Duration.ofNanos(200 * 1000)); // Score = 0.5 * 200 + 0.5 * 100 = 150
     assertEquals(150.0, tracker.getScore(), 0.001);
 
-    tracker.update(300); // Score = 0.5 * 300 + 0.5 * 150 = 225
+    tracker.update(Duration.ofNanos(300 * 1000)); // Score = 0.5 * 300 + 0.5 * 150 = 225
     assertEquals(225.0, tracker.getScore(), 0.001);
   }
 
   @Test
   public void testDefaultAlpha() {
     EwmaLatencyTracker tracker = new EwmaLatencyTracker();
-    tracker.update(100);
-    tracker.update(200);
+    tracker.update(Duration.ofNanos(100 * 1000));
+    tracker.update(Duration.ofNanos(200 * 1000));
 
     double expected =
         EwmaLatencyTracker.DEFAULT_ALPHA * 200 + (1 - EwmaLatencyTracker.DEFAULT_ALPHA) * 100;
@@ -68,9 +76,9 @@ public void testDefaultAlpha() {
   @Test
   public void testRecordError() {
     EwmaLatencyTracker tracker = new EwmaLatencyTracker(0.5);
-    tracker.update(100);
+    tracker.update(Duration.ofNanos(100 * 1000));
 
-    tracker.recordError(10000); // Score = 0.5 * 10000 + 0.5 * 100 = 5050
+    tracker.recordError(Duration.ofNanos(10000 * 1000)); // Score = 0.5 * 10000 + 0.5 * 100 = 5050
     assertEquals(5050.0, tracker.getScore(), 0.001);
   }
 
@@ -84,10 +92,10 @@ public void testInvalidAlpha() {
   @Test
   public void testAlphaOne() {
     EwmaLatencyTracker tracker = new EwmaLatencyTracker(1.0);
-    tracker.update(100);
+    tracker.update(Duration.ofNanos(100 * 1000));
     assertEquals(100.0, tracker.getScore(), 0.001);
 
-    tracker.update(200);
+    tracker.update(Duration.ofNanos(200 * 1000));
     assertEquals(200.0, tracker.getScore(), 0.001);
   }
 }