restore scheduled timetables to TransitLayer when buffer is cleared

Author: miklcct

application/src/main/java/org/opentripplanner/model/Timetable.java

@@ -20,6 +20,7 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import javax.annotation.Nullable;
 import org.opentripplanner.transit.model.framework.DataValidationException;
@@ -483,4 +484,33 @@ private static TripTimes getRepresentativeTripTimes(
       return null;
     }
   }
+
+  /**
+   * Get a copy of the scheduled timetable valid for the specified service date only
+   */
+  public Timetable copyForServiceDate(LocalDate date) {
+    if (serviceDate != null) {
+      throw new RuntimeException(
+        "Can only copy scheduled timetable for a specific date if a date hasn't been specified yet."
+      );
+    }
+    return copyOf().withServiceDate(date).build();
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (o == null || getClass() != o.getClass()) return false;
+    Timetable timetable = (Timetable) o;
+    return (
+      Objects.equals(pattern, timetable.pattern) &&
+      Objects.equals(tripTimes, timetable.tripTimes) &&
+      Objects.equals(frequencyEntries, timetable.frequencyEntries) &&
+      Objects.equals(serviceDate, timetable.serviceDate)
+    );
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(pattern, tripTimes, frequencyEntries, serviceDate);
+  }
 }

application/src/main/java/org/opentripplanner/model/TimetableSnapshot.java

