[ML] Protect against arithmetic overflow when using TimeValue.MAX_VALUE (#116749)

Author: davidkyle

muted-tests.yml

@@ -236,8 +236,6 @@ tests:
 - class: org.elasticsearch.snapshots.SnapshotShutdownIT
   method: testRestartNodeDuringSnapshot
   issue: https://github.com/elastic/elasticsearch/issues/116730
-- class: org.elasticsearch.xpack.inference.InferenceRestIT
-  issue: https://github.com/elastic/elasticsearch/issues/116740
 - class: org.elasticsearch.action.search.SearchRequestTests
   method: testSerializationConstants
   issue: https://github.com/elastic/elasticsearch/issues/116752

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/common/time/RemainingTime.java

@@ -18,8 +18,13 @@ public interface RemainingTime extends Supplier<TimeValue> {
      * Create a {@link Supplier} that returns a decreasing {@link TimeValue} on each invocation, representing the amount of time until
      * the call times out.  The timer starts when this method is called and counts down from remainingTime to 0.
      * currentTime should return the most up-to-date system time, for example Instant.now() or Clock.instant().
+     * {@link TimeValue#MAX_VALUE} is a special case where the remaining time is always TimeValue.MAX_VALUE.
      */
     static RemainingTime from(Supplier<Instant> currentTime, TimeValue remainingTime) {
+        if (remainingTime.equals(TimeValue.MAX_VALUE)) {
+            return () -> TimeValue.MAX_VALUE;
+        }
+
         var timeout = currentTime.get().plus(remainingTime.duration(), remainingTime.timeUnit().toChronoUnit());
         var maxRemainingTime = remainingTime.nanos();
         return () -> {

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/common/time/RemainingTimeTests.java

@@ -32,6 +32,11 @@ public void testRemainingTimeMaxValue() {
         assertThat(remainingTime.get(), Matchers.equalTo(TimeValue.ZERO));
     }
 
+    public void testMaxTime() {
+        var remainingTime = RemainingTime.from(Instant::now, TimeValue.MAX_VALUE);
+        assertThat(remainingTime.get(), Matchers.equalTo(TimeValue.MAX_VALUE));
+    }
+
     // always add the first value, which is read when RemainingTime.from is called, then add the test values
     private Supplier<Instant> times(Instant... instants) {
         var startTime = Stream.of(Instant.now());