[Spark] Fire processing-time timers in timestamp order

CHANGES.md

@@ -80,6 +80,7 @@
 
 ## Bugfixes
 
+* (Java) Fixed the Spark runner firing processing-time timers in reverse timestamp order ([#39824](https://github.com/apache/beam/issues/39824)).
 * (Python) Fixed incorrect profiler options handling on portable runners ([#39613](https://github.com/apache/beam/issues/39613)).
 * (Java) KafkaIO dynamic reads no longer require the obsolete `beam_fn_api` experiment ([#29998](https://github.com/apache/beam/issues/29998)).
 

runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkTimerInternals.java

@@ -25,15 +25,15 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import org.apache.beam.runners.core.StateNamespace;
 import org.apache.beam.runners.core.TimerInternals;
 import org.apache.beam.runners.spark.coders.CoderHelpers;
 import org.apache.beam.runners.spark.util.GlobalWatermarkHolder.SparkWatermarks;
 import org.apache.beam.sdk.state.TimeDomain;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Instant;
 
@@ -44,7 +44,8 @@
 public class SparkTimerInternals implements TimerInternals {
   private final Instant highWatermark;
   private final Instant synchronizedProcessingTime;
-  private final Set<TimerData> timers = Sets.newConcurrentHashSet();
+  // Timers keyed by namespace, id and family, so a later setting replaces the prior one.
+  private final Map<List<Object>, TimerData> timers = new ConcurrentHashMap<>();
 
   private Instant inputWatermark;
 
@@ -105,37 +106,46 @@ public static SparkTimerInternals global(Map<Integer, SparkWatermarks> watermark
   }
 
   public Collection<TimerData> getTimers() {
-    return timers;
+    return timers.values();
   }
 
   public void addTimers(Iterator<TimerData> timers) {
     while (timers.hasNext()) {
       TimerData timer = timers.next();
-      this.timers.add(timer);
+      // State written before setTimer replaced prior settings can carry several settings of
+      // one timer; collapse them to the setting with the latest target.
+      this.timers.merge(
+          logicalKey(timer),
+          timer,
+          (existing, restored) ->
+              restored.getTimestamp().isAfter(existing.getTimestamp()) ? restored : existing);
     }
   }
 
   @Override
   public void setTimer(TimerData timer) {
-    this.timers.add(timer);
+    // A later setting of the same timer clears the prior one, per the TimerInternals contract.
+    this.timers.put(logicalKey(timer), timer);
   }
 
   @Override
   public void deleteTimer(
       StateNamespace namespace, String timerId, String timerFamilyId, TimeDomain timeDomain) {
-    this.timers.stream()
-        .filter(
-            timer ->
-                namespace.equals(timer.getNamespace())
-                    && timerId.equals(timer.getTimerId())
-                    && timerFamilyId.equals(timer.getTimerFamilyId())
-                    && timeDomain.equals(timer.getDomain()))
-        .forEach(this::deleteTimer);
+    List<Object> key = ImmutableList.of(namespace, timerId, timerFamilyId);
+    TimerData existing = this.timers.get(key);
+    if (existing != null && timeDomain.equals(existing.getDomain())) {
+      this.timers.remove(key, existing);
+    }
   }
 
   @Override
   public void deleteTimer(TimerData timer) {
-    this.timers.remove(timer);
+    // Deletes this setting only, so a setting made by the fired callback survives.
+    this.timers.remove(logicalKey(timer), timer);
+  }
+
+  private static List<Object> logicalKey(TimerData timer) {
+    return ImmutableList.of(timer.getNamespace(), timer.getTimerId(), timer.getTimerFamilyId());
   }
 
   @Override
@@ -199,31 +209,31 @@ public static Iterator<TimerData> deserializeTimers(
    */
   public boolean hasNextProcessingTimer() {
     final Instant currentProcessingTime = this.currentProcessingTime();
-    return this.timers.stream()
+    return this.timers.values().stream()
         .anyMatch(
             (TimerData timerData) ->
                 timerData.getDomain().equals(TimeDomain.PROCESSING_TIME)
                     && currentProcessingTime.isAfter(timerData.getTimestamp()));
   }
 
   /**
-   * Finds the latest timer in {@link TimeDomain#PROCESSING_TIME} domain that has expired based on
+   * Finds the earliest timer in {@link TimeDomain#PROCESSING_TIME} domain that has expired based on
    * the current processing time.
    *
    * <p>A timer is considered expired when its timestamp is less than the current processing time.
-   * If multiple expired timers exist, the one with the latest timestamp will be returned.
+   * Expired timers fire in timestamp order.
    *
-   * @return The expired processing timer with the latest timestamp if one exists, or {@code null}
+   * @return The expired processing timer with the earliest timestamp if one exists, or {@code null}
    *     if no processing timers are ready to fire.
    */
   public @Nullable TimerData getNextProcessingTimer() {
     final Instant currentProcessingTime = this.currentProcessingTime();
-    return this.timers.stream()
+    return this.timers.values().stream()
         .filter(
             (TimerData timerData) ->
                 timerData.getDomain().equals(TimeDomain.PROCESSING_TIME)
                     && currentProcessingTime.isAfter(timerData.getTimestamp()))
-        .max(Comparator.comparing(TimerData::getTimestamp))
+        .min(Comparator.comparing(TimerData::getTimestamp))
         .orElse(null);
   }
 

runners/spark/src/main/java/org/apache/beam/runners/spark/util/TimerUtils.java

@@ -21,6 +21,7 @@
 import java.io.Serializable;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
@@ -202,10 +203,12 @@ public static <W extends BoundedWindow> void triggerExpiredTimers(
       SparkTimerInternals sparkTimerInternals,
       WindowingStrategy<?, W> windowingStrategy,
       AbstractInOutIterator<?, ?, ?> abstractInOutIterator) {
-    final Collection<TimerInternals.TimerData> expiredTimers =
+    final List<TimerInternals.TimerData> expiredTimers =
         getExpiredTimers(sparkTimerInternals, windowingStrategy);
 
     if (!expiredTimers.isEmpty()) {
+      // Timers fire in timestamp order.
+      expiredTimers.sort(Comparator.comparing(TimerInternals.TimerData::getTimestamp));
       expiredTimers.forEach(abstractInOutIterator::fireTimer);
     }
   }
@@ -221,7 +224,7 @@ public static <W extends BoundedWindow> void dropExpiredTimers(
     }
   }
 
-  private static <W extends BoundedWindow> Collection<TimerInternals.TimerData> getExpiredTimers(
+  private static <W extends BoundedWindow> List<TimerInternals.TimerData> getExpiredTimers(
       SparkTimerInternals sparkTimerInternals, WindowingStrategy<?, W> windowingStrategy) {
     return sparkTimerInternals.getTimers().stream()
         .filter(

runners/spark/src/test/java/org/apache/beam/runners/spark/stateful/SparkTimerInternalsTest.java

@@ -0,0 +1,106 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.spark.stateful;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.beam.runners.core.StateNamespaces;
+import org.apache.beam.runners.core.TimerInternals.TimerData;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.joda.time.Instant;
+import org.junit.Test;
+
+/** Tests for {@link SparkTimerInternals}. */
+public class SparkTimerInternalsTest {
+
+  private static TimerData processingTimer(String timerId, Instant timestamp) {
+    return TimerData.of(
+        timerId, "", StateNamespaces.global(), timestamp, timestamp, TimeDomain.PROCESSING_TIME);
+  }
+
+  @Test
+  public void testProcessingTimersFireInTimestampOrder() {
+    SparkTimerInternals timerInternals = SparkTimerInternals.global(null);
+
+    TimerData first = processingTimer("first", new Instant(1000));
+    TimerData second = processingTimer("second", new Instant(2000));
+    TimerData third = processingTimer("third", new Instant(3000));
+
+    // Set out of order; firing order must follow the timestamps.
+    timerInternals.setTimer(second);
+    timerInternals.setTimer(third);
+    timerInternals.setTimer(first);
+
+    // Drain the way ParDoStateUpdateFn.SparkTimerInternalsIterator does.
+    List<TimerData> fired = new ArrayList<>();
+    TimerData timer;
+    while ((timer = timerInternals.getNextProcessingTimer()) != null) {
+      fired.add(timer);
+      timerInternals.deleteTimer(timer);
+    }
+
+    assertEquals(ImmutableList.of(first, second, third), fired);
+  }
+
+  @Test
+  public void testSettingATimerAgainClearsThePriorSetting() {
+    SparkTimerInternals timerInternals = SparkTimerInternals.global(null);
+
+    timerInternals.setTimer(processingTimer("timer", new Instant(1000)));
+    TimerData latest = processingTimer("timer", new Instant(2000));
+    timerInternals.setTimer(latest);
+
+    assertEquals(ImmutableList.of(latest), ImmutableList.copyOf(timerInternals.getTimers()));
+    assertEquals(latest, timerInternals.getNextProcessingTimer());
+  }
+
+  @Test
+  public void testAddTimersKeepsTheLatestSettingOfATimer() {
+    // State written before setTimer replaced prior settings can carry several settings of one
+    // timer; the setting with the latest target wins regardless of restore order.
+    TimerData earlier = processingTimer("timer", new Instant(1000));
+    TimerData latest = processingTimer("timer", new Instant(2000));
+
+    SparkTimerInternals timerInternals = SparkTimerInternals.global(null);
+    timerInternals.addTimers(ImmutableList.of(earlier, latest).iterator());
+    assertEquals(ImmutableList.of(latest), ImmutableList.copyOf(timerInternals.getTimers()));
+
+    timerInternals = SparkTimerInternals.global(null);
+    timerInternals.addTimers(ImmutableList.of(latest, earlier).iterator());
+    assertEquals(ImmutableList.of(latest), ImmutableList.copyOf(timerInternals.getTimers()));
+  }
+
+  @Test
+  public void testGetNextProcessingTimerIgnoresEventTimeTimers() {
+    SparkTimerInternals timerInternals = SparkTimerInternals.global(null);
+    timerInternals.setTimer(
+        TimerData.of(
+            "event",
+            "",
+            StateNamespaces.global(),
+            new Instant(0),
+            new Instant(0),
+            TimeDomain.EVENT_TIME));
+
+    assertNull(timerInternals.getNextProcessingTimer());
+  }
+}

runners/spark/src/test/java/org/apache/beam/runners/spark/util/TimerUtilsTest.java

@@ -18,6 +18,8 @@
 package org.apache.beam.runners.spark.util;
 
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -36,6 +38,7 @@
 import org.joda.time.Instant;
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.InOrder;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
@@ -111,6 +114,23 @@ public void testTriggerExpiredTimersWithEmptyTimers() {
     verify(mockIterator, never()).fireTimer(any());
   }
 
+  @Test
+  public void testTriggerExpiredTimersFiresInTimestampOrder() {
+    // An even older expired timer, listed after the newer one.
+    TimerInternals.TimerData olderExpiredTimer = mock(TimerInternals.TimerData.class);
+    when(olderExpiredTimer.getTimestamp())
+        .thenReturn(NOW.minus(ALLOWED_LATENESS.plus(Duration.standardMinutes(2))));
+    when(olderExpiredTimer.getDomain()).thenReturn(TimeDomain.EVENT_TIME);
+    when(mockTimerInternals.getTimers()).thenReturn(Arrays.asList(expiredTimer, olderExpiredTimer));
+
+    TimerUtils.triggerExpiredTimers(mockTimerInternals, mockWindowingStrategy, mockIterator);
+
+    // Expired timers fire in timestamp order.
+    InOrder inOrder = inOrder(mockIterator);
+    inOrder.verify(mockIterator).fireTimer(olderExpiredTimer);
+    inOrder.verify(mockIterator).fireTimer(expiredTimer);
+  }
+
   @Test
   public void testTriggerExpiredTimersWithProcessingTimeDomain() {
     // Set up a processing-time timer