@@ -552,7 +552,39 @@ public boolean isEmpty() {
    * @return true if the timetable changed as a result of the call
    */
   private boolean clearTimetables(String feedId) {
-    return timetables.keySet().removeIf(tripPattern -> feedId.equals(tripPattern.getFeedId()));
+    var dirty = false;
+    dirtyTimetables.clear();
+
+    for (var entry : timetables.entrySet()) {
+      var pattern = entry.getKey();
+
+      if (feedId.equals(pattern.getFeedId())) {
+        var timetablesForPattern = entry.getValue();
+        var scheduledTimetable = pattern.getScheduledTimetable();
+
+        // remove scheduled timetables from the entry
+        var updatedTimetables = timetablesForPattern
+          .stream()
+          .filter(timetable ->
+            !timetable.equals(scheduledTimetable.copyForServiceDate(timetable.getServiceDate()))
+          );
+
+        // then restore updated timetables to scheduled timetables
+        var restoredTimetables = updatedTimetables
+          .map(timetable -> scheduledTimetable.copyForServiceDate(timetable.getServiceDate()))
+          .collect(ImmutableSortedSet.toImmutableSortedSet(new SortedTimetableComparator()));
+        dirty = dirty || !restoredTimetables.isEmpty();
+        restoredTimetables.forEach(updated ->
+          dirtyTimetables.put(
+            new TripPatternAndServiceDate(pattern, updated.getServiceDate()),
+            updated
+          )
+        );
+        entry.setValue(restoredTimetables);
+      }
+    }
+
+    return dirty;
   }
 
   /**

application/src/test/java/org/opentripplanner/model/TimetableSnapshotTest.java

@@ -20,6 +20,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.SortedSet;
 import java.util.concurrent.atomic.AtomicBoolean;
 import org.junit.jupiter.api.BeforeAll;
@@ -31,7 +32,10 @@
 import org.opentripplanner.transit.model.framework.Deduplicator;
 import org.opentripplanner.transit.model.framework.FeedScopedId;
 import org.opentripplanner.transit.model.framework.Result;
+import org.opentripplanner.transit.model.network.StopPattern;
 import org.opentripplanner.transit.model.network.TripPattern;
+import org.opentripplanner.transit.model.site.RegularStop;
+import org.opentripplanner.transit.model.timetable.RealTimeState;
 import org.opentripplanner.transit.model.timetable.Trip;
 import org.opentripplanner.transit.model.timetable.TripIdAndServiceDate;
 import org.opentripplanner.transit.model.timetable.TripOnServiceDate;
@@ -46,7 +50,7 @@ public class TimetableSnapshotTest {
   private static final ZoneId timeZone = ZoneIds.GMT;
   public static final LocalDate SERVICE_DATE = LocalDate.of(2024, 1, 1);
   private static Map<FeedScopedId, TripPattern> patternIndex;
-  static String feedId;
+  private static String feedId;
 
   @BeforeAll
   public static void setUp() throws Exception {
@@ -412,6 +416,130 @@ void testClear() {
     assertNotNull(snapshot.getRealtimeAddedRoute(pattern.getRoute().getId()));
   }
 
+  /**
+   * This test checks that the original timetable is given to TransitLayerUpdater for previously
+   * added patterns after the buffer is cleared.
+   * <p>
+   * Refer to bug #6197 for details.
+   */
+  @Test
+  void testTransitLayerUpdateAfterClear() {
+    var resolver = new TimetableSnapshot();
+
+    // make an updated trip
+    var pattern = patternIndex.get(new FeedScopedId(feedId, "1.1"));
+    var trip = pattern.scheduledTripsAsStream().findFirst().orElseThrow();
+    var scheduledTimetable = pattern.getScheduledTimetable();
+    var updatedTripTimes = Objects
+      .requireNonNull(scheduledTimetable.getTripTimes(trip))
+      .copyScheduledTimes();
+    for (var i = 0; i < updatedTripTimes.getNumStops(); ++i) {
+      updatedTripTimes.updateArrivalDelay(i, 30);
+      updatedTripTimes.updateDepartureDelay(i, 30);
+    }
+    updatedTripTimes.setRealTimeState(RealTimeState.UPDATED);
+    var realTimeTripUpdate = new RealTimeTripUpdate(
+      pattern,
+      updatedTripTimes,
+      SERVICE_DATE,
+      null,
+      false,
+      false
+    );
+
+    var addedDepartureStopTime = new StopTime();
+    var addedArrivalStopTime = new StopTime();
+    addedDepartureStopTime.setDepartureTime(0);
+    addedDepartureStopTime.setArrivalTime(0);
+    addedDepartureStopTime.setStop(RegularStop.of(new FeedScopedId(feedId, "XX"), () -> 0).build());
+    addedArrivalStopTime.setDepartureTime(300);
+    addedArrivalStopTime.setArrivalTime(300);
+    addedArrivalStopTime.setStop(RegularStop.of(new FeedScopedId(feedId, "YY"), () -> 1).build());
+    var addedStopTimes = List.of(addedDepartureStopTime, addedArrivalStopTime);
+    var addedStopPattern = new StopPattern(addedStopTimes);
+    var route = patternIndex.values().stream().findFirst().orElseThrow().getRoute();
+    var addedTripPattern = TripPattern
+      .of(new FeedScopedId(feedId, "1.1"))
+      .withRoute(route)
+      .withStopPattern(addedStopPattern)
+      .withCreatedByRealtimeUpdater(true)
+      .build();
+    var addedTripTimes = TripTimesFactory.tripTimes(
+      Trip.of(new FeedScopedId(feedId, "addedTrip")).withRoute(route).build(),
+      addedStopTimes,
+      new Deduplicator()
+    );
+    var addedTripUpdate = new RealTimeTripUpdate(
+      addedTripPattern,
+      addedTripTimes,
+      SERVICE_DATE,
+      null,
+      true,
+      false
+    );
+
+    var transitLayerUpdater = new TransitLayerUpdater(null) {
+      int count = 0;
+
+      /**
+       * Test that the TransitLayerUpdater receives correct updated timetables upon commit
+       * <p>
+       * This method is called 3 times.
+       * When count = 0, the buffer contains one added and one updated trip, and the timetables
+       * should reflect this fact.
+       * When count = 1, the buffer is empty, however, this method should still receive the
+       * timetables of the previous added and updated patterns in order to restore them to the
+       * initial scheduled timetable.
+       * When count = 2, the buffer is still empty, and no changes should be made.
+       */
+      @Override
+      public void update(
+        Collection<Timetable> updatedTimetables,
+        Map<TripPattern, SortedSet<Timetable>> timetables
+      ) {
+        assertThat(updatedTimetables).hasSize(count == 2 ? 0 : 2);
+
+        updatedTimetables.forEach(timetable -> {
+          var timetablePattern = timetable.getPattern();
+          assertEquals(SERVICE_DATE, timetable.getServiceDate());
+          if (timetablePattern == pattern) {
+            if (count == 1) {
+              // the timetable for the existing pattern should revert to the original
+              assertEquals(scheduledTimetable.getTripTimes(), timetable.getTripTimes());
+            } else {
+              // the timetable for the existing pattern should contain the modified times
+              assertEquals(
+                RealTimeState.UPDATED,
+                Objects.requireNonNull(timetable.getTripTimes(trip)).getRealTimeState()
+              );
+            }
+          } else if (timetablePattern == addedTripPattern) {
+            if (count == 1) {
+              // the timetable for the added pattern should be empty after clearing
+              assertThat(timetable.getTripTimes()).isEmpty();
+            } else {
+              // the timetable for the added pattern should contain the times for 1 added trip
+              assertThat(timetable.getTripTimes()).hasSize(1);
+            }
+          } else {
+            throw new RuntimeException("unknown pattern passed to transit layer updater");
+          }
+        });
+        ++count;
+      }
+    };
+
+    resolver.update(realTimeTripUpdate);
+    resolver.update(addedTripUpdate);
+    resolver.commit(transitLayerUpdater, true);
+
+    resolver.clear(feedId);
+    resolver.commit(transitLayerUpdater, true);
+
+    resolver.clear(feedId);
+    resolver.commit(transitLayerUpdater, true);
+  }
+
   private static TimetableSnapshot createCommittedSnapshot() {
     TimetableSnapshot timetableSnapshot = new TimetableSnapshot();
     return timetableSnapshot.commit(null, true);