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 kafka.server.share;
  
  import kafka.server.ReplicaManager;
  import kafka.server.share.SharePartitionManager.SharePartitionListener;
  
  import org.apache.kafka.common.KafkaException;
  import org.apache.kafka.common.TopicIdPartition;
  import org.apache.kafka.common.Uuid;
  import org.apache.kafka.common.errors.CoordinatorNotAvailableException;
  import org.apache.kafka.common.errors.FencedStateEpochException;
  import org.apache.kafka.common.errors.GroupIdNotFoundException;
  import org.apache.kafka.common.errors.InvalidRecordStateException;
  import org.apache.kafka.common.errors.InvalidRequestException;
  import org.apache.kafka.common.errors.LeaderNotAvailableException;
  import org.apache.kafka.common.errors.NotLeaderOrFollowerException;
  import org.apache.kafka.common.errors.UnknownServerException;
  import org.apache.kafka.common.errors.UnknownTopicOrPartitionException;
  import org.apache.kafka.common.message.FetchResponseData;
  import org.apache.kafka.common.message.ShareFetchResponseData.AcquiredRecords;
  import org.apache.kafka.common.protocol.Errors;
  import org.apache.kafka.common.record.ControlRecordType;
  import org.apache.kafka.common.record.Record;
  import org.apache.kafka.common.record.RecordBatch;
  import org.apache.kafka.common.utils.Time;
  import org.apache.kafka.coordinator.group.GroupConfig;
  import org.apache.kafka.coordinator.group.GroupConfigManager;
  import org.apache.kafka.coordinator.group.ShareGroupAutoOffsetResetStrategy;
  import org.apache.kafka.server.share.acknowledge.ShareAcknowledgementBatch;
  import org.apache.kafka.server.share.fetch.DelayedShareFetchGroupKey;
  import org.apache.kafka.server.share.fetch.DelayedShareFetchKey;
  import org.apache.kafka.server.share.fetch.ShareAcquiredRecords;
  import org.apache.kafka.server.share.metrics.SharePartitionMetrics;
  import org.apache.kafka.server.share.persister.GroupTopicPartitionData;
  import org.apache.kafka.server.share.persister.PartitionAllData;
  import org.apache.kafka.server.share.persister.PartitionErrorData;
  import org.apache.kafka.server.share.persister.PartitionFactory;
  import org.apache.kafka.server.share.persister.PartitionIdLeaderEpochData;
  import org.apache.kafka.server.share.persister.PartitionStateBatchData;
  import org.apache.kafka.server.share.persister.Persister;
  import org.apache.kafka.server.share.persister.PersisterStateBatch;
  import org.apache.kafka.server.share.persister.ReadShareGroupStateParameters;
  import org.apache.kafka.server.share.persister.TopicData;
  import org.apache.kafka.server.share.persister.WriteShareGroupStateParameters;
  import org.apache.kafka.server.storage.log.FetchIsolation;
  import org.apache.kafka.server.storage.log.FetchPartitionData;
  import org.apache.kafka.server.util.timer.Timer;
  import org.apache.kafka.server.util.timer.TimerTask;
  import org.apache.kafka.storage.internals.log.LogOffsetMetadata;
  
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.ArrayList;
  import java.util.Comparator;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.Iterator;
  import java.util.List;
  import java.util.Map;
  import java.util.NavigableMap;
  import java.util.Objects;
  import java.util.Optional;
  import java.util.PriorityQueue;
  import java.util.Set;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.ConcurrentSkipListMap;
  import java.util.concurrent.atomic.AtomicBoolean;
  import java.util.concurrent.atomic.AtomicReference;
  import java.util.concurrent.locks.ReadWriteLock;
  import java.util.concurrent.locks.ReentrantReadWriteLock;
  
  import static kafka.server.share.ShareFetchUtils.offsetForEarliestTimestamp;
  import static kafka.server.share.ShareFetchUtils.offsetForLatestTimestamp;
  import static kafka.server.share.ShareFetchUtils.offsetForTimestamp;
  import static kafka.server.share.ShareFetchUtils.recordLockDurationMsOrDefault;
  
  /**
   * The SharePartition is used to track the state of a partition that is shared between multiple
   * consumers. The class maintains the state of the records that have been fetched from the leader
   * and are in-flight.
   */
  @SuppressWarnings({"ClassDataAbstractionCoupling", "ClassFanOutComplexity"})
  public class SharePartition {
  
      private static final Logger log = LoggerFactory.getLogger(SharePartition.class);
  
      /**
       * empty member id used to indicate when a record is not acquired by any member.
       */
      static final String EMPTY_MEMBER_ID = Uuid.ZERO_UUID.toString();
  
      /**
       * The SharePartitionState is used to track the state of the share partition. The state of the
       * share partition determines if the partition is ready to receive requests, be initialized with
       * persisted state, or has failed to initialize.
       */
      // Visible for testing
      enum SharePartitionState {
          /**
           * The share partition is empty and has not been initialized with persisted state.
           */
          EMPTY,
          /**
           * The share partition is initializing with persisted state.
           */
          INITIALIZING,
          /**
           * The share partition is active and ready to serve requests.
           */
          ACTIVE,
          /**
           * The share partition failed to initialize with persisted state.
           */
          FAILED,
          /**
           * The share partition is fenced and cannot be used.
           */
          FENCED
      }
  
      /**
       * The RecordState is used to track the state of a record that has been fetched from the leader.
       * The state of the records determines if the records should be re-delivered, move the next fetch
       * offset, or be state persisted to disk.
       */
      // Visible for testing
      enum RecordState {
          AVAILABLE((byte) 0),
          ACQUIRED((byte) 1),
          ACKNOWLEDGED((byte) 2),
          ARCHIVED((byte) 4);
  
          public final byte id;
  
          RecordState(byte id) {
              this.id = id;
          }
  
          /**
           * Validates that the <code>newState</code> is one of the valid transition from the current
           * {@code RecordState}.
           *
           * @param newState State into which requesting to transition; must be non-<code>null</code>
           *
           * @return {@code RecordState} <code>newState</code> if validation succeeds. Returning
           *         <code>newState</code> helps state assignment chaining.
           *
           * @throws IllegalStateException if the state transition validation fails.
           */
          public RecordState validateTransition(RecordState newState) throws IllegalStateException {
              Objects.requireNonNull(newState, "newState cannot be null");
              if (this == newState) {
                  throw new IllegalStateException("The state transition is invalid as the new state is"
                      + "the same as the current state");
              }
  
              if (this == ACKNOWLEDGED || this == ARCHIVED) {
                  throw new IllegalStateException("The state transition is invalid from the current state: " + this);
              }
  
              if (this == AVAILABLE && newState != ACQUIRED) {
                  throw new IllegalStateException("The state can only be transitioned to ACQUIRED from AVAILABLE");
              }
  
              // Either the transition is from Available -> Acquired or from Acquired -> Available/
              // Acknowledged/Archived.
              return newState;
          }
  
          public static RecordState forId(byte id) {
              return switch (id) {
                  case 0 -> AVAILABLE;
                  case 1 -> ACQUIRED;
                  case 2 -> ACKNOWLEDGED;
                  case 4 -> ARCHIVED;
                  default -> throw new IllegalArgumentException("Unknown record state id: " + id);
              };
          }
  
          public byte id() {
              return this.id;
          }
      }
  
      /**
       * The DeliveryCountOps is used to specify the behavior on the delivery count: increase, decrease,
       * or do nothing.
       */
      private enum DeliveryCountOps {
          INCREASE,
          DECREASE,
          NO_OP
      }
  
      /**
       * The group id of the share partition belongs to.
       */
      private final String groupId;
  
      /**
       * The topic id partition of the share partition.
       */
      private final TopicIdPartition topicIdPartition;
  
      /**
       * The leader epoch is used to track the partition epoch.
       */
      private final int leaderEpoch;
  
      /**
       * The in-flight record is used to track the state of a record that has been fetched from the
       * leader. The state of the record is used to determine if the record should be re-fetched or if it
       * can be acknowledged or archived. Once share partition start offset is moved then the in-flight
       * records prior to the start offset are removed from the cache. The cache holds data against the
       * first offset of the in-flight batch.
       */
      private final NavigableMap<Long, InFlightBatch> cachedState;
  
      /**
       * The lock is used to synchronize access to the in-flight records. The lock is used to ensure that
       * the in-flight records are accessed in a thread-safe manner.
       */
      private final ReadWriteLock lock;
  
      /**
       * The find next fetch offset is used to indicate if the next fetch offset should be recomputed.
       */
      private final AtomicBoolean findNextFetchOffset;
  
      /**
       * The lock to ensure that the same share partition does not enter a fetch queue
       * while another one is being fetched within the queue. The caller's id that acquires the fetch
       * lock is utilized for ensuring the above.
       */
      private final AtomicReference<Uuid> fetchLock;
  
      /**
       * The max in-flight messages is used to limit the number of records that can be in-flight at any
       * given time. The max in-flight messages is used to prevent the consumer from fetching too many
       * records from the leader and running out of memory.
       */
      private final int maxInFlightMessages;
  
      /**
       * The max delivery count is used to limit the number of times a record can be delivered to the
       * consumer. The max delivery count is used to prevent the consumer re-delivering the same record
       * indefinitely.
       */
      private final int maxDeliveryCount;
  
      /**
       * The group config manager is used to retrieve the values for dynamic group configurations
       */
      private final GroupConfigManager groupConfigManager;
  
      /**
       * This is the default value which is used unless the group has a configuration which overrides it.
       * The record lock duration is used to limit the duration for which a consumer can acquire a record.
       * Once this time period is elapsed, the record will be made available or archived depending on the delivery count.
       */
      private final int defaultRecordLockDurationMs;
  
      /**
       * Timer is used to implement acquisition lock on records that guarantees the movement of records from
       * acquired to available/archived state upon timeout
       */
      private final Timer timer;
  
      /**
       * Time is used to get the currentTime.
       */
      private final Time time;
  
      /**
       * The persister is used to persist the state of the share partition to disk.
       */
      private final Persister persister;
  
      /**
       * The listener is used to notify the share partition manager when the share partition state changes.
       */
      private final SharePartitionListener listener;
  
      /**
       * The load start time is used to track the time taken to load the share partition.
       */
      private final long loadStartTimeMs;
  
      /**
       * The share partition metrics is used to track the broker-side metrics for the share partition.
       */
      private final SharePartitionMetrics sharePartitionMetrics;
  
      /**
       * The share partition start offset specifies the partition start offset from which the records
       * are cached in the cachedState of the sharePartition.
       */
      private long startOffset;
  
      /**
       * The share partition end offset specifies the partition end offset from which the records
       * are already fetched.
       */
      private long endOffset;
  
      /**
       * The initial read gap offset tracks if there are any gaps in the in-flight batch during initial
       * read of the share partition state from the persister.
       */
      private InitialReadGapOffset initialReadGapOffset;
  
      /**
       * We maintain the latest fetch offset and its metadata to estimate the minBytes requirement more efficiently.
       */
      private final OffsetMetadata fetchOffsetMetadata;
  
      /**
       * The delayed share fetch key is used to track the delayed share fetch requests that are waiting
       * for the respective share partition.
       */
      private final DelayedShareFetchKey delayedShareFetchKey;
  
      /**
       * The state epoch is used to track the version of the state of the share partition.
       */
      private int stateEpoch;
  
      /**
       * The partition state is used to track the state of the share partition.
       */
      private SharePartitionState partitionState;
  
      /**
       * The fetch lock acquired time is used to track the time when the lock for share partition is acquired.
       */
      private long fetchLockAcquiredTimeMs;
  
      /**
       * The fetch lock released time is used to track the time when the lock for share partition is released.
       */
      private long fetchLockReleasedTimeMs;
  
      /**
       * The fetch lock idle duration is used to track the time for which the fetch lock is idle.
       */
      private long fetchLockIdleDurationMs;
  
      /**
       * The replica manager is used to check to see if any delayed share fetch request can be completed because of data
       * availability due to acquisition lock timeout.
       */
      private final ReplicaManager replicaManager;
  
      SharePartition(
          String groupId,
          TopicIdPartition topicIdPartition,
          int leaderEpoch,
          int maxInFlightMessages,
          int maxDeliveryCount,
          int defaultRecordLockDurationMs,
          Timer timer,
          Time time,
          Persister persister,
          ReplicaManager replicaManager,
          GroupConfigManager groupConfigManager,
          SharePartitionListener listener
      ) {
          this(groupId, topicIdPartition, leaderEpoch, maxInFlightMessages, maxDeliveryCount, defaultRecordLockDurationMs,
              timer, time, persister, replicaManager, groupConfigManager, SharePartitionState.EMPTY, listener,
              new SharePartitionMetrics(groupId, topicIdPartition.topic(), topicIdPartition.partition()));
      }
  
      // Visible for testing
      @SuppressWarnings("ParameterNumber")
      SharePartition(
          String groupId,
          TopicIdPartition topicIdPartition,
          int leaderEpoch,
          int maxInFlightMessages,
          int maxDeliveryCount,
          int defaultRecordLockDurationMs,
          Timer timer,
          Time time,
          Persister persister,
          ReplicaManager replicaManager,
          GroupConfigManager groupConfigManager,
          SharePartitionState sharePartitionState,
          SharePartitionListener listener,
          SharePartitionMetrics sharePartitionMetrics
      ) {
          this.groupId = groupId;
          this.topicIdPartition = topicIdPartition;
          this.leaderEpoch = leaderEpoch;
          this.maxInFlightMessages = maxInFlightMessages;
          this.maxDeliveryCount = maxDeliveryCount;
          this.cachedState = new ConcurrentSkipListMap<>();
          this.lock = new ReentrantReadWriteLock();
          this.findNextFetchOffset = new AtomicBoolean(false);
          this.fetchLock = new AtomicReference<>(null);
          this.defaultRecordLockDurationMs = defaultRecordLockDurationMs;
          this.timer = timer;
          this.time = time;
          this.loadStartTimeMs = time.hiResClockMs();
          this.persister = persister;
          this.partitionState = sharePartitionState;
          this.replicaManager = replicaManager;
          this.groupConfigManager = groupConfigManager;
          this.fetchOffsetMetadata = new OffsetMetadata();
          this.delayedShareFetchKey = new DelayedShareFetchGroupKey(groupId, topicIdPartition);
          this.listener = listener;
          this.sharePartitionMetrics = sharePartitionMetrics;
          this.registerGaugeMetrics();
      }
  
      /**
       * May initialize the share partition by reading the state from the persister. The share partition
       * is initialized only if the state is in the EMPTY state. If the share partition is in ACTIVE state,
       * the method completes the future successfully. For other states, the method completes the future
       * with exception, which might be re-triable.
       *
       * @return The method returns a future which is completed when the share partition is initialized
       *         or completes with an exception if the share partition is in non-initializable state.
       */
      public CompletableFuture<Void> maybeInitialize() {
          log.debug("Maybe initialize share partition: {}-{}", groupId, topicIdPartition);
          // Check if the share partition is already initialized.
          try {
              if (initializedOrThrowException()) return CompletableFuture.completedFuture(null);
          } catch (Exception e) {
              return CompletableFuture.failedFuture(e);
          }
  
          // If code reaches here then the share partition is not initialized. Initialize the share partition.
          // All the pending requests should wait to get completed before the share partition is initialized.
          // Attain lock while updating the state to avoid any concurrent requests to be processed.
          try {
              if (!emptyToInitialState()) return CompletableFuture.completedFuture(null);
          } catch (Exception e) {
              return CompletableFuture.failedFuture(e);
          }
  
          // The share partition is not initialized, hence try to initialize it. There shall be only one
          // request trying to initialize the share partition.
          CompletableFuture<Void> future = new CompletableFuture<>();
          // Initialize the share partition by reading the state from the persister.
          persister.readState(new ReadShareGroupStateParameters.Builder()
              .setGroupTopicPartitionData(new GroupTopicPartitionData.Builder<PartitionIdLeaderEpochData>()
                  .setGroupId(this.groupId)
                  .setTopicsData(List.of(new TopicData<>(topicIdPartition.topicId(),
                      List.of(PartitionFactory.newPartitionIdLeaderEpochData(topicIdPartition.partition(), leaderEpoch)))))
                  .build())
              .build()
          ).whenComplete((result, exception) -> {
              Throwable throwable = null;
              lock.writeLock().lock();
              try {
                  if (exception != null) {
                      log.error("Failed to initialize the share partition: {}-{}", groupId, topicIdPartition, exception);
                      throwable = exception;
                      return;
                  }
  
                  if (result == null || result.topicsData() == null || result.topicsData().size() != 1) {
                      log.error("Failed to initialize the share partition: {}-{}. Invalid state found: {}.",
                          groupId, topicIdPartition, result);
                      throwable = new IllegalStateException(String.format("Failed to initialize the share partition %s-%s", groupId, topicIdPartition));
                      return;
                  }
  
                  TopicData<PartitionAllData> state = result.topicsData().get(0);
                  if (state.topicId() != topicIdPartition.topicId() || state.partitions().size() != 1) {
                      log.error("Failed to initialize the share partition: {}-{}. Invalid topic partition response: {}.",
                          groupId, topicIdPartition, result);
                      throwable = new IllegalStateException(String.format("Failed to initialize the share partition %s-%s", groupId, topicIdPartition));
                      return;
                  }
  
                  PartitionAllData partitionData = state.partitions().get(0);
                  if (partitionData.partition() != topicIdPartition.partition()) {
                      log.error("Failed to initialize the share partition: {}-{}. Invalid partition response: {}.",
                          groupId, topicIdPartition, partitionData);
                      throwable = new IllegalStateException(String.format("Failed to initialize the share partition %s-%s", groupId, topicIdPartition));
                      return;
                  }
  
                  if (partitionData.errorCode() != Errors.NONE.code()) {
                      KafkaException ex = fetchPersisterError(partitionData.errorCode(), partitionData.errorMessage());
                      log.error("Failed to initialize the share partition: {}-{}. Exception occurred: {}.",
                          groupId, topicIdPartition, partitionData);
                      throwable = ex;
                      return;
                  }
  
                  startOffset = startOffsetDuringInitialization(partitionData.startOffset());
                  stateEpoch = partitionData.stateEpoch();
  
                  List<PersisterStateBatch> stateBatches = partitionData.stateBatches();
                  long gapStartOffset = -1;
                  // The previousBatchLastOffset is used to track the last offset of the previous batch.
                  // For the first batch that should ideally start from startOffset if there are no gaps,
                  // we assume the previousBatchLastOffset to be startOffset - 1.
                  long previousBatchLastOffset = startOffset - 1;
                  for (PersisterStateBatch stateBatch : stateBatches) {
                      if (stateBatch.firstOffset() < startOffset) {
                          log.error("Invalid state batch found for the share partition: {}-{}. The base offset: {}"
                                  + " is less than the start offset: {}.", groupId, topicIdPartition,
                              stateBatch.firstOffset(), startOffset);
                          throwable = new IllegalStateException(String.format("Failed to initialize the share partition %s-%s", groupId, topicIdPartition));
                          return;
                      }
                      if (gapStartOffset == -1 && stateBatch.firstOffset() > previousBatchLastOffset + 1) {
                          gapStartOffset = previousBatchLastOffset + 1;
                      }
                      previousBatchLastOffset = stateBatch.lastOffset();
                      InFlightBatch inFlightBatch = new InFlightBatch(EMPTY_MEMBER_ID, stateBatch.firstOffset(),
                          stateBatch.lastOffset(), RecordState.forId(stateBatch.deliveryState()), stateBatch.deliveryCount(), null);
                      cachedState.put(stateBatch.firstOffset(), inFlightBatch);
                      sharePartitionMetrics.recordInFlightBatchMessageCount(stateBatch.lastOffset() - stateBatch.firstOffset() + 1);
                  }
                  // Update the endOffset of the partition.
                  if (!cachedState.isEmpty()) {
                      // If the cachedState is not empty, findNextFetchOffset flag is set to true so that any AVAILABLE records
                      // in the cached state are not missed
                      findNextFetchOffset.set(true);
                      endOffset = cachedState.lastEntry().getValue().lastOffset();
                      // initialReadGapOffset is not required, if there are no gaps in the read state response
                      if (gapStartOffset != -1) {
                          initialReadGapOffset = new InitialReadGapOffset(endOffset, gapStartOffset);
                      }
                      // In case the persister read state RPC result contains no AVAILABLE records, we can update cached state
                      // and start/end offsets.
                      maybeUpdateCachedStateAndOffsets();
                  } else {
                      endOffset = startOffset;
                  }
                  // Set the partition state to Active and complete the future.
                  partitionState = SharePartitionState.ACTIVE;
              } catch (Exception e) {
                  throwable = e;
              } finally {
                  boolean isFailed = throwable != null;
                  if (isFailed) {
                      partitionState = SharePartitionState.FAILED;
                  }
                  // Release the lock.
                  lock.writeLock().unlock();
                  // Avoid triggering the listener for waiting share fetch requests in purgatory as the
                  // share partition manager keeps track of same and will trigger the listener for the
                  // respective share partition.
                  // Complete the future.
                  if (isFailed) {
                      future.completeExceptionally(throwable);
                  } else {
                      future.complete(null);
                  }
              }
          });
  
          return future;
      }
  
      /**
       * The next fetch offset is used to determine the next offset that should be fetched from the leader.
       * The offset should be the next offset after the last fetched batch but there could be batches/
       * offsets that are either released by acknowledge API or lock timed out hence the next fetch
       * offset might be different from the last batch next offset. Hence, method checks if the next
       * fetch offset should be recomputed else returns the last computed next fetch offset.
       *
       * @return The next fetch offset that should be fetched from the leader.
       */
      public long nextFetchOffset() {
          /*
          The logic for determining the next offset to fetch data from a Share Partition hinges on a
          flag called findNextFetchOffset. If this flag is set to true, then the next fetch offset
          should be re-computed, otherwise the next fetch offset is Share Partition End Offset + 1.
          The flag is set to true in the following cases:
          1. When some previously acquired records are acknowledged with type RELEASE.
          2. When the record lock duration expires for some acquired records.
          3. When some records are released on share session close.
          The re-computation of next fetch offset is done by iterating over the cachedState and finding
          the first available record. If no available record is found, then the next fetch offset is
          set to Share Partition End Offset + 1 and findNextFetchOffset flag is set to false.
          */
          lock.writeLock().lock();
          try {
              // When none of the records in the cachedState are in the AVAILABLE state, findNextFetchOffset will be false
              if (!findNextFetchOffset.get()) {
                  if (cachedState.isEmpty() || startOffset > cachedState.lastEntry().getValue().lastOffset()) {
                      // 1. When cachedState is empty, endOffset is set to the next offset of the last
                      // offset removed from batch, which is the next offset to be fetched.
                      // 2. When startOffset has moved beyond the in-flight records, startOffset and
                      // endOffset point to the LSO, which is the next offset to be fetched.
                      log.trace("The next fetch offset for the share partition {}-{} is {}", groupId, topicIdPartition, endOffset);
                      return endOffset;
                  } else {
                      log.trace("The next fetch offset for the share partition {}-{} is {}", groupId, topicIdPartition, endOffset + 1);
                      return endOffset + 1;
                  }
              }
  
              // If this piece of code is reached, it means that findNextFetchOffset is true
              if (cachedState.isEmpty() || startOffset > cachedState.lastEntry().getValue().lastOffset()) {
                  // If cachedState is empty, there is no need of re-computing next fetch offset in future fetch requests.
                  // Same case when startOffset has moved beyond the in-flight records, startOffset and endOffset point to the LSO
                  // and the cached state is fresh.
                  findNextFetchOffset.set(false);
                  log.trace("The next fetch offset for the share partition {}-{} is {}", groupId, topicIdPartition, endOffset);
                  return endOffset;
              }
  
              long nextFetchOffset = -1;
              long gapStartOffset = isInitialReadGapOffsetWindowActive() ? initialReadGapOffset.gapStartOffset() : -1;
              for (Map.Entry<Long, InFlightBatch> entry : cachedState.entrySet()) {
                  // Check if there exists any gap in the in-flight batch which needs to be fetched. If
                  // initialReadGapOffset's endOffset is equal to the share partition's endOffset, then
                  // only the initial gaps should be considered. Once share partition's endOffset is past
                  // initial read end offset then all gaps are anyway fetched.
                  if (isInitialReadGapOffsetWindowActive()) {
                      if (entry.getKey() > gapStartOffset) {
                          nextFetchOffset = gapStartOffset;
                          break;
                      }
                      gapStartOffset = entry.getValue().lastOffset() + 1;
                  }
  
                  // Check if the state is maintained per offset or batch. If the offsetState
                  // is not maintained then the batch state is used to determine the offsets state.
                  if (entry.getValue().offsetState() == null) {
                      if (entry.getValue().batchState() == RecordState.AVAILABLE) {
                          nextFetchOffset = entry.getValue().firstOffset();
                          break;
                      }
                  } else {
                      // The offset state is maintained hence find the next available offset.
                      for (Map.Entry<Long, InFlightState> offsetState : entry.getValue().offsetState().entrySet()) {
                          if (offsetState.getValue().state == RecordState.AVAILABLE) {
                              nextFetchOffset = offsetState.getKey();
                              break;
                          }
                      }
                      // Break from the outer loop if updated.
                      if (nextFetchOffset != -1) {
                          break;
                      }
                  }
              }
  
              // If nextFetchOffset is -1, then no AVAILABLE records are found in the cachedState, so there is no need of
              // re-computing next fetch offset in future fetch requests
              if (nextFetchOffset == -1) {
                  findNextFetchOffset.set(false);
                  nextFetchOffset = endOffset + 1;
              }
              log.trace("The next fetch offset for the share partition {}-{} is {}", groupId, topicIdPartition, nextFetchOffset);
              return nextFetchOffset;
          } finally {
              lock.writeLock().unlock();
          }
      }
  
      /**
       * Acquire the fetched records for the share partition. The acquired records are added to the
       * in-flight records and the next fetch offset is updated to the next offset that should be
       * fetched from the leader.
       * <p>
       * The method always acquire the full batch records. The cache state can consist of multiple
       * full batches as a single batch. This behavior is driven by client configurations (batch size
       * and max fetch records) and allows for efficient client acknowledgements. However, partial batches
       * can exist in the cache only after a leader change and partial acknowledgements have been persisted
       * prior leader change. In such case, when a share partition loses track of a batch's start and
       * end offsets (e.g., after a leader change and partial acknowledgements), the cache stores the
       * batch based on the offset range provided by the persister. This method handles these special
       * batches by maintaining this range up to the last offset returned by the persister.  No special
       * handling is required afterward; the cache will eventually return to managing full batches.
       * <p>
       * For compacted topics, batches may be non-contiguous, and records within cached batches may contain gaps.
       * Because this method operates at the batch level, it acquires entire batches and relies on the
       * client to report any gaps in the data. Whether non-contiguous batches are acquired depends on
       * the first and last offsets of the fetched batches. Batches outside of this boundary will never
       * be acquired. For instance, if fetched batches cover offsets [0-9 and 20-29], and the configured
       * batch size and maximum fetch records are large enough (greater than 30 in this example), the
       * intervening batch [10-19] will be acquired. Since full fetched batch is acquired, the client is
       * responsible for reporting any data gaps. However, if the [0-9] and [20-29] ranges are fetched
       * in separate calls to this method, the [10-19] batch will not be acquired and cannot exist in
       * the cache.
       * <p>
       * However, for compacted topics, previously acquired batches (e.g., due to acquisition lock timeout
       * or explicit client release) might become available for acquisition again. But subsequent fetches
       * may reveal that these batches, or parts of them, have been removed by compaction. Because this
       * method works with whole batches, the disappearance of individual offsets within a batch requires
       * no special handling; the batch will be re-acquired, and the client will report the gaps. But if
       * an entire batch has been compacted away, this method must archive it in the cache to allow the
       * Share Partition Start Offset (SPSO) to progress. This is accomplished by comparing the fetchOffset
       * (the offset from which the log was read) with the first base offset of the fetch response. Any
       * batches from fetchOffset to first base offset of the fetch response are archived.
       *
       * @param memberId           The member id of the client that is fetching the record.
       * @param batchSize          The number of records per acquired records batch.
       * @param maxFetchRecords    The maximum number of records that should be acquired, this is a soft
       *                           limit and the method might acquire more records than the maxFetchRecords,
       *                           if the records are already part of the same fetch batch.
       * @param fetchOffset        The fetch offset for which the records are fetched.
       * @param fetchPartitionData The fetched records for the share partition.
       * @param isolationLevel      The isolation level for the share fetch request.
       * @return The acquired records for the share partition.
       */
      @SuppressWarnings("cyclomaticcomplexity") // Consider refactoring to avoid suppression
      public ShareAcquiredRecords acquire(
          String memberId,
          int batchSize,
          int maxFetchRecords,
          long fetchOffset,
          FetchPartitionData fetchPartitionData,
          FetchIsolation isolationLevel
      ) {
          log.trace("Received acquire request for share partition: {}-{} memberId: {}", groupId, topicIdPartition, memberId);
          if (stateNotActive() || maxFetchRecords <= 0) {
              // Nothing to acquire.
              return ShareAcquiredRecords.empty();
          }
  
          RecordBatch lastBatch = fetchPartitionData.records.lastBatch().orElse(null);
          if (lastBatch == null) {
              // Nothing to acquire.
              return ShareAcquiredRecords.empty();
          }
  
          // We require the first batch of records to get the base offset. Stop parsing further
          // batches.
          RecordBatch firstBatch = fetchPartitionData.records.batches().iterator().next();
          lock.writeLock().lock();
          try {
              long baseOffset = firstBatch.baseOffset();
  
              // There might be cached batches which are stale due to topic compaction hence archive them.
              maybeArchiveStaleBatches(fetchOffset, baseOffset);
  
              // Find the floor batch record for the request batch. The request batch could be
              // for a subset of the in-flight batch i.e. cached batch of offset 10-14 and request batch
              // of 12-13. Hence, floor entry is fetched to find the sub-map.
              Map.Entry<Long, InFlightBatch> floorEntry = cachedState.floorEntry(baseOffset);
              // We might find a batch with floor entry but not necessarily that batch has an overlap,
              // if the request batch base offset is ahead of last offset from floor entry i.e. cached
              // batch of 10-14 and request batch of 15-18, though floor entry is found but no overlap.
              // Such scenario will be handled in the next step when considering the subMap. However,
              // if the floor entry is found and the request batch base offset is within the floor entry
              // then adjust the base offset to the floor entry so that acquire method can still work on
              // previously cached batch boundaries.
              if (floorEntry != null && floorEntry.getValue().lastOffset() >= baseOffset) {
                  baseOffset = floorEntry.getKey();
              }
              // Validate if the fetch records are already part of existing batches and if available.
              NavigableMap<Long, InFlightBatch> subMap = cachedState.subMap(baseOffset, true, lastBatch.lastOffset(), true);
              // No overlap with request offsets in the cache for in-flight records. Acquire the complete
              // batch.
              if (subMap.isEmpty()) {
                  log.trace("No cached data exists for the share partition for requested fetch batch: {}-{}",
                      groupId, topicIdPartition);
                  ShareAcquiredRecords shareAcquiredRecords = acquireNewBatchRecords(memberId, fetchPartitionData.records.batches(),
                      firstBatch.baseOffset(), lastBatch.lastOffset(), batchSize, maxFetchRecords);
                  return maybeFilterAbortedTransactionalAcquiredRecords(fetchPartitionData, isolationLevel, shareAcquiredRecords);
              }
  
              log.trace("Overlap exists with in-flight records. Acquire the records if available for"
                  + " the share partition: {}-{}", groupId, topicIdPartition);
              List<AcquiredRecords> result = new ArrayList<>();
              // The acquired count is used to track the number of records acquired for the request.
              int acquiredCount = 0;
              // This tracks whether there is a gap between the subMap entries. If a gap is found, we will acquire
              // the corresponding offsets in a separate batch.
              long maybeGapStartOffset = baseOffset;
              // The fetched records are already part of the in-flight records. The records might
              // be available for re-delivery hence try acquiring same. The request batches could
              // be an exact match, subset or span over multiple already fetched batches.
              for (Map.Entry<Long, InFlightBatch> entry : subMap.entrySet()) {
                  // If the acquired count is equal to the max fetch records then break the loop.
                  if (acquiredCount >= maxFetchRecords) {
                      break;
                  }
  
                  InFlightBatch inFlightBatch = entry.getValue();
                  // If the initialReadGapOffset window is active, we need to treat the gaps in between the window as
                  // acquirable. Once the window is inactive (when we have acquired all the gaps inside the window),
                  // the remaining gaps are natural (data does not exist at those offsets) and we need not acquire them.
                  if (isInitialReadGapOffsetWindowActive()) {
                      // If nextBatchStartOffset is less than the key of the entry, this means the fetch happened for a gap in the cachedState.
                      // Thus, a new batch needs to be acquired for the gap.
                      if (maybeGapStartOffset < entry.getKey()) {
                          ShareAcquiredRecords shareAcquiredRecords = acquireNewBatchRecords(memberId, fetchPartitionData.records.batches(),
                              maybeGapStartOffset, entry.getKey() - 1, batchSize, maxFetchRecords);
                          result.addAll(shareAcquiredRecords.acquiredRecords());
                          acquiredCount += shareAcquiredRecords.count();
                      }
                      // Set nextBatchStartOffset as the last offset of the current in-flight batch + 1
                      maybeGapStartOffset = inFlightBatch.lastOffset() + 1;
                      // If the acquired count is equal to the max fetch records then break the loop.
                      if (acquiredCount >= maxFetchRecords) {
                          break;
                      }
                  }
  
                  // Compute if the batch is a full match.
                  boolean fullMatch = checkForFullMatch(inFlightBatch, firstBatch.baseOffset(), lastBatch.lastOffset());
  
                  if (!fullMatch || inFlightBatch.offsetState() != null) {
                      log.trace("Subset or offset tracked batch record found for share partition,"
                              + " batch: {} request offsets - first: {}, last: {} for the share"
                              + " partition: {}-{}", inFlightBatch, firstBatch.baseOffset(),
                          lastBatch.lastOffset(), groupId, topicIdPartition);
                      if (inFlightBatch.offsetState() == null) {
                          // Though the request is a subset of in-flight batch but the offset
                          // tracking has not been initialized yet which means that we could only
                          // acquire subset of offsets from the in-flight batch but only if the
                          // complete batch is available yet. Hence, do a pre-check to avoid exploding
                          // the in-flight offset tracking unnecessarily.
                          if (inFlightBatch.batchState() != RecordState.AVAILABLE || inFlightBatch.batchHasOngoingStateTransition()) {
                              log.trace("The batch is not available to acquire in share partition: {}-{}, skipping: {}"
                                      + " skipping offset tracking for batch as well.", groupId,
                                  topicIdPartition, inFlightBatch);
                              continue;
                          }
                          // The request batch is a subset or per offset state is managed hence update
                          // the offsets state in the in-flight batch.
                          inFlightBatch.maybeInitializeOffsetStateUpdate();
                      }
                      // Do not send max fetch records to acquireSubsetBatchRecords as we want to acquire
                      // all the records from the batch as the batch will anyway be part of the file-records
                      // response batch.
                      int acquiredSubsetCount = acquireSubsetBatchRecords(memberId, firstBatch.baseOffset(), lastBatch.lastOffset(), inFlightBatch, result);
                      acquiredCount += acquiredSubsetCount;
                      continue;
                  }
  
                  // The in-flight batch is a full match hence change the state of the complete batch.
                  if (inFlightBatch.batchState() != RecordState.AVAILABLE || inFlightBatch.batchHasOngoingStateTransition()) {
                      log.trace("The batch is not available to acquire in share partition: {}-{}, skipping: {}",
                          groupId, topicIdPartition, inFlightBatch);
                      continue;
                  }
  
                  InFlightState updateResult = inFlightBatch.tryUpdateBatchState(RecordState.ACQUIRED, DeliveryCountOps.INCREASE, maxDeliveryCount, memberId);
                  if (updateResult == null) {
                      log.info("Unable to acquire records for the batch: {} in share partition: {}-{}",
                          inFlightBatch, groupId, topicIdPartition);
                      continue;
                  }
                  // Schedule acquisition lock timeout for the batch.
                  AcquisitionLockTimerTask acquisitionLockTimeoutTask = scheduleAcquisitionLockTimeout(memberId, inFlightBatch.firstOffset(), inFlightBatch.lastOffset());
                  // Set the acquisition lock timeout task for the batch.
                  inFlightBatch.updateAcquisitionLockTimeout(acquisitionLockTimeoutTask);
  
                  result.add(new AcquiredRecords()
                      .setFirstOffset(inFlightBatch.firstOffset())
                      .setLastOffset(inFlightBatch.lastOffset())
                      .setDeliveryCount((short) inFlightBatch.batchDeliveryCount()));
                  acquiredCount += (int) (inFlightBatch.lastOffset() - inFlightBatch.firstOffset() + 1);
              }
  
              // Some of the request offsets are not found in the fetched batches. Acquire the
              // missing records as well.
              if (acquiredCount < maxFetchRecords && subMap.lastEntry().getValue().lastOffset() < lastBatch.lastOffset()) {
                  log.trace("There exists another batch which needs to be acquired as well");
                  ShareAcquiredRecords shareAcquiredRecords = acquireNewBatchRecords(memberId, fetchPartitionData.records.batches(),
                      subMap.lastEntry().getValue().lastOffset() + 1,
                      lastBatch.lastOffset(), batchSize, maxFetchRecords - acquiredCount);
                  result.addAll(shareAcquiredRecords.acquiredRecords());
                  acquiredCount += shareAcquiredRecords.count();
              }
              if (!result.isEmpty()) {
                  maybeUpdateReadGapFetchOffset(result.get(result.size() - 1).lastOffset() + 1);
                  return maybeFilterAbortedTransactionalAcquiredRecords(fetchPartitionData, isolationLevel, new ShareAcquiredRecords(result, acquiredCount));
              }
              return new ShareAcquiredRecords(result, acquiredCount);
          } finally {
              lock.writeLock().unlock();
          }
      }
  
      /**
       * Acknowledge the fetched records for the share partition. The accepted batches are removed
       * from the in-flight records once persisted. The next fetch offset is updated to the next offset
       * that should be fetched from the leader, if required.
       *
       * @param memberId               The member id of the client that is fetching the record.
       * @param acknowledgementBatches The acknowledgement batch list for the share partition.
       * @return A future which is completed when the records are acknowledged.
       */
      public CompletableFuture<Void> acknowledge(
          String memberId,
          List<ShareAcknowledgementBatch> acknowledgementBatches
      ) {
          log.trace("Acknowledgement batch request for share partition: {}-{}", groupId, topicIdPartition);
  
          CompletableFuture<Void> future = new CompletableFuture<>();
          Throwable throwable = null;
          List<InFlightState> updatedStates = new ArrayList<>();
          List<PersisterStateBatch> stateBatches = new ArrayList<>();
          lock.writeLock().lock();
          try {
              // Avoided using enhanced for loop as need to check if the last batch have offsets
              // in the range.
              for (ShareAcknowledgementBatch batch : acknowledgementBatches) {
                  // Client can either send a single entry in acknowledgeTypes which represents the state
                  // of the complete batch or can send individual offsets state.
                  Map<Long, RecordState> recordStateMap;
                  try {
                      recordStateMap = fetchRecordStateMapForAcknowledgementBatch(batch);
                  } catch (IllegalArgumentException e) {
                      log.debug("Invalid acknowledge type: {} for share partition: {}-{}",
                          batch.acknowledgeTypes(), groupId, topicIdPartition);
                      throwable = new InvalidRequestException("Invalid acknowledge type: " + batch.acknowledgeTypes());
                      break;
                  }
  
                  if (batch.lastOffset() < startOffset) {
                      log.trace("All offsets in the acknowledgement batch {} are already archived: {}-{}",
                          batch, groupId, topicIdPartition);
                      continue;
                  }
  
                  // Fetch the sub-map from the cached map for the batch to acknowledge. The sub-map can
                  // be a full match, subset or spans over multiple fetched batches.
                  NavigableMap<Long, InFlightBatch> subMap;
                  try {
                      subMap = fetchSubMapForAcknowledgementBatch(batch);
                  } catch (InvalidRecordStateException | InvalidRequestException e) {
                      throwable = e;
                      break;
                  }
  
                  // Acknowledge the records for the batch.
                  Optional<Throwable> ackThrowable = acknowledgeBatchRecords(
                      memberId,
                      batch,
                      recordStateMap,
                      subMap,
                      updatedStates,
                      stateBatches
                  );
  
                  if (ackThrowable.isPresent()) {
                      throwable = ackThrowable.get();
                      break;
                  }
              }
  
              // If the acknowledgement is successful then persist state, complete the state transition
              // and update the cached state for start offset. Else rollback the state transition.
              rollbackOrProcessStateUpdates(future, throwable, updatedStates, stateBatches);
          } finally {
              lock.writeLock().unlock();
          }
  
          return future;
      }
  
      /**
       * Release the acquired records for the share partition. The next fetch offset is updated to the next offset
       * that should be fetched from the leader.
       *
       * @param memberId The member id of the client whose records shall be released.
       * @return A future which is completed when the records are released.
       */
      public CompletableFuture<Void> releaseAcquiredRecords(String memberId) {
          log.trace("Release acquired records request for share partition: {}-{} memberId: {}", groupId, topicIdPartition, memberId);
  
          CompletableFuture<Void> future = new CompletableFuture<>();
          Throwable throwable = null;
          List<InFlightState> updatedStates = new ArrayList<>();
          List<PersisterStateBatch> stateBatches = new ArrayList<>();
  
          lock.writeLock().lock();
          try {
              RecordState recordState = RecordState.AVAILABLE;
              // Iterate over multiple fetched batches. The state can vary per offset entry
              for (Map.Entry<Long, InFlightBatch> entry : cachedState.entrySet()) {
                  InFlightBatch inFlightBatch = entry.getValue();
  
                  if (inFlightBatch.offsetState() == null
                          && inFlightBatch.batchState() == RecordState.ACQUIRED
                          && inFlightBatch.batchMemberId().equals(memberId)
                          && checkForStartOffsetWithinBatch(inFlightBatch.firstOffset(), inFlightBatch.lastOffset())) {
                      // For the case when batch.firstOffset < start offset <= batch.lastOffset, we will be having some
                      // acquired records that need to move to archived state despite their delivery count.
                      inFlightBatch.maybeInitializeOffsetStateUpdate();
                  }
  
                  if (inFlightBatch.offsetState() != null) {
                      Optional<Throwable> releaseAcquiredRecordsThrowable = releaseAcquiredRecordsForPerOffsetBatch(memberId, inFlightBatch, recordState, updatedStates, stateBatches);
                      if (releaseAcquiredRecordsThrowable.isPresent()) {
                          throwable = releaseAcquiredRecordsThrowable.get();
                          break;
                      }
                      continue;
                  }
                  Optional<Throwable> releaseAcquiredRecordsThrowable = releaseAcquiredRecordsForCompleteBatch(memberId, inFlightBatch, recordState, updatedStates, stateBatches);
                  if (releaseAcquiredRecordsThrowable.isPresent()) {
                      throwable = releaseAcquiredRecordsThrowable.get();
                      break;
                  }
              }
  
              // If the release acquired records is successful then persist state, complete the state transition
              // and update the cached state for start offset. Else rollback the state transition.
              rollbackOrProcessStateUpdates(future, throwable, updatedStates, stateBatches);
          } finally {
              lock.writeLock().unlock();
          }
          return future;
      }
  
      long loadStartTimeMs() {
          return loadStartTimeMs;
      }
  
      private Optional<Throwable> releaseAcquiredRecordsForPerOffsetBatch(String memberId,
                                                                          InFlightBatch inFlightBatch,
                                                                          RecordState recordState,
                                                                          List<InFlightState> updatedStates,
                                                                          List<PersisterStateBatch> stateBatches) {
  
          log.trace("Offset tracked batch record found, batch: {} for the share partition: {}-{}", inFlightBatch,
                  groupId, topicIdPartition);
          for (Map.Entry<Long, InFlightState> offsetState : inFlightBatch.offsetState.entrySet()) {
  
              // Check if member id is the owner of the offset.
              if (!offsetState.getValue().memberId().equals(memberId) && !offsetState.getValue().memberId().equals(EMPTY_MEMBER_ID)) {
                  log.debug("Member {} is not the owner of offset: {} in batch: {} for the share"
                          + " partition: {}-{}. Skipping offset.", memberId, offsetState.getKey(), inFlightBatch, groupId, topicIdPartition);
                  return Optional.empty();
              }
              if (offsetState.getValue().state == RecordState.ACQUIRED) {
                  // These records were fetched but they were not actually delivered to the client.
                  InFlightState updateResult = offsetState.getValue().startStateTransition(
                          offsetState.getKey() < startOffset ? RecordState.ARCHIVED : recordState,
                          DeliveryCountOps.DECREASE,
                          this.maxDeliveryCount,
                          EMPTY_MEMBER_ID
                  );
                  if (updateResult == null) {
                      log.debug("Unable to release records from acquired state for the offset: {} in batch: {}"
                                      + " for the share partition: {}-{}", offsetState.getKey(),
                              inFlightBatch, groupId, topicIdPartition);
                      return Optional.of(new InvalidRecordStateException("Unable to release acquired records for the offset"));
                  }
  
                  // Successfully updated the state of the offset.
                  updatedStates.add(updateResult);
                  stateBatches.add(new PersisterStateBatch(offsetState.getKey(), offsetState.getKey(),
                          updateResult.state.id, (short) updateResult.deliveryCount));
  
                  // If the maxDeliveryCount limit has been exceeded, the record will be transitioned to ARCHIVED state.
                  // This should not change the next fetch offset because the record is not available for acquisition
                  if (updateResult.state != RecordState.ARCHIVED) {
                      findNextFetchOffset.set(true);
                  }
              }
          }
          return Optional.empty();
      }
  
      private Optional<Throwable> releaseAcquiredRecordsForCompleteBatch(String memberId,
                                                                         InFlightBatch inFlightBatch,
                                                                         RecordState recordState,
                                                                         List<InFlightState> updatedStates,
                                                                         List<PersisterStateBatch> stateBatches) {
  
          // Check if member id is the owner of the batch.
          if (!inFlightBatch.batchMemberId().equals(memberId) && !inFlightBatch.batchMemberId().equals(EMPTY_MEMBER_ID)) {
              log.debug("Member {} is not the owner of batch record {} for share partition: {}-{}. Skipping batch.",
                      memberId, inFlightBatch, groupId, topicIdPartition);
              return Optional.empty();
          }
  
          // Change the state of complete batch since the same state exists for the entire inFlight batch.
          log.trace("Releasing acquired records for complete batch {} for the share partition: {}-{}",
                  inFlightBatch, groupId, topicIdPartition);
  
          if (inFlightBatch.batchState() == RecordState.ACQUIRED) {
              InFlightState updateResult = inFlightBatch.startBatchStateTransition(
                      inFlightBatch.lastOffset() < startOffset ? RecordState.ARCHIVED : recordState,
                      DeliveryCountOps.DECREASE,
                      this.maxDeliveryCount,
                      EMPTY_MEMBER_ID
              );
              if (updateResult == null) {
                  log.debug("Unable to release records from acquired state for the batch: {}"
                          + " for the share partition: {}-{}", inFlightBatch, groupId, topicIdPartition);
                  return Optional.of(new InvalidRecordStateException("Unable to release acquired records for the batch"));
              }
  
              // Successfully updated the state of the batch.
              updatedStates.add(updateResult);
              stateBatches.add(new PersisterStateBatch(inFlightBatch.firstOffset(), inFlightBatch.lastOffset(),
                      updateResult.state.id, (short) updateResult.deliveryCount));
  
              // If the maxDeliveryCount limit has been exceeded, the record will be transitioned to ARCHIVED state.
              // This should not change the next fetch offset because the record is not available for acquisition
              if (updateResult.state != RecordState.ARCHIVED) {
                  findNextFetchOffset.set(true);
              }
          }
          return Optional.empty();
      }
  
      /**
       * Updates the cached state, start and end offsets of the share partition as per the new log
       * start offset. The method is called when the log start offset is moved for the share partition.
       *
       * @param logStartOffset The new log start offset.
       */
      void updateCacheAndOffsets(long logStartOffset) {
          lock.writeLock().lock();
          try {
              if (logStartOffset <= startOffset) {
                  log.error("The log start offset: {} is not greater than the start offset: {} for the share partition: {}-{}",
                          logStartOffset, startOffset, groupId, topicIdPartition);
                  return;
              }
              log.debug("Updating start offset for share partition: {}-{} from: {} to: {} since LSO has moved to: {}",
                      groupId, topicIdPartition, startOffset, logStartOffset, logStartOffset);
              if (cachedState.isEmpty()) {
                  // If the cached state is empty, then the start and end offset will be the new log start offset.
                  // This can occur during the initialization of share partition if LSO has moved.
                  startOffset = logStartOffset;
                  endOffset = logStartOffset;
                  return;
              }
  
              // Archive the available records in the cached state that are before the new log start offset.
              boolean anyRecordArchived = archiveAvailableRecordsOnLsoMovement(logStartOffset);
              // If we have transitioned the state of any batch/offset from AVAILABLE to ARCHIVED,
              // then there is a chance that the next fetch offset can change.
              if (anyRecordArchived) {
                  findNextFetchOffset.set(true);
              }
  
              // The new startOffset will be the log start offset.
              startOffset = logStartOffset;
              if (endOffset < startOffset) {
                  // This case means that the cached state is completely fresh now.
                  // Example scenario - batch of 0-10 in acquired state in cached state, then LSO moves to 15,
                  // then endOffset should be 15 as well.
                  endOffset = startOffset;
              }
  
              // Note -
              // 1. We will be writing the new starOffset lazily during acknowledge/release acquired records API call.
              // 2. We will not be writing the archived state batches to the persister.
          } finally {
              lock.writeLock().unlock();
          }
      }
  
      /**
       * The method archives the available records in the cached state that are between the fetch offset
       * and the base offset of the first fetched batch. This method is required to handle the compacted
       * topics where the already fetched batch which is marked re-available, might not result in subsequent
       * fetch response from log. Hence, the batches need to be archived to allow the SPSO and next fetch
       * offset to progress.
       *
       * @param fetchOffset The fetch offset.
       * @param baseOffset  The base offset of the first fetched batch.
       */
      private void maybeArchiveStaleBatches(long fetchOffset, long baseOffset) {
          lock.writeLock().lock();
          try {
              // If the fetch happens from within a batch then fetchOffset can be ahead of base offset else
              // should be same as baseOffset of the first fetched batch. Otherwise, we might need to archive
              // some stale batches.
              if (cachedState.isEmpty() || fetchOffset >= baseOffset) {
                  // No stale batches to archive.
                  return;
              }
  
              // The fetch offset can exist in the middle of the batch. Hence, find the floor offset
              // for the fetch offset and then find the sub-map from the floor offset to the base offset.
              long floorOffset = fetchOffset;
              Map.Entry<Long, InFlightBatch> floorEntry = cachedState.floorEntry(fetchOffset);
              if (floorEntry != null && floorEntry.getValue().lastOffset() >= fetchOffset) {
                  floorOffset = floorEntry.getKey();
              }
  
              NavigableMap<Long, InFlightBatch> subMap = cachedState.subMap(floorOffset, true, baseOffset, false);
              if (subMap.isEmpty()) {
                  // No stale batches to archive.
                  return;
              }
  
              // Though such batches can be removed from the cache, but it is better to archive them so
              // that they are never acquired again.
              boolean anyRecordArchived = archiveRecords(fetchOffset, baseOffset, subMap, RecordState.AVAILABLE);
  
              // If we have transitioned the state of any batch/offset from AVAILABLE to ARCHIVED,
              // then there is a chance that the next fetch offset can change.
              if (anyRecordArchived) {
                  findNextFetchOffset.set(true);
              }
          } finally {
              lock.writeLock().unlock();
          }
      }
  
      /**
       * The method archives the available records in the cached state that are before the log start offset.
       *
       * @param logStartOffset The log start offset.
       * @return A boolean which indicates whether any record is archived or not.
       */
      private boolean archiveAvailableRecordsOnLsoMovement(long logStartOffset) {
          lock.writeLock().lock();
          try {
              return archiveRecords(startOffset, logStartOffset, cachedState, RecordState.AVAILABLE);
          } finally {
              lock.writeLock().unlock();
          }
      }
  
      /**
       * The method archive the records in a given state in the map that are before the end offset.
       *
       * @param startOffset The offset from which the records should be archived.
       * @param endOffset The offset before which the records should be archived.
       * @param map The map containing the in-flight records.
       * @param initialState The initial state of the records to be archived.
       * @return A boolean which indicates whether any record is archived or not.
       */
      private boolean archiveRecords(long startOffset, long endOffset, NavigableMap<Long, InFlightBatch> map, RecordState initialState) {
          lock.writeLock().lock();
          try {
              boolean isAnyOffsetArchived = false, isAnyBatchArchived = false;
              for (Map.Entry<Long, InFlightBatch> entry : map.entrySet()) {
                  long batchStartOffset = entry.getKey();
                  // We do not need to transition state of batches/offsets that are later than the offset.
                  if (batchStartOffset >= endOffset) {
                      break;
                  }
                  InFlightBatch inFlightBatch = entry.getValue();
                  boolean fullMatch = checkForFullMatch(inFlightBatch, startOffset, endOffset - 1);
  
                  // Maintain state per offset if the inflight batch is not a full match or the offset state is managed.
                  if (!fullMatch || inFlightBatch.offsetState() != null) {
                      log.debug("Subset or offset tracked batch record found while trying to update offsets "
                          + "and cached state map, batch: {}, offsets to update - first: {}, last: {} "
                          + "for the share partition: {}-{}", inFlightBatch, startOffset, endOffset - 1,
                          groupId, topicIdPartition);
  
                      if (inFlightBatch.offsetState() == null) {
                          if (inFlightBatch.batchState() != initialState) {
                              continue;
                          }
                          inFlightBatch.maybeInitializeOffsetStateUpdate();
                      }
                      isAnyOffsetArchived = isAnyOffsetArchived || archivePerOffsetBatchRecords(inFlightBatch, startOffset, endOffset - 1, initialState);
                      continue;
                  }
                  // The in-flight batch is a full match hence change the state of the complete batch.
                  isAnyBatchArchived = isAnyBatchArchived || archiveCompleteBatch(inFlightBatch, initialState);
              }
              return isAnyOffsetArchived || isAnyBatchArchived;
          } finally {
              lock.writeLock().unlock();
          }
      }
  
      private boolean archivePerOffsetBatchRecords(InFlightBatch inFlightBatch,
                                                   long startOffsetToArchive,
                                                   long endOffsetToArchive,
                                                   RecordState initialState
      ) {
          lock.writeLock().lock();
          try {
              boolean isAnyOffsetArchived = false;
              log.trace("Archiving offset tracked batch: {} for the share partition: {}-{}", inFlightBatch, groupId, topicIdPartition);
              for (Map.Entry<Long, InFlightState> offsetState : inFlightBatch.offsetState().entrySet()) {
                  if (offsetState.getKey() < startOffsetToArchive) {
                      continue;
                  }
                  if (offsetState.getKey() > endOffsetToArchive) {
                      // No further offsets to process.
                      break;
                  }
                  if (offsetState.getValue().state != initialState) {
                      continue;
                  }
  
                  offsetState.getValue().archive(EMPTY_MEMBER_ID);
                  if (initialState == RecordState.ACQUIRED) {
                      offsetState.getValue().cancelAndClearAcquisitionLockTimeoutTask();
                  }
                  isAnyOffsetArchived = true;
              }
              return isAnyOffsetArchived;
          } finally {
              lock.writeLock().unlock();
          }
      }
  
      private boolean archiveCompleteBatch(InFlightBatch inFlightBatch, RecordState initialState) {
          lock.writeLock().lock();
          try {
              log.trace("Archiving complete batch: {} for the share partition: {}-{}", inFlightBatch, groupId, topicIdPartition);
              if (inFlightBatch.batchState() == initialState) {
                  // Change the state of complete batch since the same state exists for the entire inFlight batch.
                  inFlightBatch.archiveBatch(EMPTY_MEMBER_ID);
                  if (initialState == RecordState.ACQUIRED) {
                      inFlightBatch.batchState.cancelAndClearAcquisitionLockTimeoutTask();
                  }
                  return true;
              }
          } finally {
              lock.writeLock().unlock();
          }
          return false;
      }
  
      /**
       * Checks if the records can be acquired for the share partition. The records can be acquired if
       * the number of records in-flight is less than the max in-flight messages. Or if the fetch is
       * to happen somewhere in between the record states cached in the share partition i.e. re-acquire
       * the records that are already fetched before.
       *
       * @return A boolean which indicates whether more records can be acquired or not.
       */
      boolean canAcquireRecords() {
          if (nextFetchOffset() != endOffset() + 1) {
              return true;
          }
          return numInFlightRecords() < maxInFlightMessages;
      }
  
      /**
       * Prior to fetching records from the leader, the fetch lock is acquired to ensure that the same
       * share partition is not fetched concurrently by multiple clients. The fetch lock is released once
       * the records are fetched and acquired.
       *
       * @param fetchId - the caller's id that is trying to acquire the fetch lock.
       * @return A boolean which indicates whether the fetch lock is acquired.
       */
      public boolean maybeAcquireFetchLock(Uuid fetchId) {
          if (stateNotActive()) {
              return false;
          }
          boolean acquired = fetchLock.compareAndSet(null, Objects.requireNonNull(fetchId));
          if (acquired) {
              long currentTime = time.hiResClockMs();
              fetchLockAcquiredTimeMs = currentTime;
              fetchLockIdleDurationMs = fetchLockReleasedTimeMs != 0 ? currentTime - fetchLockReleasedTimeMs : 0;
          }
          return acquired;
      }
  
      /**
       * Release the fetch lock once the records are fetched from the leader. It is imperative that the caller
       * that acquired the fetch lock should be the one releasing it.
       * @param fetchId - The caller's id that is trying to release the fetch lock.
       */
      void releaseFetchLock(Uuid fetchId) {
          // Register the metric for the duration the fetch lock was held. Do not register the metric
          // if the fetch lock was not acquired.
          long currentTime = time.hiResClockMs();
          if (!fetchLock.compareAndSet(Objects.requireNonNull(fetchId), null)) {
              // This code should not be reached unless we are in error-prone scenarios. Since we are releasing the fetch
              // lock for multiple share partitions at different places in DelayedShareFetch (due to tackling remote
              // storage fetch and local log fetch from a single purgatory), in order to safeguard ourselves from bad code,
              // we are logging when an instance that does not hold the fetch lock tries to release it.
              Uuid fetchLockAcquiredBy = fetchLock.getAndSet(null);
              log.info("Instance {} does not hold the fetch lock, yet trying to release it for share partition {}-{}. The lock was held by {}",
                  fetchId, groupId, topicIdPartition, fetchLockAcquiredBy);
          }
          long acquiredDurationMs = currentTime - fetchLockAcquiredTimeMs;
          // Update the metric for the fetch lock time.
          sharePartitionMetrics.recordFetchLockTimeMs(acquiredDurationMs);
          // Update fetch lock ratio metric.
          recordFetchLockRatioMetric(acquiredDurationMs);
          fetchLockReleasedTimeMs = currentTime;
      }
  
      /**
       * Marks the share partition as fenced.
       */
      void markFenced() {
          lock.writeLock().lock();
          try {
              partitionState = SharePartitionState.FENCED;
          } finally {
              lock.writeLock().unlock();
          }
      }
  
      /**
       * Returns the share partition listener.
       *
       * @return The share partition listener.
       */
      SharePartitionListener listener() {
          return this.listener;
      }
  
      int leaderEpoch() {
          return leaderEpoch;
      }
  
      /**
       * Records the fetch lock ratio metric. The metric is the ratio of the time duration the fetch
       * lock was acquired to the total time since the last lock acquisition. The total time is calculated
       * by adding the duration of the fetch lock idle time to the time the fetch lock was acquired.
       *
       * @param acquiredDurationMs The time duration the fetch lock was acquired.
       */
      // Visible for testing
      void recordFetchLockRatioMetric(long acquiredDurationMs) {
          if (fetchLockIdleDurationMs < 0) {
              // This is just a safe check to avoid negative time for fetch lock idle duration. This
              // should not happen in any scenarios. If it does then just return from the method and
              // no metric update is an indicator of the issue.
              return;
          }
  
          // Update the total fetch lock acquired time.
          double fetchLockToTotalTime;
          if (acquiredDurationMs + fetchLockIdleDurationMs == 0) {
              // If the total time is 0 then the ratio is 1 i.e. the fetch lock was acquired for the complete time.
              fetchLockToTotalTime = 1.0;
          } else if (acquiredDurationMs == 0) {
              // If the acquired duration is 0 then the ratio is the calculated by the idle duration.
              fetchLockToTotalTime = 1.0 / fetchLockIdleDurationMs;
          } else {
              fetchLockToTotalTime = acquiredDurationMs * (1.0 / (acquiredDurationMs + fetchLockIdleDurationMs));
          }
          sharePartitionMetrics.recordFetchLockRatio((int) (fetchLockToTotalTime * 100));
      }
  
      private void registerGaugeMetrics() {
          sharePartitionMetrics.registerInFlightMessageCount(this::numInFlightRecords);
          sharePartitionMetrics.registerInFlightBatchCount(this.cachedState::size);
      }
  
      private long numInFlightRecords() {
          lock.readLock().lock();
          long numRecords;
          try {
              if (cachedState.isEmpty()) {
                  numRecords = 0;
              } else {
                  numRecords = this.endOffset - this.startOffset + 1;
              }
          } finally {
              lock.readLock().unlock();
          }
          return numRecords;
      }
  
      private boolean stateNotActive() {
          return  partitionState() != SharePartitionState.ACTIVE;
      }
  
      private boolean emptyToInitialState() {
          lock.writeLock().lock();
          try {
              if (initializedOrThrowException()) return false;
              partitionState = SharePartitionState.INITIALIZING;
              return true;
          } finally {
              lock.writeLock().unlock();
          }
      }
  
      private boolean initializedOrThrowException() {
          SharePartitionState currentState = partitionState();
          return switch (currentState) {
              case ACTIVE -> true;
              case FAILED -> throw new IllegalStateException(
                  String.format("Share partition failed to load %s-%s", groupId, topicIdPartition));
              case INITIALIZING -> throw new LeaderNotAvailableException(
                  String.format("Share partition is already initializing %s-%s", groupId, topicIdPartition));
              case FENCED -> throw new FencedStateEpochException(
                  String.format("Share partition is fenced %s-%s", groupId, topicIdPartition));
              case EMPTY ->
                  // The share partition is not yet initialized.
                  false;
          };
      }
  
      // Method to reduce the window that tracks gaps in the cachedState
      private void maybeUpdateReadGapFetchOffset(long offset) {
          lock.writeLock().lock();
          try {
              if (initialReadGapOffset != null) {
                  if (initialReadGapOffset.endOffset() == endOffset) {
                      initialReadGapOffset.gapStartOffset(offset);
                  } else {
                      // The initial read gap offset is not valid anymore as the end offset has moved
                      // beyond the initial read gap offset. Hence, reset the initial read gap offset.
                      initialReadGapOffset = null;
                  }
              }
          } finally {
              lock.writeLock().unlock();
          }
      }
  
      private ShareAcquiredRecords acquireNewBatchRecords(
          String memberId,
          Iterable<? extends RecordBatch> batches,
          long firstOffset,
          long lastOffset,
          int batchSize,
          int maxFetchRecords
      ) {
          lock.writeLock().lock();
          try {
              // If same batch is fetched and previous batch is removed from the cache then we need to
              // update the batch first offset to endOffset, only if endOffset passed the firstOffset.
              // For an initial start of the share fetch from a topic partition the endOffset will be initialized
              // to 0 but firstOffset can be higher than 0.
              long firstAcquiredOffset = firstOffset;
              if (cachedState.isEmpty() && endOffset > firstAcquiredOffset) {
                  firstAcquiredOffset = endOffset;
              }
  
              // Check how many messages can be acquired from the batch.
              long lastAcquiredOffset = lastOffset;
              if (maxFetchRecords < lastAcquiredOffset - firstAcquiredOffset + 1) {
                  // The max messages to acquire is less than the complete available batches hence
                  // limit the acquired records. The last offset shall be the batches last offset
                  // which falls under the max messages limit. As the max fetch records is the soft
                  // limit, the last offset can be higher than the max messages.
                  lastAcquiredOffset = lastOffsetFromBatchWithRequestOffset(batches, firstAcquiredOffset + maxFetchRecords - 1);
              }
  
              // Create batches of acquired records.
              List<AcquiredRecords> acquiredRecords = createBatches(memberId, batches, firstAcquiredOffset, lastAcquiredOffset, batchSize);
              // if the cachedState was empty before acquiring the new batches then startOffset needs to be updated
              if (cachedState.firstKey() == firstAcquiredOffset)  {
                  startOffset = firstAcquiredOffset;
              }
  
              // If the new batch acquired is part of a gap in the cachedState, then endOffset should not be updated.
              // Ex. if startOffset is 10, endOffset is 30, there is a gap from 10 to 20, and an inFlight batch from 21 to 30.
              // In this case, the nextFetchOffset results in 10 and the records are fetched. A new batch is acquired from
              // 10 to 20, but the endOffset remains at 30.
              if (lastAcquiredOffset > endOffset) {
                  endOffset = lastAcquiredOffset;
              }
              maybeUpdateReadGapFetchOffset(lastAcquiredOffset + 1);
              return new ShareAcquiredRecords(acquiredRecords, (int) (lastAcquiredOffset - firstAcquiredOffset + 1));
          } finally {
              lock.writeLock().unlock();
          }
      }
  
      private List<AcquiredRecords> createBatches(
          String memberId,
          Iterable<? extends RecordBatch> batches,
          long firstAcquiredOffset,
          long lastAcquiredOffset,
          int batchSize
      ) {
          lock.writeLock().lock();
          try {
              List<AcquiredRecords> result = new ArrayList<>();
              long currentFirstOffset = firstAcquiredOffset;
              // No split of batches is required if the batch size is greater than records which
              // can be acquired, else split the batch into multiple batches.
              if (lastAcquiredOffset - firstAcquiredOffset + 1 > batchSize) {
                  // The batch is split into multiple batches considering batch size.
                  // Note: Try reading only the baseOffset of the batch and avoid reading the lastOffset
                  // as lastOffset call of RecordBatch is expensive (loads headers).
                  for (RecordBatch batch : batches) {
                      long batchBaseOffset = batch.baseOffset();
                      // Check if the batch is already past the last acquired offset then break.
                      if (batchBaseOffset > lastAcquiredOffset) {
                          // Break the loop and the last batch will be processed outside the loop.
                          break;
                      }
  
                      // Create new batch once the batch size is reached.
                      if (batchBaseOffset - currentFirstOffset >= batchSize) {
                          result.add(new AcquiredRecords()
                              .setFirstOffset(currentFirstOffset)
                              .setLastOffset(batchBaseOffset - 1)
                              .setDeliveryCount((short) 1));
                          currentFirstOffset = batchBaseOffset;
                      }
                  }
              }
              // Add the last batch or the only batch if the batch size is greater than the records which
              // can be acquired.
              result.add(new AcquiredRecords()
                  .setFirstOffset(currentFirstOffset)
                  .setLastOffset(lastAcquiredOffset)
                  .setDeliveryCount((short) 1));
  
              result.forEach(acquiredRecords -> {
                  // Schedule acquisition lock timeout for the batch.
                  AcquisitionLockTimerTask timerTask = scheduleAcquisitionLockTimeout(memberId, acquiredRecords.firstOffset(), acquiredRecords.lastOffset());
                  // Add the new batch to the in-flight records along with the acquisition lock timeout task for the batch.
                  cachedState.put(acquiredRecords.firstOffset(), new InFlightBatch(
                      memberId,
                      acquiredRecords.firstOffset(),
                      acquiredRecords.lastOffset(),
                      RecordState.ACQUIRED,
                      1,
                      timerTask));
                  // Update the in-flight batch message count metrics for the share partition.
                  sharePartitionMetrics.recordInFlightBatchMessageCount(acquiredRecords.lastOffset() - acquiredRecords.firstOffset() + 1);
              });
              return result;
          } finally {
              lock.writeLock().unlock();
          }
      }
  
      private int acquireSubsetBatchRecords(
          String memberId,
          long requestFirstOffset,
          long requestLastOffset,
          InFlightBatch inFlightBatch,
          List<AcquiredRecords> result
      ) {
          lock.writeLock().lock();
          int acquiredCount = 0;
          try {
              for (Map.Entry<Long, InFlightState> offsetState : inFlightBatch.offsetState.entrySet()) {
                  // For the first batch which might have offsets prior to the request base
                  // offset i.e. cached batch of 10-14 offsets and request batch of 12-13.
                  if (offsetState.getKey() < requestFirstOffset) {
                      continue;
                  }
  
                  if (offsetState.getKey() > requestLastOffset) {
                      // No further offsets to process.
                      break;
                  }
  
                  if (offsetState.getValue().state != RecordState.AVAILABLE || offsetState.getValue().hasOngoingStateTransition()) {
                      log.trace("The offset {} is not available in share partition: {}-{}, skipping: {}",
                          offsetState.getKey(), groupId, topicIdPartition, inFlightBatch);
                      continue;
                  }
  
                  InFlightState updateResult =  offsetState.getValue().tryUpdateState(RecordState.ACQUIRED, DeliveryCountOps.INCREASE,
                      maxDeliveryCount, memberId);
                  if (updateResult == null) {
                      log.trace("Unable to acquire records for the offset: {} in batch: {}"
                              + " for the share partition: {}-{}", offsetState.getKey(), inFlightBatch,
                          groupId, topicIdPartition);
                      continue;
                  }
                  // Schedule acquisition lock timeout for the offset.
                  AcquisitionLockTimerTask acquisitionLockTimeoutTask = scheduleAcquisitionLockTimeout(memberId, offsetState.getKey(), offsetState.getKey());
                  // Update acquisition lock timeout task for the offset.
                  offsetState.getValue().updateAcquisitionLockTimeoutTask(acquisitionLockTimeoutTask);
  
                  // TODO: Maybe we can club the continuous offsets here.
                  result.add(new AcquiredRecords()
                      .setFirstOffset(offsetState.getKey())
                      .setLastOffset(offsetState.getKey())
                      .setDeliveryCount((short) offsetState.getValue().deliveryCount));
                  acquiredCount++;
              }
          } finally {
              lock.writeLock().unlock();
          }
          return acquiredCount;
      }
  
      /**
       * Check if the in-flight batch is a full match with the request offsets. The full match represents
       * complete overlap of the in-flight batch with the request offsets.
       *
       * @param inFlightBatch The in-flight batch to check for full match.
       * @param firstOffsetToCompare The first offset of the request batch.
       * @param lastOffsetToCompare The last offset of the request batch.
       *
       * @return True if the in-flight batch is a full match with the request offsets, false otherwise.
       */
      private boolean checkForFullMatch(InFlightBatch inFlightBatch, long firstOffsetToCompare, long lastOffsetToCompare) {
          return inFlightBatch.firstOffset() >= firstOffsetToCompare && inFlightBatch.lastOffset() <= lastOffsetToCompare;
      }
  
      /**
       * Check if the start offset has moved and within the request first and last offset.
       *
       * @param batchFirstOffset The first offset of the batch.
       * @param batchLastOffset The last offset of the batch.
       *
       * @return True if the start offset has moved and within the request first and last offset, false otherwise.
       */
      private boolean checkForStartOffsetWithinBatch(long batchFirstOffset, long batchLastOffset) {
          long localStartOffset = startOffset();
          return batchFirstOffset < localStartOffset && batchLastOffset >= localStartOffset;
      }
  
      private Map<Long, RecordState> fetchRecordStateMapForAcknowledgementBatch(
          ShareAcknowledgementBatch batch) {
          // Client can either send a single entry in acknowledgeTypes which represents the state
          // of the complete batch or can send individual offsets state. Construct a map with record state
          // for each offset in the batch, if single acknowledge type is sent, the map will have only one entry.
          Map<Long, RecordState> recordStateMap = new HashMap<>();
          for (int index = 0; index < batch.acknowledgeTypes().size(); index++) {
              recordStateMap.put(batch.firstOffset() + index,
                  fetchRecordState(batch.acknowledgeTypes().get(index)));
          }
          return recordStateMap;
      }
  
      private static RecordState fetchRecordState(byte acknowledgeType) {
          return switch (acknowledgeType) {
              case 1 /* ACCEPT */ -> RecordState.ACKNOWLEDGED;
              case 2 /* RELEASE */ -> RecordState.AVAILABLE;
              case 3, 0 /* REJECT / GAP */ -> RecordState.ARCHIVED;
              default -> throw new IllegalArgumentException("Invalid acknowledge type: " + acknowledgeType);
          };
      }
  
      private NavigableMap<Long, InFlightBatch> fetchSubMapForAcknowledgementBatch(
          ShareAcknowledgementBatch batch
      ) {
          lock.writeLock().lock();
          try {
              // Find the floor batch record for the request batch. The request batch could be
              // for a subset of the batch i.e. cached batch of offset 10-14 and request batch
              // of 12-13. Hence, floor entry is fetched to find the sub-map.
              Map.Entry<Long, InFlightBatch> floorOffset = cachedState.floorEntry(batch.firstOffset());
              if (floorOffset == null) {
                  boolean hasStartOffsetMoved = checkForStartOffsetWithinBatch(batch.firstOffset(), batch.lastOffset());
                  if (hasStartOffsetMoved) {
                      // If the start offset has been moved and within the request batch then fetch
                      // the floor entry from start offset and acknowledge cached offsets. Consider
                      // the case where the start offset has moved from 0 to 10, with the cached batch
                      // of 0 - 5, 5 - 10, 10 - 12, 12 - 15. The request batch for acknowledge is 5 - 15,
                      // then post acquisition lock timeout the cache will have data from only from 10 to 15.
                      // Hence, we need to fetch the floor entry from start offset.
                      floorOffset = cachedState.floorEntry(startOffset);
                  } else {
                      log.debug("Batch record {} not found for share partition: {}-{}", batch, groupId,
                          topicIdPartition);
                      throw new InvalidRecordStateException(
                          "Batch record not found. The request batch offsets are not found in the cache.");
                  }
              }
  
              NavigableMap<Long, InFlightBatch> subMap = cachedState.subMap(floorOffset.getKey(), true, batch.lastOffset(), true);
              // Validate if the request batch has the first offset greater than the last offset of the last
              // fetched cached batch, then there will be no offsets in the request that can be acknowledged.
              if (subMap.lastEntry().getValue().lastOffset < batch.firstOffset()) {
                  log.debug("Request batch: {} has offsets which are not found for share partition: {}-{}", batch, groupId, topicIdPartition);
                  throw new InvalidRequestException("Batch record not found. The first offset in request is past acquired records.");
              }
  
              // Validate if the request batch has the last offset greater than the last offset of
              // the last fetched cached batch, then there will be offsets in the request than cannot
              // be found in the fetched batches.
              if (batch.lastOffset() > subMap.lastEntry().getValue().lastOffset) {
                  log.debug("Request batch: {} has offsets which are not found for share partition: {}-{}", batch, groupId, topicIdPartition);
                  throw new InvalidRequestException("Batch record not found. The last offset in request is past acquired records.");
              }
  
              return subMap;
          } finally {
              lock.writeLock().unlock();
          }
      }
  
      private Optional<Throwable> acknowledgeBatchRecords(
          String memberId,
          ShareAcknowledgementBatch batch,
          Map<Long, RecordState> recordStateMap,
          NavigableMap<Long, InFlightBatch> subMap,
          final List<InFlightState> updatedStates,
          List<PersisterStateBatch> stateBatches
      ) {
          Optional<Throwable> throwable;
          lock.writeLock().lock();
          try {
              // The acknowledgement batch either is exact fetch equivalent batch (mostly), subset
              // or spans over multiple fetched batches. The state can vary per offset itself from
              // the fetched batch in case of subset or client sent individual offsets state.
              for (Map.Entry<Long, InFlightBatch> entry : subMap.entrySet()) {
                  InFlightBatch inFlightBatch = entry.getValue();
  
                  // If startOffset has moved ahead of the in-flight batch, skip the batch.
                  if (inFlightBatch.lastOffset() < startOffset) {
                      log.trace("All offsets in the inflight batch {} are already archived: {}-{}",
                          inFlightBatch, groupId, topicIdPartition);
                      continue;
                  }
  
                  // Validate if the requested member id is the owner of the batch.
                  if (inFlightBatch.offsetState() == null) {
                      throwable = validateAcknowledgementBatchMemberId(memberId, inFlightBatch);
                      if (throwable.isPresent()) {
                          return throwable;
                      }
                  }
  
                  // Determine if the in-flight batch is a full match from the request batch.
                  boolean fullMatch = checkForFullMatch(inFlightBatch, batch.firstOffset(), batch.lastOffset());
                  boolean isPerOffsetClientAck = batch.acknowledgeTypes().size() > 1;
                  boolean hasStartOffsetMoved = checkForStartOffsetWithinBatch(inFlightBatch.firstOffset(), inFlightBatch.lastOffset());
  
                  // Maintain state per offset if the inflight batch is not a full match or the
                  // offset state is managed or client sent individual offsets state or
                  // the start offset is within this in-flight batch.
                  if (!fullMatch || inFlightBatch.offsetState() != null || isPerOffsetClientAck || hasStartOffsetMoved) {
                      log.debug("Subset or offset tracked batch record found for acknowledgement,"
                              + " batch: {}, request offsets - first: {}, last: {}, client per offset"
                              + "state {} for the share partition: {}-{}", inFlightBatch, batch.firstOffset(),
                          batch.lastOffset(), isPerOffsetClientAck, groupId, topicIdPartition);
                      if (inFlightBatch.offsetState() == null) {
                          // Though the request is a subset of in-flight batch but the offset
                          // tracking has not been initialized yet which means that we could only
                          // acknowledge subset of offsets from the in-flight batch but only if the
                          // complete batch is acquired yet. Hence, do a pre-check to avoid exploding
                          // the in-flight offset tracking unnecessarily.
                          if (inFlightBatch.batchState() != RecordState.ACQUIRED) {
                              log.debug("The batch is not in the acquired state: {} for share partition: {}-{}",
                                  inFlightBatch, groupId, topicIdPartition);
                              return Optional.of(new InvalidRecordStateException("The batch cannot be acknowledged. The subset batch is not in the acquired state."));
                          }
                          // The request batch is a subset and requires per offset state hence initialize
                          // the offsets state in the in-flight batch.
                          inFlightBatch.maybeInitializeOffsetStateUpdate();
                      }
  
                      throwable = acknowledgePerOffsetBatchRecords(memberId, batch, inFlightBatch,
                          recordStateMap, updatedStates, stateBatches);
                  } else {
                      // The in-flight batch is a full match hence change the state of the complete batch.
                      throwable = acknowledgeCompleteBatch(batch, inFlightBatch,
                          recordStateMap.get(batch.firstOffset()), updatedStates, stateBatches);
                  }
  
                  if (throwable.isPresent()) {
                      return throwable;
                  }
              }
          } finally {
              lock.writeLock().unlock();
          }
          return Optional.empty();
      }
  
      private Optional<Throwable> validateAcknowledgementBatchMemberId(
          String memberId,
          InFlightBatch inFlightBatch
      ) {
          // EMPTY_MEMBER_ID is used to indicate that the batch is not in acquired state.
          if (inFlightBatch.batchMemberId().equals(EMPTY_MEMBER_ID)) {
              log.debug("The batch is not in the acquired state: {} for share partition: {}-{}. Empty member id for batch.",
                  inFlightBatch, groupId, topicIdPartition);
              return Optional.of(new InvalidRecordStateException("The batch cannot be acknowledged. The batch is not in the acquired state."));
          }
  
          if (!inFlightBatch.batchMemberId().equals(memberId)) {
              log.debug("Member {} is not the owner of batch record {} for share partition: {}-{}",
                  memberId, inFlightBatch, groupId, topicIdPartition);
              return Optional.of(new InvalidRecordStateException("Member is not the owner of batch record"));
          }
          return Optional.empty();
      }
  
      private Optional<Throwable> acknowledgePerOffsetBatchRecords(
          String memberId,
          ShareAcknowledgementBatch batch,
          InFlightBatch inFlightBatch,
          Map<Long, RecordState> recordStateMap,
          List<InFlightState> updatedStates,
          List<PersisterStateBatch> stateBatches
      ) {
          lock.writeLock().lock();
          try {
              // Fetch the first record state from the map to be used as default record state in case the
              // offset record state is not provided by client.
              RecordState recordStateDefault = recordStateMap.get(batch.firstOffset());
              for (Map.Entry<Long, InFlightState> offsetState : inFlightBatch.offsetState.entrySet()) {
  
                  // 1. For the first batch which might have offsets prior to the request base
                  // offset i.e. cached batch of 10-14 offsets and request batch of 12-13.
                  // 2. Skip the offsets which are below the start offset of the share partition
                  if (offsetState.getKey() < batch.firstOffset() || offsetState.getKey() < startOffset) {
                      continue;
                  }
  
                  if (offsetState.getKey() > batch.lastOffset()) {
                      // No further offsets to process.
                      break;
                  }
  
                  if (offsetState.getValue().state != RecordState.ACQUIRED) {
                      log.debug("The offset is not acquired, offset: {} batch: {} for the share"
                              + " partition: {}-{}", offsetState.getKey(), inFlightBatch, groupId,
                          topicIdPartition);
                      return Optional.of(new InvalidRecordStateException(
                          "The batch cannot be acknowledged. The offset is not acquired."));
                  }
  
                  // Check if member id is the owner of the offset.
                  if (!offsetState.getValue().memberId.equals(memberId)) {
                      log.debug("Member {} is not the owner of offset: {} in batch: {} for the share"
                              + " partition: {}-{}", memberId, offsetState.getKey(), inFlightBatch,
                          groupId, topicIdPartition);
                      return Optional.of(
                          new InvalidRecordStateException("Member is not the owner of offset"));
                  }
  
                  // Determine the record state for the offset. If the per offset record state is not provided
                  // by the client, then use the batch record state.
                  RecordState recordState =
                      recordStateMap.size() > 1 ? recordStateMap.get(offsetState.getKey()) :
                          recordStateDefault;
                  InFlightState updateResult = offsetState.getValue().startStateTransition(
                      recordState,
                      DeliveryCountOps.NO_OP,
                      this.maxDeliveryCount,
                      EMPTY_MEMBER_ID
                  );
                  if (updateResult == null) {
                      log.debug("Unable to acknowledge records for the offset: {} in batch: {}"
                              + " for the share partition: {}-{}", offsetState.getKey(),
                          inFlightBatch, groupId, topicIdPartition);
                      return Optional.of(new InvalidRecordStateException(
                          "Unable to acknowledge records for the batch"));
                  }
                  // Successfully updated the state of the offset.
                  updatedStates.add(updateResult);
                  stateBatches.add(new PersisterStateBatch(offsetState.getKey(), offsetState.getKey(),
                      updateResult.state.id, (short) updateResult.deliveryCount));
                  // If the maxDeliveryCount limit has been exceeded, the record will be transitioned to ARCHIVED state.
                  // This should not change the next fetch offset because the record is not available for acquisition
                  if (recordState == RecordState.AVAILABLE
                      && updateResult.state != RecordState.ARCHIVED) {
                      findNextFetchOffset.set(true);
                  }
              }
          } finally {
              lock.writeLock().unlock();
          }
          return Optional.empty();
      }
  
      private Optional<Throwable> acknowledgeCompleteBatch(
          ShareAcknowledgementBatch batch,
          InFlightBatch inFlightBatch,
          RecordState recordState,
          List<InFlightState> updatedStates,
          List<PersisterStateBatch> stateBatches
      ) {
          lock.writeLock().lock();
          try {
              // The in-flight batch is a full match hence change the state of the complete.
              log.trace("Acknowledging complete batch record {} for the share partition: {}-{}",
                  batch, groupId, topicIdPartition);
              if (inFlightBatch.batchState() != RecordState.ACQUIRED) {
                  log.debug("The batch is not in the acquired state: {} for share partition: {}-{}",
                      inFlightBatch, groupId, topicIdPartition);
                  return Optional.of(new InvalidRecordStateException(
                      "The batch cannot be acknowledged. The batch is not in the acquired state."));
              }
  
              // Change the state of complete batch since the same state exists for the entire inFlight batch.
              // The member id is reset to EMPTY_MEMBER_ID irrespective of the acknowledge type as the batch is
              // either released or moved to a state where member id existence is not important. The member id
              // is only important when the batch is acquired.
              InFlightState updateResult = inFlightBatch.startBatchStateTransition(
                  recordState,
                  DeliveryCountOps.NO_OP,
                  this.maxDeliveryCount,
                  EMPTY_MEMBER_ID
              );
              if (updateResult == null) {
                  log.debug("Unable to acknowledge records for the batch: {} with state: {}"
                          + " for the share partition: {}-{}", inFlightBatch, recordState, groupId,
                      topicIdPartition);
                  return Optional.of(
                      new InvalidRecordStateException("Unable to acknowledge records for the batch"));
              }
  
              // Successfully updated the state of the batch.
              updatedStates.add(updateResult);
              stateBatches.add(
                  new PersisterStateBatch(inFlightBatch.firstOffset, inFlightBatch.lastOffset,
                      updateResult.state.id, (short) updateResult.deliveryCount));
  
              // If the maxDeliveryCount limit has been exceeded, the record will be transitioned to ARCHIVED state.
              // This should not change the nextFetchOffset because the record is not available for acquisition
              if (recordState == RecordState.AVAILABLE
                  && updateResult.state != RecordState.ARCHIVED) {
                  findNextFetchOffset.set(true);
              }
          } finally {
              lock.writeLock().unlock();
          }
          return Optional.empty();
      }
  
      protected void updateFetchOffsetMetadata(long nextFetchOffset, LogOffsetMetadata logOffsetMetadata) {
          lock.writeLock().lock();
          try {
              fetchOffsetMetadata.updateOffsetMetadata(nextFetchOffset, logOffsetMetadata);
          } finally {
              lock.writeLock().unlock();
          }
      }
  
      protected Optional<LogOffsetMetadata> fetchOffsetMetadata(long nextFetchOffset) {
          lock.readLock().lock();
          try {
              if (fetchOffsetMetadata.offsetMetadata() == null || fetchOffsetMetadata.offset() != nextFetchOffset)
                  return Optional.empty();
              return Optional.of(fetchOffsetMetadata.offsetMetadata());
          } finally {
              lock.readLock().unlock();
          }
      }
  
      // Visible for testing
      SharePartitionState partitionState() {
          lock.readLock().lock();
          try {
              return partitionState;
          } finally {
              lock.readLock().unlock();
          }
      }
  
      // Visible for testing
      void rollbackOrProcessStateUpdates(
          CompletableFuture<Void> future,
          Throwable throwable,
          List<InFlightState> updatedStates,
          List<PersisterStateBatch> stateBatches
      ) {
          lock.writeLock().lock();
          try {
              if (throwable != null) {
                  // Log in DEBUG to avoid flooding of logs for a faulty client.
                  log.debug("Request failed for updating state, rollback any changed state"
                      + " for the share partition: {}-{}", groupId, topicIdPartition);
                  updatedStates.forEach(state -> state.completeStateTransition(false));
                  future.completeExceptionally(throwable);
                  return;
              }
  
              if (stateBatches.isEmpty() && updatedStates.isEmpty()) {
                  future.complete(null);
                  return;
              }
          } finally {
              lock.writeLock().unlock();
          }
  
          writeShareGroupState(stateBatches).whenComplete((result, exception) -> {
              // There can be a pending delayed share fetch requests for the share partition which are waiting
              // on the startOffset to move ahead, hence track if the state is updated in the cache. If
              // yes, then notify the delayed share fetch purgatory to complete the pending requests.
              boolean cacheStateUpdated = false;
              lock.writeLock().lock();
              try {
                  if (exception != null) {
                      log.error("Failed to write state to persister for the share partition: {}-{}",
                          groupId, topicIdPartition, exception);
                      updatedStates.forEach(state -> state.completeStateTransition(false));
                      future.completeExceptionally(exception);
                      return;
                  }
  
                  log.trace("State change request successful for share partition: {}-{}",
                      groupId, topicIdPartition);
                  updatedStates.forEach(state -> {
                      state.completeStateTransition(true);
                      // Cancel the acquisition lock timeout task for the state since it is acknowledged/released successfully.
                      state.cancelAndClearAcquisitionLockTimeoutTask();
                  });
                  // Update the cached state and start and end offsets after acknowledging/releasing the acquired records.
                  cacheStateUpdated  = maybeUpdateCachedStateAndOffsets();
                  future.complete(null);
              } finally {
                  lock.writeLock().unlock();
                  // Maybe complete the delayed share fetch request if the state has been changed in cache
                  // which might have moved start offset ahead. Hence, the pending delayed share fetch
                  // request can be completed. The call should be made outside the lock to avoid deadlock.
                  maybeCompleteDelayedShareFetchRequest(cacheStateUpdated);
              }
          });
      }
  
      private boolean maybeUpdateCachedStateAndOffsets() {
          lock.writeLock().lock();
          try {
              if (!canMoveStartOffset()) {
                  return false;
              }
  
              // This will help to find the next position for the startOffset.
              // The new position of startOffset will be lastOffsetAcknowledged + 1
              long lastOffsetAcknowledged = findLastOffsetAcknowledged();
              // If lastOffsetAcknowledged is -1, this means we cannot move startOffset ahead
              if (lastOffsetAcknowledged == -1) {
                  return false;
              }
  
              // This is true if all records in the cachedState have been acknowledged (either Accept or Reject).
              // The resulting action should be to empty the cachedState altogether
              long lastCachedOffset = cachedState.lastEntry().getValue().lastOffset();
              if (lastOffsetAcknowledged == lastCachedOffset) {
                  startOffset = lastCachedOffset + 1; // The next offset that will be fetched and acquired in the share partition
                  endOffset = lastCachedOffset + 1;
                  cachedState.clear();
                  // Nothing further to do.
                  return true;
              }
  
              /*
               The cachedState contains some records that are yet to be acknowledged, and thus should
               not be removed. Only a subMap will be removed from the cachedState. The logic to remove
               batches from cachedState is as follows:
               a) Only full batches can be removed from the cachedState, For example if there is batch (0-99)
               and 0-49 records are acknowledged (ACCEPT or REJECT), the first 50 records will not be removed
               from the cachedState. Instead, the startOffset will be moved to 50, but the batch will only
               be removed once all the messages (0-99) are acknowledged (ACCEPT or REJECT).
              */
  
              // Since only a subMap will be removed, we need to find the first and last keys of that subMap
              long firstKeyToRemove = cachedState.firstKey();
              long lastKeyToRemove;
              NavigableMap.Entry<Long, InFlightBatch> entry = cachedState.floorEntry(lastOffsetAcknowledged);
              // If the lastOffsetAcknowledged is equal to the last offset of entry, then the entire batch can potentially be removed.
              if (lastOffsetAcknowledged == entry.getValue().lastOffset()) {
                  startOffset = cachedState.higherKey(lastOffsetAcknowledged);
                  if (isInitialReadGapOffsetWindowActive()) {
                      // This case will arise if we have a situation where there is an acquirable gap after the lastOffsetAcknowledged.
                      // Ex, the cachedState has following state batches -> {(0, 10), (11, 20), (31,40)} and all these batches are acked.
                      // There is a gap from 21 to 30. Let the initialReadGapOffset.gapStartOffset be 21. In this case,
                      // lastOffsetAcknowledged will be 20, but we cannot simply move the start offset to the first offset
                      // of next cachedState batch (next cachedState batch is 31 to 40). There is an acquirable gap in between (21 to 30)
                      // and The startOffset should be at 21. Hence, we set startOffset to the minimum of initialReadGapOffset.gapStartOffset
                      // and higher key of lastOffsetAcknowledged
                      startOffset = Math.min(initialReadGapOffset.gapStartOffset(), startOffset);
                  }
                  lastKeyToRemove = entry.getKey();
              } else {
                  // The code will reach this point only if lastOffsetAcknowledged is in the middle of some stateBatch. In this case
                  // we can simply move the startOffset to the next offset of lastOffsetAcknowledged and should consider any read gap offsets.
                  startOffset = lastOffsetAcknowledged + 1;
                  if (entry.getKey().equals(cachedState.firstKey())) {
                      // If the first batch in cachedState has some records yet to be acknowledged,
                      // then nothing should be removed from cachedState
                      lastKeyToRemove = -1;
                  } else {
                      lastKeyToRemove = cachedState.lowerKey(entry.getKey());
                  }
              }
  
              if (lastKeyToRemove != -1) {
                  cachedState.subMap(firstKeyToRemove, true, lastKeyToRemove, true).clear();
              }
              return true;
          } finally {
              lock.writeLock().unlock();
          }
      }
  
      private boolean canMoveStartOffset() {
          // The Share Partition Start Offset may be moved after acknowledge request is complete.
          // The following conditions need to be met to move the startOffset:
          // 1. When the cachedState is not empty.
          // 2. When the acknowledgement type for the records is either ACCEPT or REJECT.
          // 3. When all the previous records have been acknowledged (ACCEPT or REJECT).
          if (cachedState.isEmpty()) {
              return false;
          }
  
          NavigableMap.Entry<Long, InFlightBatch> entry = cachedState.floorEntry(startOffset);
          if (entry == null) {
              // The start offset is not found in the cached state when there is a gap starting at the start offset.
              // For example, if the start offset is 10 and the cached state has batches -> { (21, 30), (31, 40) }.
              // This case arises only when the share partition is initialized and the read state response results in
              // state batches containing gaps. This situation is possible in the case where in the previous instance
              // of this share partition, the gap offsets were fetched but not acknowledged, and the next batch of offsets
              // were fetched as well as acknowledged. In the above example, possibly in the previous instance of the share
              // partition, the batch 10-20 was fetched but not acknowledged and the batch 21-30 was fetched and acknowledged.
              // Thus, the persister has no clue about what happened with the batch 10-20. During the re-initialization of
              // the share partition, the start offset is set to 10 and the cached state has the batch 21-30, resulting in a gap.
              log.debug("The start offset: {} is not found in the cached state for share partition: {}-{} " +
                  "as there is an acquirable gap at the beginning. Cannot move the start offset.", startOffset, groupId, topicIdPartition);
              return false;
          }
          RecordState startOffsetState = entry.getValue().offsetState == null ?
              entry.getValue().batchState() :
              entry.getValue().offsetState().get(startOffset).state();
          return isRecordStateAcknowledged(startOffsetState);
      }
  
      private boolean isInitialReadGapOffsetWindowActive() {
          return initialReadGapOffset != null && initialReadGapOffset.endOffset() == endOffset;
      }
  
      /**
       * The record state is considered acknowledged if it is either acknowledged or archived.
       * These are terminal states for the record.
       *
       * @param recordState The record state to check.
       *
       * @return True if the record state is acknowledged or archived, false otherwise.
       */
      private boolean isRecordStateAcknowledged(RecordState recordState) {
          return recordState == RecordState.ACKNOWLEDGED || recordState == RecordState.ARCHIVED;
      }
  
      // Visible for testing
      long findLastOffsetAcknowledged() {
          long lastOffsetAcknowledged = -1;
          lock.readLock().lock();
          try {
              for (NavigableMap.Entry<Long, InFlightBatch> entry : cachedState.entrySet()) {
                  InFlightBatch inFlightBatch = entry.getValue();
  
                  if (isInitialReadGapOffsetWindowActive() && inFlightBatch.lastOffset() >= initialReadGapOffset.gapStartOffset()) {
                      return lastOffsetAcknowledged;
                  }
  
                  if (inFlightBatch.offsetState() == null) {
                      if (!isRecordStateAcknowledged(inFlightBatch.batchState())) {
                          return lastOffsetAcknowledged;
                      }
                      lastOffsetAcknowledged = inFlightBatch.lastOffset();
                  } else {
                      for (Map.Entry<Long, InFlightState> offsetState : inFlightBatch.offsetState.entrySet()) {
                          if (!isRecordStateAcknowledged(offsetState.getValue().state())) {
                              return lastOffsetAcknowledged;
                          }
                          lastOffsetAcknowledged = offsetState.getKey();
                      }
                  }
              }
          } finally {
              lock.readLock().unlock();
          }
          return lastOffsetAcknowledged;
      }
  
      /**
       * Find the last offset from the batch which contains the request offset. If found, return the last offset
       * of the batch, otherwise return the request offset.
       *
       * @param batches The batches to search for the request offset.
       * @param offset The request offset to find.
       * @return The last offset of the batch which contains the request offset, otherwise the request offset.
       */
      private long lastOffsetFromBatchWithRequestOffset(
          Iterable<? extends RecordBatch> batches,
          long offset
      ) {
          // Fetch the last batch which might contains the request offset. Avoid calling lastOffset()
          // on the batches as it loads the header in memory.
          RecordBatch previousBatch = null;
          for (RecordBatch batch : batches) {
              if (offset >= batch.baseOffset()) {
                  previousBatch =  batch;
                  continue;
              }
              break;
          }
          if (previousBatch != null && offset <= previousBatch.lastOffset())
              return previousBatch.lastOffset();
          return offset;
      }
  
      // Visible for testing
      CompletableFuture<Void> writeShareGroupState(List<PersisterStateBatch> stateBatches) {
          CompletableFuture<Void> future = new CompletableFuture<>();
          persister.writeState(new WriteShareGroupStateParameters.Builder()
              .setGroupTopicPartitionData(new GroupTopicPartitionData.Builder<PartitionStateBatchData>()
                  .setGroupId(this.groupId)
                  .setTopicsData(List.of(new TopicData<>(topicIdPartition.topicId(),
                      List.of(PartitionFactory.newPartitionStateBatchData(
                          topicIdPartition.partition(), stateEpoch, startOffset, leaderEpoch, stateBatches))))
                  ).build()).build())
              .whenComplete((result, exception) -> {
                  if (exception != null) {
                      log.error("Failed to write the share group state for share partition: {}-{}", groupId, topicIdPartition, exception);
                      future.completeExceptionally(new IllegalStateException(String.format("Failed to write the share group state for share partition %s-%s",
                          groupId, topicIdPartition), exception));
                      return;
                  }
  
                  if (result == null || result.topicsData() == null || result.topicsData().size() != 1) {
                      log.error("Failed to write the share group state for share partition: {}-{}. Invalid state found: {}",
                          groupId, topicIdPartition, result);
                      future.completeExceptionally(new IllegalStateException(String.format("Failed to write the share group state for share partition %s-%s",
                          groupId, topicIdPartition)));
                      return;
                  }
  
                  TopicData<PartitionErrorData> state = result.topicsData().get(0);
                  if (state.topicId() != topicIdPartition.topicId() || state.partitions().size() != 1
                      || state.partitions().get(0).partition() != topicIdPartition.partition()) {
                      log.error("Failed to write the share group state for share partition: {}-{}. Invalid topic partition response: {}",
                          groupId, topicIdPartition, result);
                      future.completeExceptionally(new IllegalStateException(String.format("Failed to write the share group state for share partition %s-%s",
                          groupId, topicIdPartition)));
                      return;
                  }
  
                  PartitionErrorData partitionData = state.partitions().get(0);
                  if (partitionData.errorCode() != Errors.NONE.code()) {
                      KafkaException ex = fetchPersisterError(partitionData.errorCode(), partitionData.errorMessage());
                      log.error("Failed to write the share group state for share partition: {}-{} due to exception",
                          groupId, topicIdPartition, ex);
                      future.completeExceptionally(ex);
                      return;
                  }
                  future.complete(null);
              });
          return future;
      }
  
      private KafkaException fetchPersisterError(short errorCode, String errorMessage) {
          Errors error = Errors.forCode(errorCode);
          return switch (error) {
              case NOT_COORDINATOR, COORDINATOR_NOT_AVAILABLE, COORDINATOR_LOAD_IN_PROGRESS ->
                  new CoordinatorNotAvailableException(errorMessage);
              case GROUP_ID_NOT_FOUND ->
                  new GroupIdNotFoundException(errorMessage);
              case UNKNOWN_TOPIC_OR_PARTITION ->
                  new UnknownTopicOrPartitionException(errorMessage);
              case FENCED_STATE_EPOCH ->
                  new FencedStateEpochException(errorMessage);
              case FENCED_LEADER_EPOCH ->
                  new NotLeaderOrFollowerException(errorMessage);
              default ->
                  new UnknownServerException(errorMessage);
          };
      }
  
      // Visible for testing
      AcquisitionLockTimerTask scheduleAcquisitionLockTimeout(String memberId, long firstOffset, long lastOffset) {
          // The recordLockDuration value would depend on whether the dynamic config SHARE_RECORD_LOCK_DURATION_MS in
          // GroupConfig.java is set or not. If dynamic config is set, then that is used, otherwise the value of
          // SHARE_GROUP_RECORD_LOCK_DURATION_MS_CONFIG defined in ShareGroupConfig is used
          int recordLockDurationMs = recordLockDurationMsOrDefault(groupConfigManager, groupId, defaultRecordLockDurationMs);
          return scheduleAcquisitionLockTimeout(memberId, firstOffset, lastOffset, recordLockDurationMs);
      }
  
      /**
       * Apply acquisition lock to acquired records.
       *
       * @param memberId The member id of the client that is putting the acquisition lock.
       * @param firstOffset The first offset of the acquired records.
       * @param lastOffset The last offset of the acquired records.
       * @param delayMs The delay in milliseconds after which the acquisition lock will be released.
       */
      private AcquisitionLockTimerTask scheduleAcquisitionLockTimeout(
          String memberId,
          long firstOffset,
          long lastOffset,
          long delayMs
      ) {
          AcquisitionLockTimerTask acquisitionLockTimerTask = acquisitionLockTimerTask(memberId, firstOffset, lastOffset, delayMs);
          timer.add(acquisitionLockTimerTask);
          return acquisitionLockTimerTask;
      }
  
      private AcquisitionLockTimerTask acquisitionLockTimerTask(
          String memberId,
          long firstOffset,
          long lastOffset,
          long delayMs
      ) {
          return new AcquisitionLockTimerTask(delayMs, memberId, firstOffset, lastOffset);
      }
  
      private void releaseAcquisitionLockOnTimeout(String memberId, long firstOffset, long lastOffset) {
          List<PersisterStateBatch> stateBatches;
          lock.writeLock().lock();
          try {
              Map.Entry<Long, InFlightBatch> floorOffset = cachedState.floorEntry(firstOffset);
              if (floorOffset == null) {
                  log.error("Base offset {} not found for share partition: {}-{}", firstOffset, groupId, topicIdPartition);
                  return;
              }
              stateBatches = new ArrayList<>();
              NavigableMap<Long, InFlightBatch> subMap = cachedState.subMap(floorOffset.getKey(), true, lastOffset, true);
              for (Map.Entry<Long, InFlightBatch> entry : subMap.entrySet()) {
                  InFlightBatch inFlightBatch = entry.getValue();
  
                  if (inFlightBatch.offsetState() == null
                          && inFlightBatch.batchState() == RecordState.ACQUIRED
                          && checkForStartOffsetWithinBatch(inFlightBatch.firstOffset(), inFlightBatch.lastOffset())) {
  
                      // For the case when batch.firstOffset < start offset <= batch.lastOffset, we will be having some
                      // acquired records that need to move to archived state despite their delivery count.
                      inFlightBatch.maybeInitializeOffsetStateUpdate();
                  }
  
                  // Case when the state of complete batch is valid
                  if (inFlightBatch.offsetState() == null) {
                      releaseAcquisitionLockOnTimeoutForCompleteBatch(inFlightBatch, stateBatches, memberId);
                  } else { // Case when batch has a valid offset state map.
                      releaseAcquisitionLockOnTimeoutForPerOffsetBatch(inFlightBatch, stateBatches, memberId, firstOffset, lastOffset);
                  }
              }
  
              if (!stateBatches.isEmpty()) {
                  writeShareGroupState(stateBatches).whenComplete((result, exception) -> {
                      if (exception != null) {
                          log.error("Failed to write the share group state on acquisition lock timeout for share partition: {}-{} memberId: {}",
                              groupId, topicIdPartition, memberId, exception);
                      }
                      // Even if write share group state RPC call fails, we will still go ahead with the state transition.
                      // Update the cached state and start and end offsets after releasing the acquisition lock on timeout.
                      maybeUpdateCachedStateAndOffsets();
                  });
              }
          } finally {
              lock.writeLock().unlock();
          }
  
          // If we have an acquisition lock timeout for a share-partition, then we should check if
          // there is a pending share fetch request for the share-partition and complete it.
          // Skip null check for stateBatches, it should always be initialized if reached here.
          maybeCompleteDelayedShareFetchRequest(!stateBatches.isEmpty());
      }
  
      private void releaseAcquisitionLockOnTimeoutForCompleteBatch(InFlightBatch inFlightBatch,
                                                                   List<PersisterStateBatch> stateBatches,
                                                                   String memberId) {
          if (inFlightBatch.batchState() == RecordState.ACQUIRED) {
              InFlightState updateResult = inFlightBatch.tryUpdateBatchState(
                      inFlightBatch.lastOffset() < startOffset ? RecordState.ARCHIVED : RecordState.AVAILABLE,
                      DeliveryCountOps.NO_OP,
                      maxDeliveryCount,
                      EMPTY_MEMBER_ID);
              if (updateResult == null) {
                  log.error("Unable to release acquisition lock on timeout for the batch: {}"
                          + " for the share partition: {}-{} memberId: {}", inFlightBatch, groupId, topicIdPartition, memberId);
                  return;
              }
              stateBatches.add(new PersisterStateBatch(inFlightBatch.firstOffset(), inFlightBatch.lastOffset(),
                      updateResult.state.id, (short) updateResult.deliveryCount));
  
              // Cancel the acquisition lock timeout task for the batch since it is completed now.
              updateResult.cancelAndClearAcquisitionLockTimeoutTask();
              if (updateResult.state != RecordState.ARCHIVED) {
                  findNextFetchOffset.set(true);
              }
              return;
          }
          log.debug("The batch is not in acquired state while release of acquisition lock on timeout, skipping, batch: {}"
                  + " for the share partition: {}-{} memberId: {}", inFlightBatch, groupId, topicIdPartition, memberId);
      }
  
      private void releaseAcquisitionLockOnTimeoutForPerOffsetBatch(InFlightBatch inFlightBatch,
                                                                    List<PersisterStateBatch> stateBatches,
                                                                    String memberId,
                                                                    long firstOffset,
                                                                    long lastOffset) {
          for (Map.Entry<Long, InFlightState> offsetState : inFlightBatch.offsetState().entrySet()) {
  
              // For the first batch which might have offsets prior to the request base
              // offset i.e. cached batch of 10-14 offsets and request batch of 12-13.
              if (offsetState.getKey() < firstOffset) {
                  continue;
              }
              if (offsetState.getKey() > lastOffset) {
                  // No further offsets to process.
                  break;
              }
              if (offsetState.getValue().state != RecordState.ACQUIRED) {
                  log.debug("The offset is not in acquired state while release of acquisition lock on timeout, skipping, offset: {} batch: {}"
                                  + " for the share partition: {}-{} memberId: {}", offsetState.getKey(), inFlightBatch,
                          groupId, topicIdPartition, memberId);
                  continue;
              }
              InFlightState updateResult = offsetState.getValue().tryUpdateState(
                      offsetState.getKey() < startOffset ? RecordState.ARCHIVED : RecordState.AVAILABLE,
                      DeliveryCountOps.NO_OP,
                      maxDeliveryCount,
                      EMPTY_MEMBER_ID);
              if (updateResult == null) {
                  log.error("Unable to release acquisition lock on timeout for the offset: {} in batch: {}"
                                  + " for the share partition: {}-{} memberId: {}", offsetState.getKey(), inFlightBatch,
                          groupId, topicIdPartition, memberId);
                  continue;
              }
              stateBatches.add(new PersisterStateBatch(offsetState.getKey(), offsetState.getKey(),
                      updateResult.state.id, (short) updateResult.deliveryCount));
  
              // Cancel the acquisition lock timeout task for the offset since it is completed now.
              updateResult.cancelAndClearAcquisitionLockTimeoutTask();
              if (updateResult.state != RecordState.ARCHIVED) {
                  findNextFetchOffset.set(true);
              }
          }
      }
  
      private void maybeCompleteDelayedShareFetchRequest(boolean shouldComplete) {
          if (shouldComplete) {
              replicaManager.completeDelayedShareFetchRequest(delayedShareFetchKey);
          }
      }
  
      private long startOffsetDuringInitialization(long partitionDataStartOffset) {
          // Set the state epoch and end offset from the persisted state.
          if (partitionDataStartOffset != PartitionFactory.UNINITIALIZED_START_OFFSET) {
              return partitionDataStartOffset;
          }
          ShareGroupAutoOffsetResetStrategy offsetResetStrategy;
          if (groupConfigManager.groupConfig(groupId).isPresent()) {
              offsetResetStrategy = groupConfigManager.groupConfig(groupId).get().shareAutoOffsetReset();
          } else {
              offsetResetStrategy = GroupConfig.defaultShareAutoOffsetReset();
          }
  
          if (offsetResetStrategy.type() == ShareGroupAutoOffsetResetStrategy.StrategyType.LATEST) {
              return offsetForLatestTimestamp(topicIdPartition, replicaManager, leaderEpoch);
          } else if (offsetResetStrategy.type() == ShareGroupAutoOffsetResetStrategy.StrategyType.EARLIEST) {
              return offsetForEarliestTimestamp(topicIdPartition, replicaManager, leaderEpoch);
          } else {
              // offsetResetStrategy type is BY_DURATION
              return offsetForTimestamp(topicIdPartition, replicaManager, offsetResetStrategy.timestamp(), leaderEpoch);
          }
      }
  
      private ShareAcquiredRecords maybeFilterAbortedTransactionalAcquiredRecords(
          FetchPartitionData fetchPartitionData,
          FetchIsolation isolationLevel,
          ShareAcquiredRecords shareAcquiredRecords
      ) {
          if (isolationLevel != FetchIsolation.TXN_COMMITTED || fetchPartitionData.abortedTransactions.isEmpty() || fetchPartitionData.abortedTransactions.get().isEmpty())
              return shareAcquiredRecords;
  
          // When FetchIsolation.TXN_COMMITTED is used as isolation level by the share group, we need to filter any
          // transactions that were aborted/did not commit due to timeout.
          List<AcquiredRecords> result = filterAbortedTransactionalAcquiredRecords(fetchPartitionData.records.batches(),
              shareAcquiredRecords.acquiredRecords(), fetchPartitionData.abortedTransactions.get());
          int acquiredCount = 0;
          for (AcquiredRecords records : result) {
              acquiredCount += (int) (records.lastOffset() - records.firstOffset() + 1);
          }
          return new ShareAcquiredRecords(result, acquiredCount);
      }
  
      private List<AcquiredRecords> filterAbortedTransactionalAcquiredRecords(
          Iterable<? extends RecordBatch> batches,
          List<AcquiredRecords> acquiredRecords,
          List<FetchResponseData.AbortedTransaction> abortedTransactions
      ) {
          // The record batches that need to be archived in cachedState because they were a part of aborted transactions.
          List<RecordBatch> recordsToArchive = fetchAbortedTransactionRecordBatches(batches, abortedTransactions);
          for (RecordBatch recordBatch : recordsToArchive) {
              // Archive the offsets/batches in the cached state.
              NavigableMap<Long, InFlightBatch> subMap = fetchSubMap(recordBatch);
              archiveRecords(recordBatch.baseOffset(), recordBatch.lastOffset() + 1, subMap, RecordState.ACQUIRED);
          }
          return filterRecordBatchesFromAcquiredRecords(acquiredRecords, recordsToArchive);
      }
  
      /**
       * This function filters out the offsets present in the acquired records list that are also a part of batches that need to be archived.
       * It follows an iterative refinement of acquired records to eliminate batches to be archived.
       * @param acquiredRecordsList The list containing acquired records. This list is sorted by the firstOffset of the acquired batch.
       * @param batchesToArchive The list containing record batches to archive. This list is sorted by the baseOffset of the record batch.
       * @return The list containing filtered acquired records offsets.
       */
      List<AcquiredRecords> filterRecordBatchesFromAcquiredRecords(
          List<AcquiredRecords> acquiredRecordsList,
          List<RecordBatch> batchesToArchive
      ) {
          Iterator<RecordBatch> batchesToArchiveIterator = batchesToArchive.iterator();
          if (!batchesToArchiveIterator.hasNext())
              return acquiredRecordsList;
          List<AcquiredRecords> result = new ArrayList<>();
          Iterator<AcquiredRecords> acquiredRecordsListIter = acquiredRecordsList.iterator();
          RecordBatch batchToArchive = batchesToArchiveIterator.next();
          AcquiredRecords unresolvedAcquiredRecords = null;
  
          while (unresolvedAcquiredRecords != null || acquiredRecordsListIter.hasNext()) {
              if (unresolvedAcquiredRecords == null)
                  unresolvedAcquiredRecords = acquiredRecordsListIter.next();
  
              long unresolvedFirstOffset = unresolvedAcquiredRecords.firstOffset();
              long unresolvedLastOffset = unresolvedAcquiredRecords.lastOffset();
              short unresolvedDeliveryCount = unresolvedAcquiredRecords.deliveryCount();
  
              if (batchToArchive == null) {
                  result.add(unresolvedAcquiredRecords);
                  unresolvedAcquiredRecords = null;
                  continue;
              }
  
              // Non-overlap check - unresolvedFirstOffset offsets lie before the batchToArchive offsets. No need to filter out the offsets in such a scenario.
              if (unresolvedLastOffset < batchToArchive.baseOffset()) {
                  // Offsets in unresolvedAcquiredRecords do not overlap with batchToArchive, hence it should not get filtered out.
                  result.add(unresolvedAcquiredRecords);
                  unresolvedAcquiredRecords = null;
              }
  
              // Overlap check - unresolvedFirstOffset offsets overlap with the batchToArchive offsets. We need to filter out the overlapping
              // offsets in such a scenario.
              if (unresolvedFirstOffset <= batchToArchive.lastOffset() &&
                  unresolvedLastOffset >= batchToArchive.baseOffset()) {
                  unresolvedAcquiredRecords = null;
                  // Split the unresolvedFirstOffset into parts - before and after the overlapping record batchToArchive.
                  if (unresolvedFirstOffset < batchToArchive.baseOffset()) {
                      // The offsets in unresolvedAcquiredRecords that are present before batchToArchive's baseOffset should not get filtered out.
                      result.add(new AcquiredRecords()
                          .setFirstOffset(unresolvedFirstOffset)
                          .setLastOffset(batchToArchive.baseOffset() - 1)
                          .setDeliveryCount(unresolvedDeliveryCount));
                  }
                  if (unresolvedLastOffset > batchToArchive.lastOffset()) {
                      // The offsets in unresolvedAcquiredRecords that are present after batchToArchive's lastOffset should not get filtered out
                      // and should be taken forward for further processing since they could potentially contain offsets that need to be archived.
                      unresolvedAcquiredRecords = new AcquiredRecords()
                          .setFirstOffset(batchToArchive.lastOffset() + 1)
                          .setLastOffset(unresolvedLastOffset)
                          .setDeliveryCount(unresolvedDeliveryCount);
                  }
              }
  
              // There is at least one offset in unresolvedFirstOffset which lies after the batchToArchive. Hence, we move forward
              // the batchToArchive to the next element in batchesToArchiveIterator.
              if (unresolvedLastOffset > batchToArchive.lastOffset()) {
                  if (batchesToArchiveIterator.hasNext())
                      batchToArchive = batchesToArchiveIterator.next();
                  else
                      batchToArchive = null;
              }
          }
          return result;
      }
  
      /**
       * This function fetches the sub map from cachedState where all the offset details present in the recordBatch can be referred to
       * OR it gives an exception if those offsets are not present in cachedState.
       * @param recordBatch The record batch for which we want to find the sub map.
       * @return the sub map containing all the offset details.
       */
      private NavigableMap<Long, InFlightBatch> fetchSubMap(RecordBatch recordBatch) {
          lock.readLock().lock();
          try {
              Map.Entry<Long, InFlightBatch> floorEntry = cachedState.floorEntry(recordBatch.baseOffset());
              if (floorEntry == null) {
                  log.debug("Fetched batch record {} not found for share partition: {}-{}", recordBatch, groupId,
                      topicIdPartition);
                  throw new IllegalStateException(
                      "Batch record not found. The request batch offsets are not found in the cache.");
              }
              return cachedState.subMap(floorEntry.getKey(), true, recordBatch.lastOffset(), true);
          } finally {
              lock.readLock().unlock();
          }
      }
  
      // Visible for testing.
      List<RecordBatch> fetchAbortedTransactionRecordBatches(
          Iterable<? extends RecordBatch> batches,
          List<FetchResponseData.AbortedTransaction> abortedTransactions
      ) {
          PriorityQueue<FetchResponseData.AbortedTransaction> orderedAbortedTransactions = orderedAbortedTransactions(abortedTransactions);
          Set<Long> abortedProducerIds = new HashSet<>();
          List<RecordBatch> recordsToArchive = new ArrayList<>();
  
          for (RecordBatch currentBatch : batches) {
              if (currentBatch.hasProducerId()) {
                  // remove from the aborted transactions queue, all aborted transactions which have begun before the
                  // current batch's last offset and add the associated producerIds to the aborted producer set.
                  while (!orderedAbortedTransactions.isEmpty() && orderedAbortedTransactions.peek().firstOffset() <= currentBatch.lastOffset()) {
                      FetchResponseData.AbortedTransaction abortedTransaction = orderedAbortedTransactions.poll();
                      abortedProducerIds.add(abortedTransaction.producerId());
                  }
                  long producerId = currentBatch.producerId();
                  if (containsAbortMarker(currentBatch)) {
                      abortedProducerIds.remove(producerId);
                  } else if (isBatchAborted(currentBatch, abortedProducerIds)) {
                      log.debug("Skipping aborted record batch for share partition: {}-{} with producerId {} and " +
                          "offsets {} to {}", groupId, topicIdPartition, producerId, currentBatch.baseOffset(), currentBatch.lastOffset());
                      recordsToArchive.add(currentBatch);
                  }
              }
          }
          return recordsToArchive;
      }
  
      private PriorityQueue<FetchResponseData.AbortedTransaction> orderedAbortedTransactions(List<FetchResponseData.AbortedTransaction> abortedTransactions) {
          PriorityQueue<FetchResponseData.AbortedTransaction> orderedAbortedTransactions = new PriorityQueue<>(
              abortedTransactions.size(), Comparator.comparingLong(FetchResponseData.AbortedTransaction::firstOffset)
          );
          orderedAbortedTransactions.addAll(abortedTransactions);
          return orderedAbortedTransactions;
      }
  
      private boolean isBatchAborted(RecordBatch batch, Set<Long> abortedProducerIds) {
          return batch.isTransactional() && abortedProducerIds.contains(batch.producerId());
      }
  
      // Visible for testing.
      boolean containsAbortMarker(RecordBatch batch) {
          if (!batch.isControlBatch())
              return false;
  
          Iterator<Record> batchIterator = batch.iterator();
          if (!batchIterator.hasNext())
              return false;
  
          Record firstRecord = batchIterator.next();
          return ControlRecordType.ABORT == ControlRecordType.parse(firstRecord.key());
      }
  
      // Visible for testing. Should only be used for testing purposes.
      NavigableMap<Long, InFlightBatch> cachedState() {
          return new ConcurrentSkipListMap<>(cachedState);
      }
  
      // Visible for testing.
      boolean findNextFetchOffset() {
          return findNextFetchOffset.get();
      }
  
      // Visible for testing. Should only be used for testing purposes.
      void findNextFetchOffset(boolean findNextOffset) {
          findNextFetchOffset.getAndSet(findNextOffset);
      }
  
      // Visible for testing
      long startOffset() {
          lock.readLock().lock();
          try {
              return this.startOffset;
          } finally {
              lock.readLock().unlock();
          }
      }
  
      // Visible for testing
      long endOffset() {
          lock.readLock().lock();
          try {
              return this.endOffset;
          } finally {
              lock.readLock().unlock();
          }
      }
  
      // Visible for testing.
      int stateEpoch() {
          return stateEpoch;
      }
  
      // Visible for testing.
      Timer timer() {
          return timer;
      }
  
      // Visible for testing
      InitialReadGapOffset initialReadGapOffset() {
          return initialReadGapOffset;
      }
  
      // Visible for testing.
      Uuid fetchLock() {
          return fetchLock.get();
      }
  
      /**
       * The InitialReadGapOffset class is used to record the gap start and end offset of the probable gaps
       * of available records which are neither known to Persister nor to SharePartition. Share Partition
       * will use this information to determine the next fetch offset and should try to fetch the records
       * in the gap.
       */
      // Visible for Testing
      static class InitialReadGapOffset {
          private final long endOffset;
          private long gapStartOffset;
  
          InitialReadGapOffset(long endOffset, long gapStartOffset) {
              this.endOffset = endOffset;
              this.gapStartOffset = gapStartOffset;
          }
  
          long endOffset() {
              return endOffset;
          }
  
          long gapStartOffset() {
              return gapStartOffset;
          }
  
          void gapStartOffset(long gapStartOffset) {
              this.gapStartOffset = gapStartOffset;
          }
      }
  
      // Visible for testing
      final class AcquisitionLockTimerTask extends TimerTask {
          private final long expirationMs;
          private final String memberId;
          private final long firstOffset;
          private final long lastOffset;
  
          AcquisitionLockTimerTask(long delayMs, String memberId, long firstOffset, long lastOffset) {
              super(delayMs);
              this.expirationMs = time.hiResClockMs() + delayMs;
              this.memberId = memberId;
              this.firstOffset = firstOffset;
              this.lastOffset = lastOffset;
          }
  
          long expirationMs() {
              return expirationMs;
          }
  
          /**
           * The task is executed when the acquisition lock timeout is reached. The task releases the acquired records.
           */
          @Override
          public void run() {
              sharePartitionMetrics.recordAcquisitionLockTimeoutPerSec(lastOffset - firstOffset + 1);
              releaseAcquisitionLockOnTimeout(memberId, firstOffset, lastOffset);
          }
      }
  
      /**
       * The InFlightBatch maintains the in-memory state of the fetched records i.e. in-flight records.
       */
      final class InFlightBatch {
          // The offset of the first record in the batch that is fetched from the log.
          private final long firstOffset;
          // The last offset of the batch that is fetched from the log.
          private final long lastOffset;
  
          // The batch state of the fetched records. If the offset state map is empty then batchState
          // determines the state of the complete batch else individual offset determines the state of
          // the respective records.
          private InFlightState batchState;
  
          // The offset state map is used to track the state of the records per offset. However, the
          // offset state map is only required when the state of the offsets within same batch are
          // different. The states can be different when explicit offset acknowledgment is done which
          // is different from the batch state.
          private NavigableMap<Long, InFlightState> offsetState;
  
          InFlightBatch(String memberId, long firstOffset, long lastOffset, RecordState state,
              int deliveryCount, AcquisitionLockTimerTask acquisitionLockTimeoutTask
          ) {
              this.firstOffset = firstOffset;
              this.lastOffset = lastOffset;
              this.batchState = new InFlightState(state, deliveryCount, memberId, acquisitionLockTimeoutTask);
          }
  
          // Visible for testing.
          long firstOffset() {
              return firstOffset;
          }
  
          // Visible for testing.
          long lastOffset() {
              return lastOffset;
          }
  
          // Visible for testing.
          RecordState batchState() {
              return inFlightState().state;
          }
  
          // Visible for testing.
          String batchMemberId() {
              if (batchState == null) {
                  throw new IllegalStateException("The batch member id is not available as the offset state is maintained");
              }
              return batchState.memberId;
          }
  
          // Visible for testing.
          int batchDeliveryCount() {
              if (batchState == null) {
                  throw new IllegalStateException("The batch delivery count is not available as the offset state is maintained");
              }
              return batchState.deliveryCount;
          }
  
          // Visible for testing.
          AcquisitionLockTimerTask batchAcquisitionLockTimeoutTask() {
              return inFlightState().acquisitionLockTimeoutTask;
          }
  
          // Visible for testing.
          NavigableMap<Long, InFlightState> offsetState() {
              return offsetState;
          }
  
          private InFlightState inFlightState() {
              if (batchState == null) {
                  throw new IllegalStateException("The batch state is not available as the offset state is maintained");
              }
              return batchState;
          }
  
          private boolean batchHasOngoingStateTransition() {
              return inFlightState().hasOngoingStateTransition();
          }
  
          private void archiveBatch(String newMemberId) {
              inFlightState().archive(newMemberId);
          }
  
          private InFlightState tryUpdateBatchState(RecordState newState, DeliveryCountOps ops, int maxDeliveryCount, String newMemberId) {
              if (batchState == null) {
                  throw new IllegalStateException("The batch state update is not available as the offset state is maintained");
              }
              return batchState.tryUpdateState(newState, ops, maxDeliveryCount, newMemberId);
          }
  
          private InFlightState startBatchStateTransition(RecordState newState, DeliveryCountOps ops, int maxDeliveryCount,
                                                          String newMemberId) {
              if (batchState == null) {
                  throw new IllegalStateException("The batch state update is not available as the offset state is maintained");
              }
              return batchState.startStateTransition(newState, ops, maxDeliveryCount, newMemberId);
          }
  
          private void maybeInitializeOffsetStateUpdate() {
              if (offsetState == null) {
                  offsetState = new ConcurrentSkipListMap<>();
                  // The offset state map is not initialized hence initialize the state of the offsets
                  // from the first offset to the last offset. Mark the batch inflightState to null as
                  // the state of the records is maintained in the offset state map now.
                  for (long offset = this.firstOffset; offset <= this.lastOffset; offset++) {
                      if (batchState.acquisitionLockTimeoutTask != null) {
                          // The acquisition lock timeout task is already scheduled for the batch, hence we need to schedule
                          // the acquisition lock timeout task for the offset as well.
                          long delayMs = batchState.acquisitionLockTimeoutTask.expirationMs() - time.hiResClockMs();
                          AcquisitionLockTimerTask timerTask = acquisitionLockTimerTask(batchState.memberId, offset, offset, delayMs);
                          offsetState.put(offset, new InFlightState(batchState.state, batchState.deliveryCount, batchState.memberId, timerTask));
                          timer.add(timerTask);
                      } else {
                          offsetState.put(offset, new InFlightState(batchState.state, batchState.deliveryCount, batchState.memberId));
                      }
                  }
                  // Cancel the acquisition lock timeout task for the batch as the offset state is maintained.
                  if (batchState.acquisitionLockTimeoutTask != null) {
                      batchState.cancelAndClearAcquisitionLockTimeoutTask();
                  }
                  batchState = null;
              }
          }
  
          private void updateAcquisitionLockTimeout(AcquisitionLockTimerTask acquisitionLockTimeoutTask) {
              inFlightState().acquisitionLockTimeoutTask = acquisitionLockTimeoutTask;
          }
  
          @Override
          public String toString() {
              return "InFlightBatch(" +
                  "firstOffset=" + firstOffset +
                  ", lastOffset=" + lastOffset +
                  ", inFlightState=" + batchState +
                  ", offsetState=" + ((offsetState == null) ? "null" : offsetState) +
                  ")";
          }
      }
  
      /**
       * The InFlightState is used to track the state and delivery count of a record that has been
       * fetched from the leader. The state of the record is used to determine if the record should
       * be re-deliver or if it can be acknowledged or archived.
       */
      static final class InFlightState {
  
          // The state of the fetch batch records.
          private RecordState state;
          // The number of times the records has been delivered to the client.
          private int deliveryCount;
          // The member id of the client that is fetching/acknowledging the record.
          private String memberId;
          // The state of the records before the transition. In case we need to revert an in-flight state, we revert the above
          // attributes of InFlightState to this state, namely - state, deliveryCount and memberId.
          private InFlightState rollbackState;
          // The timer task for the acquisition lock timeout.
          private AcquisitionLockTimerTask acquisitionLockTimeoutTask;
  
  
          InFlightState(RecordState state, int deliveryCount, String memberId) {
              this(state, deliveryCount, memberId, null);
          }
  
          InFlightState(RecordState state, int deliveryCount, String memberId, AcquisitionLockTimerTask acquisitionLockTimeoutTask) {
              this.state = state;
              this.deliveryCount = deliveryCount;
              this.memberId = memberId;
              this.acquisitionLockTimeoutTask = acquisitionLockTimeoutTask;
          }
  
          // Visible for testing.
          RecordState state() {
              return state;
          }
  
          String memberId() {
              return memberId;
          }
  
          // Visible for testing.
          TimerTask acquisitionLockTimeoutTask() {
              return acquisitionLockTimeoutTask;
          }
  
          void updateAcquisitionLockTimeoutTask(AcquisitionLockTimerTask acquisitionLockTimeoutTask) throws IllegalArgumentException {
              if (this.acquisitionLockTimeoutTask != null) {
                  throw new IllegalArgumentException("Existing acquisition lock timeout exists, cannot override.");
              }
              this.acquisitionLockTimeoutTask = acquisitionLockTimeoutTask;
          }
  
          void cancelAndClearAcquisitionLockTimeoutTask() {
              acquisitionLockTimeoutTask.cancel();
              acquisitionLockTimeoutTask = null;
          }
  
          private boolean hasOngoingStateTransition() {
              if (rollbackState == null) {
                  // This case could occur when the batch/offset hasn't transitioned even once or the state transitions have
                  // been committed.
                  return false;
              }
              return rollbackState.state != null;
          }
  
          /**
           * Try to update the state of the records. The state of the records can only be updated if the
           * new state is allowed to be transitioned from old state. The delivery count is not changed
           * if the state update is unsuccessful.
           *
           * @param newState The new state of the records.
           * @param ops      The behavior on the delivery count.
           *
           * @return {@code InFlightState} if update succeeds, null otherwise. Returning state
           *         helps update chaining.
           */
          private InFlightState tryUpdateState(RecordState newState, DeliveryCountOps ops, int maxDeliveryCount, String newMemberId) {
              try {
                  if (newState == RecordState.AVAILABLE && ops != DeliveryCountOps.DECREASE && deliveryCount >= maxDeliveryCount) {
                      newState = RecordState.ARCHIVED;
                  }
                  state = state.validateTransition(newState);
                  if (newState != RecordState.ARCHIVED) {
                      deliveryCount = updatedDeliveryCount(ops);
                  }
                  memberId = newMemberId;
                  return this;
              } catch (IllegalStateException e) {
                  log.error("Failed to update state of the records", e);
                  return null;
              }
          }
  
          private int updatedDeliveryCount(DeliveryCountOps ops) {
              return switch (ops) {
                  case INCREASE -> deliveryCount + 1;
                  case DECREASE -> deliveryCount - 1;
                  // do nothing
                  case NO_OP -> deliveryCount;
              };
          }
  
          private void archive(String newMemberId) {
              state = RecordState.ARCHIVED;
              memberId = newMemberId;
          }
  
          private InFlightState startStateTransition(RecordState newState, DeliveryCountOps ops, int maxDeliveryCount, String newMemberId) {
              rollbackState = new InFlightState(state, deliveryCount, memberId, acquisitionLockTimeoutTask);
              return tryUpdateState(newState, ops, maxDeliveryCount, newMemberId);
          }
  
          private void completeStateTransition(boolean commit) {
              if (commit) {
                  rollbackState = null;
                  return;
              }
              state = rollbackState.state;
              deliveryCount = rollbackState.deliveryCount;
              memberId = rollbackState.memberId;
              rollbackState = null;
          }
  
          @Override
          public int hashCode() {
              return Objects.hash(state, deliveryCount, memberId);
          }
  
          @Override
          public boolean equals(Object o) {
              if (this == o) {
                  return true;
              }
              if (o == null || getClass() != o.getClass()) {
                  return false;
              }
              InFlightState that = (InFlightState) o;
              return state == that.state && deliveryCount == that.deliveryCount && memberId.equals(that.memberId);
          }
  
          @Override
          public String toString() {
              return "InFlightState(" +
                  "state=" + state.toString() +
                  ", deliveryCount=" + deliveryCount +
                  ", memberId=" + memberId +
                  ")";
          }
      }
  
      /**
       * FetchOffsetMetadata class is used to cache offset and its log metadata.
       */
      static final class OffsetMetadata {
          // This offset could be different from offsetMetadata.messageOffset if it's in the middle of a batch.
          private long offset;
          private LogOffsetMetadata offsetMetadata;
  
          OffsetMetadata() {
              offset = -1;
          }
  
          long offset() {
              return offset;
          }
  
          LogOffsetMetadata offsetMetadata() {
              return offsetMetadata;
          }
  
          void updateOffsetMetadata(long offset, LogOffsetMetadata offsetMetadata) {
              this.offset = offset;
              this.offsetMetadata = offsetMetadata;
          }
      }
  }
  