Fix copy-on-write in TimetableSnapshot

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

@@ -1,6 +1,7 @@
 package org.opentripplanner.model;
 
 import com.google.common.collect.HashMultimap;
+import com.google.common.collect.ImmutableSortedSet;
 import com.google.common.collect.SetMultimap;
 import java.time.LocalDate;
 import java.util.Collection;
@@ -186,25 +187,7 @@ public Result<UpdateSuccess, UpdateError> update(
     Timetable tt = resolve(pattern, serviceDate);
     // we need to perform the copy of Timetable here rather than in Timetable.update()
     // to avoid repeatedly copying in case several updates are applied to the same timetable
-    if (!dirtyTimetables.contains(tt)) {
-      Timetable old = tt;
-      tt = new Timetable(tt, serviceDate);
-      SortedSet<Timetable> sortedTimetables = timetables.get(pattern);
-      if (sortedTimetables == null) {
-        sortedTimetables = new TreeSet<>(new SortedTimetableComparator());
-      } else {
-        SortedSet<Timetable> temp = new TreeSet<>(new SortedTimetableComparator());
-        temp.addAll(sortedTimetables);
-        sortedTimetables = temp;
-      }
-      if (old.getServiceDate() != null) {
-        sortedTimetables.remove(old);
-      }
-      sortedTimetables.add(tt);
-      timetables.put(pattern, sortedTimetables);
-      dirtyTimetables.add(tt);
-      dirty = true;
-    }
+    tt = copyTimetable(pattern, serviceDate, tt);
 
     // Assume all trips in a pattern are from the same feed, which should be the case.
     // Find trip index
@@ -330,10 +313,11 @@ public boolean revertTripToScheduledTripPattern(FeedScopedId tripId, LocalDate s
         }
 
         if (tripTimesToRemove != null) {
-          for (Timetable sortedTimetable : sortedTimetables) {
-            boolean isDirty = sortedTimetable.getTripTimes().remove(tripTimesToRemove);
+          for (Timetable originalTimetable : sortedTimetables) {
+            boolean isDirty = originalTimetable.getTripTimes().contains(tripTimesToRemove);
             if (isDirty) {
-              dirtyTimetables.add(sortedTimetable);
+              Timetable updatedTimetable = copyTimetable(pattern, serviceDate, originalTimetable);
+              updatedTimetable.getTripTimes().remove(tripTimesToRemove);
             }
           }
         }
@@ -455,6 +439,39 @@ private void addPatternToIndex(TripPattern tripPattern) {
     }
   }
 
+  /**
+   * Make a copy of the given timetable for a given pattern and service date.
+   * If the timetable was already copied-on write in this snapshot, the same instance will be
+   * returned. The SortedSet that holds the collection of Timetables for that pattern
+   * (sorted by service date) is shared between multiple snapshots and must be copied as well.<br/>
+   * Note on performance: if  multiple Timetables are modified in a SortedSet, the SortedSet will be
+   * copied multiple times. The impact on memory/garbage collection is assumed to be minimal
+   * since the collection is small.
+   * The SortedSet is made immutable to prevent change after snapshot publication.
+   */
+  private Timetable copyTimetable(TripPattern pattern, LocalDate serviceDate, Timetable tt) {
+    if (!dirtyTimetables.contains(tt)) {
+      Timetable old = tt;
+      tt = new Timetable(tt, serviceDate);
+      SortedSet<Timetable> sortedTimetables = timetables.get(pattern);
+      if (sortedTimetables == null) {
+        sortedTimetables = new TreeSet<>(new SortedTimetableComparator());
+      } else {
+        SortedSet<Timetable> temp = new TreeSet<>(new SortedTimetableComparator());
+        temp.addAll(sortedTimetables);
+        sortedTimetables = temp;
+      }
+      if (old.getServiceDate() != null) {
+        sortedTimetables.remove(old);
+      }
+      sortedTimetables.add(tt);
+      timetables.put(pattern, ImmutableSortedSet.copyOfSorted(sortedTimetables));
+      dirtyTimetables.add(tt);
+      dirty = true;
+    }
+    return tt;
+  }
+
   protected static class SortedTimetableComparator implements Comparator<Timetable> {
 
     @Override