content: /*
   * 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.cassandra.service.consensus.migration;
  
  import java.io.IOException;
  import java.util.AbstractMap;
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.LinkedHashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.NavigableMap;
  import java.util.Objects;
  import java.util.Set;
  import java.util.stream.Collectors;
  import javax.annotation.Nonnull;
  import javax.annotation.Nullable;
  
  import com.google.common.collect.ImmutableList;
  import com.google.common.collect.ImmutableMap;
  import com.google.common.collect.ImmutableSortedMap;
  import com.google.common.collect.Iterables;
  import com.google.common.collect.Maps;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.cassandra.config.DatabaseDescriptor;
  import org.apache.cassandra.db.DecoratedKey;
  import org.apache.cassandra.db.PartitionPosition;
  import org.apache.cassandra.dht.IPartitioner;
  import org.apache.cassandra.dht.NormalizedRanges;
  import org.apache.cassandra.dht.Range;
  import org.apache.cassandra.dht.Token;
  import org.apache.cassandra.io.util.DataInputPlus;
  import org.apache.cassandra.io.util.DataOutputPlus;
  import org.apache.cassandra.schema.TableId;
  import org.apache.cassandra.tcm.Epoch;
  import org.apache.cassandra.tcm.serialization.MetadataSerializer;
  import org.apache.cassandra.tcm.serialization.Version;
  
  import static com.google.common.base.Preconditions.checkArgument;
  import static com.google.common.base.Preconditions.checkState;
  import static com.google.common.collect.ImmutableList.toImmutableList;
  import static org.apache.cassandra.db.TypeSizes.sizeof;
  import static org.apache.cassandra.dht.NormalizedRanges.normalizedRanges;
  import static org.apache.cassandra.dht.Range.normalize;
  import static org.apache.cassandra.dht.Range.subtract;
  import static org.apache.cassandra.utils.CollectionSerializers.deserializeMap;
  import static org.apache.cassandra.utils.CollectionSerializers.deserializeSet;
  import static org.apache.cassandra.utils.CollectionSerializers.serializeCollection;
  import static org.apache.cassandra.utils.CollectionSerializers.serializeMap;
  import static org.apache.cassandra.utils.CollectionSerializers.serializedCollectionSize;
  import static org.apache.cassandra.utils.CollectionSerializers.serializedMapSize;
  
  // TODO Move this into the schema for the table once this is based off of TrM
  public class TableMigrationState
  {
      private static final Logger logger = LoggerFactory.getLogger(TableMigrationState.class);
  
      @Nonnull
      public final String keyspaceName;
  
      @Nonnull
      public final String tableName;
  
      @Nonnull
      public final TableId tableId;
  
      @Nonnull
      public final ConsensusMigrationTarget targetProtocol;
  
      /**
       * Migrated means that both phases are completed when migrating to Accord. Paxos only has one phase.
       */
      @Nonnull
      public final NormalizedRanges<Token> migratedRanges;
  
      /*
       * Necessary to track which ranges started migrating at which epoch
       * in order to know whether a repair qualifies in terms of finishing
       * migration of the range.
       */
      @Nonnull
      public final NavigableMap<Epoch, NormalizedRanges<Token>> migratingRangesByEpoch;
  
      /**
       * Ranges that are migrating and could be in either phase when migrating to Accord. Paxos only has one phase.
       */
      @Nonnull
      public final NormalizedRanges<Token> migratingRanges;
  
      /**
       * These are ranges that are migrating and have not been repaired yet when migrating to Accord so Accord can't read from them. These ranges
       * should continue to be operated on by Paxos
       *
       * When migrating to Accord a repair can only move the range to migrated if it is already in repairCompletedRanges
       *
       * Additionally Paxos continues to operate on a migrating range and key migration is not performed
       *
       * This is skipped when migrating from Accord to Paxos.
       */
      @Nonnull
      public final NormalizedRanges<Token> repairPendingRanges;
      public final NormalizedRanges<Token> repairedRanges;
  
      /**
       * Ranges that are migrating could be in either phase when migrating to Accord. Paxos only has one phase.
       */
      @Nonnull
      public final NormalizedRanges<Token> migratingAndMigratedRanges;
  
      /**
       * Same as migratingAndMigratedRanges if migrating to Paxos, otherwise migratingAndMigratedRanges.subtract(repairPendingRanges)
       *
       * Not included in equals or hashCode because it is inferred from other fields
       */
      @Nonnull
      public final NormalizedRanges<Token> accordSafeToReadRanges;
  
      public TableMigrationState(@Nonnull String keyspaceName,
                                 @Nonnull String tableName,
                                 @Nonnull TableId tableId,
                                 @Nonnull ConsensusMigrationTarget targetProtocol,
                                 @Nonnull Collection<Range<Token>> migratedRanges,
                                 @Nonnull Collection<Range<Token>> repairPendingRanges,
                                 @Nonnull Map<Epoch, ? extends List<Range<Token>>> migratingRangesByEpoch)
      {
          this.keyspaceName = keyspaceName;
          this.tableName = tableName;
          this.tableId = tableId;
          this.targetProtocol = targetProtocol;
          this.migratedRanges = normalizedRanges(migratedRanges);
          this.repairPendingRanges = normalizedRanges(repairPendingRanges);
          this.migratingRangesByEpoch = ImmutableSortedMap.copyOf(
          migratingRangesByEpoch.entrySet()
                                .stream()
                                .map(entry -> new AbstractMap.SimpleEntry<>(entry.getKey(), normalizedRanges(entry.getValue())))
                                .collect(Collectors.toList()));
          this.migratingRanges = normalizedRanges(migratingRangesByEpoch.values().stream().flatMap(Collection::stream).collect(Collectors.toList()));
          this.repairedRanges = migratingRanges.subtract(this.repairPendingRanges);
          this.migratingAndMigratedRanges = normalizedRanges(ImmutableList.<Range<Token>>builder().addAll(migratedRanges).addAll(migratingRanges).build());
          this.accordSafeToReadRanges = !repairPendingRanges.isEmpty() ? migratingAndMigratedRanges.subtract(this.repairPendingRanges) : migratingAndMigratedRanges;
      }
  
      static List<Range<Token>> initialRepairPendingRanges(ConsensusMigrationTarget target, List<Range<Token>> initialMigratingRanges)
      {
          return target == ConsensusMigrationTarget.accord ? ImmutableList.copyOf(initialMigratingRanges) : ImmutableList.of();
      }
  
      public TableMigrationState reverseMigration(ConsensusMigrationTarget target, Epoch epoch)
      {
          IPartitioner partitioner = DatabaseDescriptor.getPartitioner();
          Range<Token> fullRange = new Range<>(partitioner.getMinimumToken(), partitioner.getMinimumToken());
          List<Range<Token>> allTouched = new ArrayList<>(migratedRanges);
          allTouched.addAll(migratingRanges);
          allTouched = Range.deoverlap(allTouched);
          return new TableMigrationState(keyspaceName, tableName, tableId, target,
                                         Range.normalize(fullRange.subtractAll(allTouched)),
                                         initialRepairPendingRanges(target, migratingRanges),
                                         Collections.singletonMap(epoch, migratingRanges));
      }
  
      public boolean hasMigratedFullTokenRange(IPartitioner partitioner)
      {
          // migrated ranges are normalized
          if (!migratingRanges.isEmpty() || migratedRanges.size() != 1)
              return false;
  
          Range<Token> fullRange = new Range<>(partitioner.getMinimumToken(), partitioner.getMinimumToken());
          return migratedRanges.get(0).contains(fullRange);
      }
  
      @Nonnull
      public List<Range<Token>> migratingRanges() {
  
          return migratingRanges;
      }
  
      @Nonnull
      public List<Range<Token>> repairPendingRanges()
      {
          return repairPendingRanges;
      }
  
      public TableMigrationState withRangesMigrating(@Nonnull Collection<Range<Token>> ranges,
                                                     @Nonnull ConsensusMigrationTarget target)
      {
          checkState(!migratingRangesByEpoch.containsKey(Epoch.EMPTY), "Shouldn't already have an entry for the empty epoch");
          // Doesn't matter which epoch the range started migrating in for this context so merge them all
          Collection<Range<Token>> migratingRanges = normalize(migratingRangesByEpoch.values().stream().flatMap(Collection::stream).collect(Collectors.toList()));
          checkArgument(target == targetProtocol, "Requested migration to target protocol " + target + " conflicts with in progress migration to protocol " + targetProtocol);
          List<Range<Token>> normalizedRanges = normalize(ranges);
          if (subtract(normalizedRanges, migratingRanges).isEmpty())
              logger.warn("Range " + ranges + " is already being migrated");
          Set<Range<Token>> withoutAlreadyMigrated = subtract(normalizedRanges, migratedRanges);
          if (withoutAlreadyMigrated.isEmpty())
              logger.warn("Range " + ranges + " is already migrated");
          Set<Range<Token>> withoutBoth = subtract(withoutAlreadyMigrated, migratingRanges);
          if (withoutBoth.isEmpty())
              logger.warn("Range " + ranges + " is already migrating/migrated");
  
          if (!Range.equals(normalizedRanges, withoutBoth))
              logger.warn("Ranges " + normalizedRanges + " to start migrating is already partially migrating/migrated " + withoutBoth);
  
          Map<Epoch, List<Range<Token>>> newMigratingRanges = new HashMap<>(migratingRangesByEpoch.size() + 1);
          newMigratingRanges.putAll(migratingRangesByEpoch);
          newMigratingRanges.put(Epoch.EMPTY, normalizedRanges);
  
          List<Range<Token>> newRepairPendingRanges = new ArrayList<>(repairPendingRanges);
          if (target == ConsensusMigrationTarget.accord)
              newRepairPendingRanges.addAll(withoutBoth);
  
          return new TableMigrationState(keyspaceName, tableName, tableId, targetProtocol, migratedRanges, newRepairPendingRanges, newMigratingRanges);
      }
  
      public TableMigrationState withReplacementForEmptyEpoch(@Nonnull Epoch replacementEpoch)
      {
          if (!migratingRangesByEpoch.containsKey(Epoch.EMPTY))
              return this;
          Map<Epoch, List<Range<Token>>> newMigratingRangesByEpoch = new HashMap<>(migratingRangesByEpoch.size());
          migratingRangesByEpoch.forEach((epoch, ranges) -> {
              if (epoch.equals(Epoch.EMPTY))
                  newMigratingRangesByEpoch.put(replacementEpoch, ranges);
              else
                  newMigratingRangesByEpoch.put(epoch, ranges);
          });
  
          if (newMigratingRangesByEpoch != null)
              return new TableMigrationState(keyspaceName, tableName, tableId, targetProtocol, migratedRanges, repairPendingRanges, newMigratingRangesByEpoch);
          else
              return this;
      }
  
      public NormalizedRanges<Token> normalizedIntersectionWithMigratingAtEpoch(Epoch epoch, Collection<Range<Token>> ranges)
      {
          // This should be inclusive because the epoch we store in the map is the epoch in which the range has been marked migrating
          // in startMigrationToConsensusProtocol
          NormalizedRanges<Token> normalizedRepairedRanges = normalizedRanges(ranges);
          NavigableMap<Epoch, NormalizedRanges<Token>> coveredEpochs = migratingRangesByEpoch.headMap(epoch, true);
          NormalizedRanges<Token> normalizedMigratingRanges = normalizedRanges(coveredEpochs.values().stream().flatMap(Collection::stream).collect(Collectors.toList()));
          // These are the ranges that are impacted by this repair based on the epoch filtering needed to make sure this repair applies to this migration
          return normalizedRepairedRanges.intersection(normalizedMigratingRanges);
      }
  
      public TableMigrationState withRangesRepairedAtEpoch(@Nonnull Collection<Range<Token>> ranges,
                                                           @Nonnull Epoch epoch,
                                                           @Nonnull ConsensusMigrationRepairType repairType)
      {
          checkState(!migratingRangesByEpoch.containsKey(Epoch.EMPTY), "Shouldn't have an entry for the empty epoch");
          checkArgument(epoch.isAfter(Epoch.EMPTY), "Epoch shouldn't be empty");
  
          NormalizedRanges<Token> normalizedRepairedRanges = normalizedRanges(ranges);
          // This should be inclusive because the epoch we store in the map is the epoch in which the range has been marked migrating
          // in startMigrationToConsensusProtocol
          NavigableMap<Epoch, NormalizedRanges<Token>> coveredEpochs = migratingRangesByEpoch.headMap(epoch, true);
          NormalizedRanges<Token> normalizedMigratingRanges = normalizedRanges(coveredEpochs.values().stream().flatMap(Collection::stream).collect(Collectors.toList()));
          // These are the ranges that are impacted by this repair based on the epoch filtering needed to make sure this repair applies to this migration
          NormalizedRanges<Token> normalizedRepairedIntersection = normalizedRepairedRanges.intersection(normalizedMigratingRanges);
          checkState(!normalizedRepairedIntersection.isEmpty(), "None of Ranges " + ranges + " were being migrated");
  
          // Any ranges that were repair pending can't actually be fully migrated yet, they will be subtracted from repairPendingRanges after
          // the new migratingRangesByEpoch and migratedRanges are constructed
          NormalizedRanges<Token> actuallyMigratedRanges = normalizedRepairedIntersection.subtract(repairPendingRanges);
  
          List<Range<Token>> newMigratedRanges = migratedRanges;
          Map<Epoch, NormalizedRanges<Token>> newMigratingRangesByEpoch = migratingRangesByEpoch;
  
          // Not all repairs are capable of completing the migration to a given target
          if ((targetProtocol == ConsensusMigrationTarget.accord && repairType.repairsPaxos())
              || (targetProtocol == ConsensusMigrationTarget.paxos && repairType.repairedAccord))
          {
              newMigratingRangesByEpoch = new HashMap<>();
              // Everything in this epoch or later can't have been migrated so re-add all of them
              newMigratingRangesByEpoch.putAll(migratingRangesByEpoch.tailMap(epoch, false));
              // Include anything still remaining to be migrated after subtracting what was repaired (and not excluded to due repairPendingRanges)
              for (Map.Entry<Epoch, NormalizedRanges<Token>> e : coveredEpochs.entrySet())
              {
                  // Epoch when these ranges started migrating
                  Epoch rangesEpoch = e.getKey();
                  NormalizedRanges<Token> epochMigratingRanges = e.getValue();
                  NormalizedRanges<Token> remainingRanges = epochMigratingRanges.subtract(actuallyMigratedRanges);
                  if (!remainingRanges.isEmpty())
                      newMigratingRangesByEpoch.put(rangesEpoch, remainingRanges);
              }
  
              newMigratedRanges = new ArrayList<>(normalizedMigratingRanges.size() + ranges.size());
              newMigratedRanges.addAll(migratedRanges);
              newMigratedRanges.addAll(actuallyMigratedRanges);
          }
  
          // After this repair any ranges in normalizedRepairedIntersection is repaired and no longer repair pending
          // Accord can safely read from them if this is a migration to Accord (after Paxos key migration)
          List<Range<Token>> newRepairPendingRanges = repairPendingRanges;
          if (repairType.repairedData)
          {
              List<Range<Token>> repairedRangesSatisfyingEpoch = new ArrayList<>();
              for (Map.Entry<Epoch, NormalizedRanges<Token>> e : coveredEpochs.entrySet())
                  repairedRangesSatisfyingEpoch.addAll(normalizedRepairedRanges.intersection(e.getValue()));
              newRepairPendingRanges = repairPendingRanges.subtract(normalizedRanges(repairedRangesSatisfyingEpoch));
          }
  
          return new TableMigrationState(keyspaceName, tableName, tableId, targetProtocol, newMigratedRanges, newRepairPendingRanges, newMigratingRangesByEpoch);
      }
  
      public boolean paxosReadSatisfiedByKeyMigrationAtEpoch(DecoratedKey key, ConsensusMigratedAt consensusMigratedAt)
      {
          // This check is being done from a Paxos read attempt which needs to
          // check if Accord needs to resolve any in flight accord transactions
          // if the migration target is Accord then nothing needs to be done
          if (targetProtocol != ConsensusMigrationTarget.paxos)
              return true;
  
          return satisfiedByKeyMigrationAtEpoch(key, consensusMigratedAt);
      }
  
      public boolean satisfiedByKeyMigrationAtEpoch(@Nonnull DecoratedKey key, @Nullable ConsensusMigratedAt consensusMigratedAt)
      {
          if (consensusMigratedAt == null)
          {
              // It hasn't been migrated and needs migration if it is in a migrating range
              return migratingRanges.intersects(key.getToken());
          }
          else
          {
              // It has been migrated and might be from a late enough epoch to satisfy this migration
              return consensusMigratedAt.migratedAtTarget == targetProtocol
                     && migratingRangesByEpoch.headMap(consensusMigratedAt.migratedAtEpoch, true).values()
                                              .stream()
                                              .flatMap(List::stream)
                                              .anyMatch(range -> range.contains(key.getToken()));
          }
      }
  
      public Epoch minMigrationEpoch(Token token)
      {
          // TODO should there be an index to make this more efficient?
          for (Map.Entry<Epoch, NormalizedRanges<Token>> e : migratingRangesByEpoch.entrySet())
          {
              if (e.getValue().intersects(token))
                  return e.getKey();
          }
          return Epoch.EMPTY;
      }
  
  
      public @Nonnull TableId getTableId()
      {
          return tableId;
      }
  
      public Map<String, Object> toMap()
      {
          ImmutableMap.Builder<String, Object> builder = ImmutableMap.builder();
          builder.put("keyspace", keyspaceName);
          builder.put("table", tableName);
          builder.put("tableId", tableId.toString());
          builder.put("targetProtocol", targetProtocol.toString());
          builder.put("migratedRanges", migratedRanges.stream().map(Objects::toString).collect(toImmutableList()));
          Map<Long, List<String>> rangesByEpoch = new LinkedHashMap<>();
          for (Map.Entry<Epoch, NormalizedRanges<Token>> entry : migratingRangesByEpoch.entrySet())
          {
              rangesByEpoch.put(entry.getKey().getEpoch(), entry.getValue().stream().map(Objects::toString).collect(toImmutableList()));
          }
          builder.put("migratingRangesByEpoch", rangesByEpoch);
          builder.put("repairPendingRanges", repairPendingRanges.stream().map(Objects::toString).collect(toImmutableList()));
          return builder.build();
      }
  
      @Override
      public boolean equals(Object o)
      {
          if (this == o) return true;
          if (o == null || getClass() != o.getClass()) return false;
          TableMigrationState that = (TableMigrationState) o;
          return keyspaceName.equals(that.keyspaceName) && tableName.equals(that.tableName) && tableId.equals(that.tableId) && targetProtocol == that.targetProtocol && migratedRanges.equals(that.migratedRanges) && migratingRangesByEpoch.equals(that.migratingRangesByEpoch) && migratingRanges.equals(that.migratingRanges) && migratingAndMigratedRanges.equals(that.migratingAndMigratedRanges);
      }
  
      @Override
      public int hashCode()
      {
          return Objects.hash(keyspaceName, tableName, tableId, targetProtocol, migratedRanges, migratingRangesByEpoch, migratingRanges, migratingAndMigratedRanges);
      }
  
      public static final MetadataSerializer<TableMigrationState> serializer = new MetadataSerializer<TableMigrationState>()
      {
          @Override
          public void serialize(TableMigrationState t, DataOutputPlus out, Version version) throws IOException
          {
              out.write(t.targetProtocol.value);
              out.writeUTF(t.keyspaceName);
              out.writeUTF(t.tableName);
              t.tableId.serialize(out);
              serializeCollection(t.migratedRanges, out, version, Range.serializer);
              serializeCollection(t.repairPendingRanges, out, version, Range.serializer);
              serializeMap(t.migratingRangesByEpoch, out, version, Epoch.serializer, ConsensusTableMigration.rangesSerializer);
          }
  
          @Override
          public TableMigrationState deserialize(DataInputPlus in, Version version) throws IOException
          {
              ConsensusMigrationTarget targetProtocol = ConsensusMigrationTarget.fromValue(in.readByte());
              String keyspaceName = in.readUTF();
              String tableName = in.readUTF();
              TableId tableId = TableId.deserialize(in);
              Set<Range<Token>> migratedRanges = deserializeSet(in, version, Range.serializer);
              Set<Range<Token>> repairPendingRanges = deserializeSet(in, version, Range.serializer);
              Map<Epoch, NormalizedRanges<Token>> migratingRangesByEpoch = deserializeMap(in, version, Epoch.serializer, ConsensusTableMigration.rangesSerializer, Maps::newHashMapWithExpectedSize);
              return new TableMigrationState(keyspaceName, tableName, tableId, targetProtocol, migratedRanges, repairPendingRanges, migratingRangesByEpoch);
          }
  
          @Override
          public long serializedSize(TableMigrationState t, Version version)
          {
              return sizeof(t.targetProtocol.value)
                     + sizeof(t.keyspaceName)
                     + sizeof(t.tableName)
                     + t.tableId.serializedSize()
                     + serializedCollectionSize(t.migratedRanges, version, Range.serializer)
                     + serializedCollectionSize(t.repairPendingRanges, version, Range.serializer)
                     + serializedMapSize(t.migratingRangesByEpoch, version, Epoch.serializer, ConsensusTableMigration.rangesSerializer);
          }
      };
  
      public Iterable<Range<PartitionPosition>> migratedRangesAsPartitionPosition()
      {
          return Iterables.transform(migratedRanges, range -> new Range<>(range.left.maxKeyBound(), range.right.maxKeyBound()));
      }
  
      public Iterable<Range<PartitionPosition>> migratingAndMigratedRangesAsPartitionPosition()
      {
          return Iterables.transform(migratingAndMigratedRanges, range -> new Range<>(range.left.maxKeyBound(), range.right.maxKeyBound()));
      }
  
      @Override
      public String toString()
      {
          return "TableMigrationState{" +
                 "keyspaceName='" + keyspaceName + '\'' +
                 ", tableName='" + tableName + '\'' +
                 ", tableId=" + tableId +
                 ", targetProtocol=" + targetProtocol +
                 ", migratedRanges=" + migratedRanges +
                 ", migratingRangesByEpoch=" + migratingRangesByEpoch +
                 ", migratingRanges=" + migratingRanges +
                 ", repairPendingRanges=" + repairPendingRanges +
                 ", migratingAndMigratedRanges=" + migratingAndMigratedRanges +
                 '}';
      }
  }
  