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.kafka.clients.producer.internals;
  
  import org.apache.kafka.common.TopicPartition;
  import org.apache.kafka.common.requests.ProduceResponse;
  import org.apache.kafka.common.utils.LogContext;
  import org.apache.kafka.common.utils.ProducerIdAndEpoch;
  
  import org.slf4j.Logger;
  
  import java.util.HashMap;
  import java.util.Map;
  import java.util.OptionalInt;
  import java.util.OptionalLong;
  
  class TxnPartitionMap {
  
      private final Logger log;
  
      private final Map<TopicPartition, TxnPartitionEntry> topicPartitions = new HashMap<>();
  
      TxnPartitionMap(LogContext logContext) {
          this.log = logContext.logger(TxnPartitionMap.class);
      }
  
      TxnPartitionEntry get(TopicPartition topicPartition) {
          TxnPartitionEntry ent = topicPartitions.get(topicPartition);
          if (ent == null) {
              throw new IllegalStateException("Trying to get txnPartitionEntry for " + topicPartition +
                  ", but it was never set for this partition.");
          }
          return ent;
      }
  
      TxnPartitionEntry getOrCreate(TopicPartition topicPartition) {
          return topicPartitions.computeIfAbsent(topicPartition, TxnPartitionEntry::new);
      }
  
      boolean contains(TopicPartition topicPartition) {
          return topicPartitions.containsKey(topicPartition);
      }
  
      void reset() {
          topicPartitions.clear();
      }
  
      OptionalLong lastAckedOffset(TopicPartition topicPartition) {
          TxnPartitionEntry entry = topicPartitions.get(topicPartition);
          if (entry != null)
              return entry.lastAckedOffset();
          return OptionalLong.empty();
      }
  
      OptionalInt lastAckedSequence(TopicPartition topicPartition) {
          TxnPartitionEntry entry = topicPartitions.get(topicPartition);
          if (entry != null)
              return entry.lastAckedSequence();
          return OptionalInt.empty();
      }
  
      void startSequencesAtBeginning(TopicPartition topicPartition, ProducerIdAndEpoch newProducerIdAndEpoch) {
          TxnPartitionEntry entry = get(topicPartition);
          if (entry != null)
              entry.startSequencesAtBeginning(newProducerIdAndEpoch);
      }
  
      void remove(TopicPartition topicPartition) {
          topicPartitions.remove(topicPartition);
      }
  
  
      void updateLastAckedOffset(TopicPartition topicPartition, boolean isTransactional, long lastOffset) {
          OptionalLong lastAckedOffset = lastAckedOffset(topicPartition);
          // It might happen that the TransactionManager has been reset while a request was reenqueued and got a valid
          // response for this. This can happen only if the producer is only idempotent (not transactional) and in
          // this case there will be no tracked bookkeeper entry about it, so we have to insert one.
          if (lastAckedOffset.isEmpty() && !isTransactional)
              getOrCreate(topicPartition);
          if (lastOffset > lastAckedOffset.orElse(ProduceResponse.INVALID_OFFSET))
              get(topicPartition).setLastAckedOffset(lastOffset);
          else
              log.trace("Partition {} keeps lastOffset at {}", topicPartition, lastOffset);
      }
  
      // If a batch is failed fatally, the sequence numbers for future batches bound for the partition must be adjusted
      // so that they don't fail with the OutOfOrderSequenceException.
      //
      // This method must only be called when we know that the batch is question has been unequivocally failed by the broker,
      // ie. it has received a confirmed fatal status code like 'Message Too Large' or something similar.
      void adjustSequencesDueToFailedBatch(ProducerBatch batch) {
          if (!contains(batch.topicPartition))
              // Sequence numbers are not being tracked for this partition. This could happen if the producer id was just
              // reset due to a previous OutOfOrderSequenceException.
              return;
          log.debug("producerId: {}, send to partition {} failed fatally. Reducing future sequence numbers by {}",
                  batch.producerId(), batch.topicPartition, batch.recordCount);
  
          get(batch.topicPartition).adjustSequencesDueToFailedBatch(batch.baseSequence(), batch.recordCount);
      }
  
      int maybeUpdateLastAckedSequence(TopicPartition topicPartition, int sequence) {
          TxnPartitionEntry entry = topicPartitions.get(topicPartition);
          if (entry != null)
              return entry.maybeUpdateLastAckedSequence(sequence);
          return TxnPartitionEntry.NO_LAST_ACKED_SEQUENCE_NUMBER;
      }
  
      ProducerBatch nextBatchBySequence(TopicPartition topicPartition) {
          return get(topicPartition).nextBatchBySequence();
      }
  
      void removeInFlightBatch(ProducerBatch batch) {
          get(batch.topicPartition).removeInFlightBatch(batch);
      }
  }
  