Properly set lastModifiedEpoch on multistep operations

Patch by Marcus Eriksson; reviewed by Sam Tunnicliffe for
CASSANDRA-19538

src/java/org/apache/cassandra/tcm/MultiStepOperation.java

@@ -155,7 +155,7 @@ public abstract class MultiStepOperation<CONTEXT>
     public static Transformation.Result applyMultipleTransformations(ClusterMetadata metadata, Transformation.Kind next, List<Transformation> transformations)
     {
         ImmutableSet.Builder<MetadataKey> modifiedKeys = ImmutableSet.builder();
-        Epoch lastModifiedEpoch = metadata.epoch;
+        Epoch lastModifiedEpoch = metadata.epoch.nextEpoch();
         boolean foundStart = false;
         for (Transformation nextTransformation : transformations)
         {
@@ -169,7 +169,7 @@ public abstract class MultiStepOperation<CONTEXT>
                 modifiedKeys.addAll(result.success().affectedMetadata);
             }
         }
-        return new Transformation.Success(metadata.forceEpoch(lastModifiedEpoch.nextEpoch()), LockedRanges.AffectedRanges.EMPTY, modifiedKeys.build());
+        return new Transformation.Success(metadata, LockedRanges.AffectedRanges.EMPTY, modifiedKeys.build());
     }
 
     /**

src/java/org/apache/cassandra/tcm/ownership/DataPlacement.java

@@ -155,6 +155,12 @@ public class DataPlacement
         }
     }
 
+    public DataPlacement withCappedLastModified(Epoch lastModified)
+    {
+        return new DataPlacement(reads.withCappedLastModified(lastModified),
+                                 writes.withCappedLastModified(lastModified));
+    }
+
     @Override
     public String toString()
     {

src/java/org/apache/cassandra/tcm/ownership/DataPlacements.java

@@ -25,6 +25,7 @@ import java.util.Map;
 import java.util.Objects;
 import java.util.function.BiConsumer;
 
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
 
@@ -122,7 +123,7 @@ public class DataPlacements extends ReplicationMap<DataPlacement> implements Met
     @Override
     public DataPlacements withLastModified(Epoch epoch)
     {
-        return new DataPlacements(epoch, asMap());
+        return new DataPlacements(epoch, capLastModified(epoch, map));
     }
 
     @Override
@@ -255,6 +256,13 @@ public class DataPlacements extends ReplicationMap<DataPlacement> implements Met
         }
     }
 
+    public static ImmutableMap<ReplicationParams, DataPlacement> capLastModified(Epoch lastModified, Map<ReplicationParams, DataPlacement> placements)
+    {
+        ImmutableMap.Builder<ReplicationParams, DataPlacement> builder = ImmutableMap.builder();
+        placements.forEach((params, placement) -> builder.put(params, placement.withCappedLastModified(lastModified)));
+        return builder.build();
+    }
+
     public void dumpDiff(DataPlacements other)
     {
         if (!map.equals(other.map))

src/java/org/apache/cassandra/tcm/ownership/PlacementForRange.java

@@ -152,6 +152,20 @@ public class PlacementForRange
         return builder.build();
     }
 
+    public PlacementForRange withCappedLastModified(Epoch lastModified)
+    {
+        SortedMap<Range<Token>, VersionedEndpoints.ForRange> copy = new TreeMap<>();
+        for (Map.Entry<Range<Token>, VersionedEndpoints.ForRange> entry : replicaGroups.entrySet())
+        {
+            Range<Token> range = entry.getKey();
+            VersionedEndpoints.ForRange forRange = entry.getValue();
+            if (forRange.lastModified().isAfter(lastModified))
+                forRange = forRange.withLastModified(lastModified);
+            copy.put(range, forRange);
+        }
+        return new PlacementForRange(copy);
+    }
+
     @Override
     public String toString()
     {

src/java/org/apache/cassandra/tcm/ownership/VersionedEndpoints.java

@@ -63,7 +63,7 @@ public interface VersionedEndpoints<E extends Endpoints<E>> extends MetadataValu
 
         public ForRange withLastModified(Epoch epoch)
         {
-            return new ForRange(lastModified, endpointsForRange);
+            return new ForRange(epoch, endpointsForRange);
         }
 
         public Epoch lastModified()
@@ -151,7 +151,7 @@ public interface VersionedEndpoints<E extends Endpoints<E>> extends MetadataValu
 
         public ForToken withLastModified(Epoch epoch)
         {
-            return new ForToken(lastModified, endpointsForToken);
+            return new ForToken(epoch, endpointsForToken);
         }
 
         public ForToken map(Function<EndpointsForToken, EndpointsForToken> fn)
@@ -159,13 +159,11 @@ public interface VersionedEndpoints<E extends Endpoints<E>> extends MetadataValu
             return new ForToken(lastModified, fn.apply(endpointsForToken));
         }
 
-
         public ForToken without(Set<InetAddressAndPort> remove)
         {
             return map(e -> e.without(remove));
         }
 
-
         public Epoch lastModified()
         {
             return lastModified;