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.Collection;
  import java.util.Collections;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  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.base.Preconditions;
  import com.google.common.collect.ImmutableList;
  import com.google.common.collect.ImmutableMap;
  import com.google.common.collect.Maps;
  import com.google.common.collect.Sets;
  
  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.DistributedSchema;
  import org.apache.cassandra.schema.TableId;
  import org.apache.cassandra.schema.TableMetadata;
  import org.apache.cassandra.tcm.Epoch;
  import org.apache.cassandra.tcm.MetadataValue;
  import org.apache.cassandra.tcm.serialization.MetadataSerializer;
  import org.apache.cassandra.tcm.serialization.Version;
  import org.apache.cassandra.utils.PojoToString;
  
  import static com.google.common.base.Preconditions.checkNotNull;
  import static com.google.common.base.Preconditions.checkState;
  import static org.apache.cassandra.service.consensus.migration.TableMigrationState.initialRepairPendingRanges;
  import static org.apache.cassandra.utils.CollectionSerializers.deserializeMap;
  import static org.apache.cassandra.utils.CollectionSerializers.serializeMap;
  import static org.apache.cassandra.utils.CollectionSerializers.serializedMapSize;
  
  // TODO this will mostly go away once we can move TableMigrationState into the table schema
  public class ConsensusMigrationState implements MetadataValue<ConsensusMigrationState>
  {
      public static ConsensusMigrationState EMPTY = new ConsensusMigrationState(Epoch.EMPTY, ImmutableMap.of());
  
      @Nonnull
      public final Map<TableId, TableMigrationState> tableStates;
  
      public final Epoch lastModified;
  
      public ConsensusMigrationState(@Nonnull Epoch lastModified, @Nonnull Map<TableId, TableMigrationState> tableStates)
      {
          checkNotNull(tableStates, "tableStates is null");
          checkNotNull(lastModified, "lastModified is null");
          this.lastModified = lastModified;
          this.tableStates = ImmutableMap.copyOf(tableStates);
      }
  
      public Map<String, Object> toMap(@Nullable Set<String> keyspaceNames, @Nullable Set<String> tableNames)
      {
          return ImmutableMap.of("lastModifiedEpoch", lastModified.getEpoch(),
                                 "tableStates", tableStatesAsMaps(keyspaceNames, tableNames),
                                 "version", PojoToString.CURRENT_VERSION);
      }
  
      public Collection<TableMigrationState> tableStates()
      {
          return tableStates.values();
      }
  
      public List<TableMigrationState> tableStatesFor(List<TableId> tableIDs)
      {
          return tableIDs.stream().map(tableStates::get).collect(Collectors.toList());
      }
  
      private List<Map<String, Object>> tableStatesAsMaps(@Nullable Set<String> keyspaceNames,
                                                          @Nullable Set<String> tableNames)
      {
          ImmutableList.Builder<Map<String, Object>> builder = ImmutableList.builder();
          for (TableMigrationState tms : tableStates.values())
          {
              if (keyspaceNames != null && !keyspaceNames.contains(tms.keyspaceName))
                  continue;
              if (tableNames != null && !tableNames.contains(tms.tableName))
                  continue;
              builder.add(tms.toMap());
          }
          return builder.build();
      }
  
      @Override
      public boolean equals(Object o)
      {
          if (this == o) return true;
          if (o == null || getClass() != o.getClass()) return false;
          ConsensusMigrationState that = (ConsensusMigrationState) o;
          return tableStates.equals(that.tableStates);
      }
  
      public ConsensusMigrationState withReversedMigrations(Map<TableId, TableMetadata> tables, Epoch epoch)
      {
          if (tables.isEmpty())
              return this;
  
          ImmutableMap.Builder<TableId, TableMigrationState> updated = ImmutableMap.builder();
  
          tableStates.forEach((id, state) -> {
              if (!tables.containsKey(id))
                  updated.put(id, state);
          });
  
          tables.values().forEach(metadata -> {
              TableMigrationState state = tableStates.get(metadata.id);
              if (state != null)
                  updated.put(metadata.id, state.reverseMigration(ConsensusMigrationTarget.fromTransactionalMode(metadata.params.transactionalMode), epoch));
          });
  
          return new ConsensusMigrationState(lastModified, updated.build());
      }
  
      private static void withRangesMigrating(Map<TableId, TableMigrationState> current, ImmutableMap.Builder<TableId, TableMigrationState> next, TableMetadata metadata, List<Range<Token>> ranges, boolean overwrite)
      {
          TableMigrationState tableState = current.get(metadata.id);
          checkState(tableState != null || overwrite, "Can't begin migrating a table without first altering the schema to set transactional mode");
          TransactionalMigrationFromMode migrationFromMode = metadata.params.transactionalMigrationFrom;
          ConsensusMigrationTarget target = ConsensusMigrationTarget.fromTransactionalMode(metadata.params.transactionalMode);
          checkState(migrationFromMode != null && migrationFromMode != TransactionalMigrationFromMode.none, "Table transactional migration from can't be null or none");
  
          Map<Epoch, List<Range<Token>>> migratingRangesByEpoch = ImmutableMap.of();
          if (!ranges.isEmpty())
              migratingRangesByEpoch = ImmutableMap.of(Epoch.EMPTY, ranges);
  
          if (overwrite)
              tableState = new TableMigrationState(metadata.keyspace, metadata.name, metadata.id, target, ImmutableList.of(), initialRepairPendingRanges(target, ranges), migratingRangesByEpoch);
          else
              tableState = tableState.withRangesMigrating(ranges, target);
  
          next.put(metadata.id, tableState);
      }
  
      private static void putUnchanged(Map<TableId, TableMigrationState> current, ImmutableMap.Builder<TableId, TableMigrationState> next, Set<TableId> changed)
      {
          current.forEach((id, migrationState) -> {
              if (!changed.contains(id))
                  next.put(id, migrationState);
          });
      }
  
      private static void putUnchanged(Map<TableId, TableMigrationState> current, ImmutableMap.Builder<TableId, TableMigrationState> next, Collection<TableMetadata> changed)
      {
          Set<TableId> changedIds = changed.stream().map(TableMetadata::id).collect(Collectors.toSet());
  
          putUnchanged(current, next, changedIds);
      }
  
      public ConsensusMigrationState withRangesMigrating(Collection<TableMetadata> tables, List<Range<Token>> ranges, boolean overwrite)
      {
          ImmutableMap.Builder<TableId, TableMigrationState> updated = ImmutableMap.builder();
          putUnchanged(tableStates, updated, tables);
          tables.forEach(metadata -> withRangesMigrating(tableStates, updated, metadata, ranges, overwrite));
          return new ConsensusMigrationState(lastModified, updated.build());
      }
  
      public ConsensusMigrationState withMigrationsCompletedFor(Collection<TableId> completed)
      {
          ImmutableMap.Builder<TableId, TableMigrationState> updated = ImmutableMap.builder();
          putUnchanged(tableStates, updated, new HashSet<>(completed));
          return new ConsensusMigrationState(lastModified, updated.build());
      }
  
      public ConsensusMigrationState withRangesRepairedAtEpoch(TableMetadata metadata, List<Range<Token>> ranges, Epoch minEpoch, ConsensusMigrationRepairType repairType)
      {
          TableMigrationState state = Preconditions.checkNotNull(tableStates.get(metadata.id));
          state = state.withRangesRepairedAtEpoch(ranges, minEpoch, repairType);
  
          if (state.hasMigratedFullTokenRange(metadata.partitioner))
          {
              return withMigrationsCompletedFor(Collections.singleton(metadata.id));
          }
          else
          {
              ImmutableMap.Builder<TableId, TableMigrationState> updated = ImmutableMap.builder();
              putUnchanged(tableStates, updated, Collections.singleton(metadata.id));
              updated.put(metadata.id, state);
              return new ConsensusMigrationState(lastModified, updated.build());
          }
      }
  
      public ConsensusMigrationState withMigrationsRemovedFor(Set<TableId> removed)
      {
          if (tableStates.isEmpty() || Sets.intersection(tableStates.keySet(), removed).isEmpty())
              return this;
          ImmutableMap.Builder<TableId, TableMigrationState> updated = ImmutableMap.builder();
          putUnchanged(tableStates, updated, removed);
          return new ConsensusMigrationState(lastModified, updated.build());
      }
  
      @Override
      public int hashCode()
      {
          return Objects.hash(tableStates);
      }
  
      @Override
      public ConsensusMigrationState withLastModified(Epoch epoch)
      {
          ImmutableMap.Builder<TableId, TableMigrationState> newMap = ImmutableMap.builderWithExpectedSize(tableStates.size());
          tableStates.forEach((tableId, tableState) -> {
              newMap.put(tableId, tableState.withReplacementForEmptyEpoch(epoch));
          });
          return new ConsensusMigrationState(epoch, newMap.build());
      }
  
      @Override
      public Epoch lastModified()
      {
          return lastModified;
      }
  
      public void validateAgainstSchema(DistributedSchema schema)
      {
          tableStates.forEach((id, migrationState) -> {
              TableMetadata metadata = schema.getTableMetadata(id);
              checkState(ConsensusMigrationTarget.fromTransactionalMode(metadata.params.transactionalMode).equals(migrationState.targetProtocol));
          });
      }
  
      public static final MetadataSerializer<ConsensusMigrationState> serializer = new MetadataSerializer<ConsensusMigrationState>()
      {
          @Override
          public void serialize(ConsensusMigrationState consensusMigrationState, DataOutputPlus out, Version version) throws IOException
          {
              Epoch.serializer.serialize(consensusMigrationState.lastModified, out, version);
              serializeMap(consensusMigrationState.tableStates, out, version, TableId.metadataSerializer, TableMigrationState.serializer);
          }
  
          @Override
          public ConsensusMigrationState deserialize(DataInputPlus in, Version version) throws IOException
          {
              Epoch lastModified = Epoch.serializer.deserialize(in, version);
              Map<TableId, TableMigrationState> tableMigrationStates = deserializeMap(in, version, TableId.metadataSerializer, TableMigrationState.serializer, Maps::newHashMapWithExpectedSize);
              return new ConsensusMigrationState(lastModified, tableMigrationStates);
          }
  
          @Override
          public long serializedSize(ConsensusMigrationState t, Version version)
          {
              return Epoch.serializer.serializedSize(t.lastModified, version)
                     + serializedMapSize(t.tableStates, version, TableId.metadataSerializer, TableMigrationState.serializer);
          }
      };
  
      @Override
      public String toString()
      {
          return "ConsensusMigrationState{" +
                 "tableStates=" + tableStates +
                 ", lastModified=" + lastModified +
                 '}';
      }
  }
  