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.cql3.statements;
  
  import java.nio.ByteBuffer;
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  
  import com.google.common.collect.HashMultiset;
  import com.google.common.collect.ImmutableMap;
  import com.google.common.collect.Maps;
  
  import org.apache.cassandra.db.ConsistencyLevel;
  import org.apache.cassandra.db.CounterMutation;
  import org.apache.cassandra.db.DecoratedKey;
  import org.apache.cassandra.db.IMutation;
  import org.apache.cassandra.db.Mutation;
  import org.apache.cassandra.db.ReadCommand.PotentialTxnConflicts;
  import org.apache.cassandra.db.RegularAndStaticColumns;
  import org.apache.cassandra.db.commitlog.CommitLogSegment;
  import org.apache.cassandra.db.partitions.PartitionUpdate;
  import org.apache.cassandra.db.virtual.VirtualMutation;
  import org.apache.cassandra.net.MessagingService;
  import org.apache.cassandra.schema.TableId;
  import org.apache.cassandra.schema.TableMetadata;
  import org.apache.cassandra.service.ClientState;
  
  import static org.apache.cassandra.utils.MonotonicClock.Global.approxTime;
  
  /**
   * Utility class to collect updates.
   *
   * <p>In a batch statement we don't want to recreate mutations every time as this is particularly inefficient when
   * applying multiple batch to the same partition (see #6737). </p>
   *
   */
  final class BatchUpdatesCollector implements UpdatesCollector
  {
      /**
       * The columns that will be updated for each table (keyed by the table ID).
       */
      private final Map<TableId, RegularAndStaticColumns> updatedColumns;
  
      /**
       * The number of updated rows per table and key.
       */
      private final Map<TableId, HashMultiset<ByteBuffer>> perPartitionKeyCounts;
  
      /**
       * The mutations per keyspace.
       *
       * optimised for the common single-keyspace case
       *
       * Key is keyspace name, then we have an IMutationBuilder for each touched partition key in that keyspace
       *
       * MutationBuilder holds a PartitionUpdate.Builder
       */
      private final Map<String, Map<ByteBuffer, IMutationBuilder>> mutationBuilders = Maps.newHashMapWithExpectedSize(1);
  
  
      BatchUpdatesCollector(Map<TableId, RegularAndStaticColumns> updatedColumns, Map<TableId, HashMultiset<ByteBuffer>> perPartitionKeyCounts)
      {
          super();
          this.updatedColumns = updatedColumns;
          this.perPartitionKeyCounts = perPartitionKeyCounts;
      }
  
      /**
       * Gets the <code>PartitionUpdate.Builder</code> for the specified column family and key. If the builder does not
       * exist it will be created.
       *
       * @param metadata the column family meta data
       * @param dk the partition key
       * @param consistency the consistency level
       * @return the <code>PartitionUpdate.Builder</code> for the specified column family and key
       */
      public PartitionUpdate.Builder getPartitionUpdateBuilder(TableMetadata metadata, DecoratedKey dk, ConsistencyLevel consistency)
      {
          IMutationBuilder mut = getMutationBuilder(metadata, dk, consistency);
          PartitionUpdate.Builder upd = mut.get(metadata.id);
          if (upd == null)
          {
              RegularAndStaticColumns columns = updatedColumns.get(metadata.id);
              assert columns != null;
              upd = new PartitionUpdate.Builder(metadata, dk, columns, perPartitionKeyCounts.get(metadata.id).count(dk.getKey()));
              mut.add(upd);
          }
          return upd;
      }
  
      private IMutationBuilder getMutationBuilder(TableMetadata metadata, DecoratedKey dk, ConsistencyLevel consistency)
      {
          Map<ByteBuffer, IMutationBuilder> ksMap = keyspaceMap(metadata.keyspace);
          IMutationBuilder mutationBuilder = ksMap.get(dk.getKey());
          if (mutationBuilder == null)
          {
              mutationBuilder = makeMutationBuilder(metadata, dk, consistency);
              ksMap.put(dk.getKey(), mutationBuilder);
          }
          return mutationBuilder;
      }
  
      private IMutationBuilder makeMutationBuilder(TableMetadata metadata, DecoratedKey partitionKey, ConsistencyLevel cl)
      {
          if (metadata.isVirtual())
          {
              return new VirtualMutationBuilder(metadata.keyspace, partitionKey);
          }
          else
          {
              MutationBuilder builder = new MutationBuilder(metadata.keyspace, partitionKey, 1);
              return metadata.isCounter() ? new CounterMutationBuilder(builder, cl) : builder;
          }
      }
  
      /**
       * Returns a collection containing all the mutations.
       * 
       * @param state state related to the client connection
       * 
       * @return a collection containing all the mutations.
       */
      @Override
      public List<IMutation> toMutations(ClientState state, PotentialTxnConflicts potentialTxnConflicts)
      {
          List<IMutation> ms = new ArrayList<>();
          for (Map<ByteBuffer, IMutationBuilder> ksMap : mutationBuilders.values())
          {
              for (IMutationBuilder builder : ksMap.values())
              {
                  IMutation mutation = builder.build(potentialTxnConflicts);
                  mutation.validateIndexedColumns(state);
                  mutation.validateSize(MessagingService.current_version, CommitLogSegment.ENTRY_OVERHEAD_SIZE);
                  ms.add(mutation);
              }
          }
          return ms;
      }
  
      /**
       * Returns the key-mutation mappings for the specified keyspace.
       *
       * @param ksName the keyspace name
       * @return the key-mutation mappings for the specified keyspace.
       */
      private Map<ByteBuffer, IMutationBuilder> keyspaceMap(String ksName)
      {
          Map<ByteBuffer, IMutationBuilder> ksMap = mutationBuilders.get(ksName);
          if (ksMap == null)
          {
              ksMap = Maps.newHashMapWithExpectedSize(1);
              mutationBuilders.put(ksName, ksMap);
          }
          return ksMap;
      }
  
      private interface IMutationBuilder
      {
          /**
           * Add a new PartitionUpdate builder to this mutation builder
           * @param builder the builder to add
           * @return this
           */
          IMutationBuilder add(PartitionUpdate.Builder builder);
  
          /**
           * Build the immutable mutation
           */
          IMutation build(PotentialTxnConflicts potentialTxnConflicts);
  
          /**
           * Get the builder for the given tableId
           */
          PartitionUpdate.Builder get(TableId tableId);
      }
  
      private static class MutationBuilder implements IMutationBuilder
      {
          private final Map<TableId, PartitionUpdate.Builder> modifications;
          private final DecoratedKey key;
          private final String keyspaceName;
          private final long createdAt = approxTime.now();
  
          private MutationBuilder(String keyspaceName, DecoratedKey key, int initialSize)
          {
              this.keyspaceName = keyspaceName;
              this.key = key;
              this.modifications = Maps.newHashMapWithExpectedSize(initialSize);
          }
  
          public MutationBuilder add(PartitionUpdate.Builder updateBuilder)
          {
              assert updateBuilder != null;
              assert updateBuilder.partitionKey().getPartitioner() == key.getPartitioner();
              PartitionUpdate.Builder prev = modifications.put(updateBuilder.metadata().id, updateBuilder);
              if (prev != null)
                  // developer error
                  throw new IllegalArgumentException("Table " + updateBuilder.metadata().name + " already has modifications in this mutation: " + prev);
              return this;
          }
  
          public Mutation build(PotentialTxnConflicts potentialTxnConflicts)
          {
              ImmutableMap.Builder<TableId, PartitionUpdate> updates = new ImmutableMap.Builder<>();
              for (Map.Entry<TableId, PartitionUpdate.Builder> updateEntry : modifications.entrySet())
              {
                  PartitionUpdate update = updateEntry.getValue().build();
                  updates.put(updateEntry.getKey(), update);
              }
              return new Mutation(keyspaceName, key, updates.build(), createdAt, potentialTxnConflicts);
          }
  
          public PartitionUpdate.Builder get(TableId tableId)
          {
              return modifications.get(tableId);
          }
  
          public DecoratedKey key()
          {
              return key;
          }
  
          public boolean isEmpty()
          {
              return modifications.isEmpty();
          }
  
          public String getKeyspaceName()
          {
              return keyspaceName;
          }
      }
  
      private static class CounterMutationBuilder implements IMutationBuilder
      {
          private final MutationBuilder mutationBuilder;
          private final ConsistencyLevel cl;
  
          private CounterMutationBuilder(MutationBuilder mutationBuilder, ConsistencyLevel cl)
          {
              this.mutationBuilder = mutationBuilder;
              this.cl = cl;
          }
  
          public IMutationBuilder add(PartitionUpdate.Builder builder)
          {
              return mutationBuilder.add(builder);
          }
  
          public IMutation build(PotentialTxnConflicts potentialTxnConflicts)
          {
              return new CounterMutation(mutationBuilder.build(potentialTxnConflicts), cl);
          }
  
          public PartitionUpdate.Builder get(TableId id)
          {
              return mutationBuilder.get(id);
          }
      }
  
      private static class VirtualMutationBuilder implements IMutationBuilder
      {
          private final String keyspaceName;
          private final DecoratedKey partitionKey;
  
          private final HashMap<TableId, PartitionUpdate.Builder> modifications = new HashMap<>();
  
          private VirtualMutationBuilder(String keyspaceName, DecoratedKey partitionKey)
          {
              this.keyspaceName = keyspaceName;
              this.partitionKey = partitionKey;
          }
  
          @Override
          public VirtualMutationBuilder add(PartitionUpdate.Builder builder)
          {
              PartitionUpdate.Builder prev = modifications.put(builder.metadata().id, builder);
              if (null != prev)
                  throw new IllegalStateException();
              return this;
          }
  
          @Override
          public VirtualMutation build(PotentialTxnConflicts potentialTxnConflicts)
          {
              ImmutableMap.Builder<TableId, PartitionUpdate> updates = new ImmutableMap.Builder<>();
              modifications.forEach((tableId, updateBuilder) -> updates.put(tableId, updateBuilder.build()));
              return new VirtualMutation(keyspaceName, partitionKey, updates.build());
          }
  
          @Override
          public PartitionUpdate.Builder get(TableId tableId)
          {
              return modifications.get(tableId);
          }
      }
  }
  