content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements. See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to You under the Apache License, Version 2.0
   * (the "License"); you may not use this file except in compliance with
   * the License. You may obtain a copy of the License at
   *
   *    http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.kafka.raft.internals;
  
  import org.apache.kafka.common.TopicPartition;
  import org.apache.kafka.common.Uuid;
  import org.apache.kafka.common.message.KRaftVersionRecord;
  import org.apache.kafka.common.metrics.KafkaMetric;
  import org.apache.kafka.common.metrics.Metrics;
  import org.apache.kafka.common.record.MemoryRecords;
  import org.apache.kafka.common.utils.BufferSupplier;
  import org.apache.kafka.common.utils.LogContext;
  import org.apache.kafka.raft.ExternalKRaftMetrics;
  import org.apache.kafka.raft.MockLog;
  import org.apache.kafka.raft.VoterSet;
  import org.apache.kafka.raft.VoterSetTest;
  import org.apache.kafka.server.common.KRaftVersion;
  import org.apache.kafka.server.common.OffsetAndEpoch;
  import org.apache.kafka.server.common.serialization.RecordSerde;
  import org.apache.kafka.snapshot.RecordsSnapshotWriter;
  
  import org.junit.jupiter.api.Test;
  import org.mockito.Mockito;
  
  import java.util.Optional;
  import java.util.stream.IntStream;
  
  import static org.junit.jupiter.api.Assertions.assertEquals;
  
  final class KRaftControlRecordStateMachineTest {
      private static final RecordSerde<String> STRING_SERDE = new StringSerde();
  
      private static MockLog buildLog() {
          return new MockLog(new TopicPartition("partition", 0), Uuid.randomUuid(), new LogContext());
      }
  
      private static KRaftControlRecordStateMachine buildPartitionListener(
          MockLog log,
          VoterSet staticVoterSet,
          KafkaRaftMetrics raftMetrics,
          ExternalKRaftMetrics externalMetrics
      ) {
          return new KRaftControlRecordStateMachine(
              staticVoterSet,
              log,
              STRING_SERDE,
              BufferSupplier.NO_CACHING,
              1024,
              new LogContext(),
              raftMetrics,
              externalMetrics
          );
      }
  
      private static KafkaMetric getNumberOfVoters(final Metrics metrics) {
          return metrics.metrics().get(metrics.metricName("number-of-voters", "raft-metrics"));
      }
  
      @Test
      void testEmptyPartition() {
          Metrics metrics = new Metrics();
          KafkaRaftMetrics raftMetrics = new KafkaRaftMetrics(metrics, "raft");
          ExternalKRaftMetrics externalMetrics = Mockito.mock(ExternalKRaftMetrics.class);
          MockLog log = buildLog();
          VoterSet voterSet = VoterSetTest.voterSet(VoterSetTest.voterMap(IntStream.of(1, 2, 3), true));
  
          KRaftControlRecordStateMachine partitionState = buildPartitionListener(log, voterSet, raftMetrics, externalMetrics);
  
          // This should be a no-op operation
          partitionState.updateState();
          Mockito.verifyNoInteractions(externalMetrics);
  
          assertEquals(voterSet, partitionState.lastVoterSet());
          assertEquals(3, getNumberOfVoters(metrics).metricValue());
      }
  
      @Test
      void testEmptyPartitionWithNoStaticVoters() {
          Metrics metrics = new Metrics();
          KafkaRaftMetrics raftMetrics = new KafkaRaftMetrics(metrics, "raft");
          ExternalKRaftMetrics externalMetrics = Mockito.mock(ExternalKRaftMetrics.class);
          MockLog log = buildLog();
  
          KRaftControlRecordStateMachine partitionState = buildPartitionListener(log, VoterSet.empty(), raftMetrics, externalMetrics);
  
          // This should be a no-op operation
          partitionState.updateState();
          Mockito.verifyNoInteractions(externalMetrics);
  
          assertEquals(VoterSet.empty(), partitionState.lastVoterSet());
          assertEquals(0, getNumberOfVoters(metrics).metricValue());
      }
  
      @Test
      void testUpdateWithoutSnapshot() {
          Metrics metrics = new Metrics();
          KafkaRaftMetrics raftMetrics = new KafkaRaftMetrics(metrics, "raft");
          ExternalKRaftMetrics externalMetrics = Mockito.mock(ExternalKRaftMetrics.class);
          MockLog log = buildLog();
          VoterSet staticVoterSet = VoterSetTest.voterSet(VoterSetTest.voterMap(IntStream.of(1, 2, 3), true));
          BufferSupplier bufferSupplier = BufferSupplier.NO_CACHING;
          int epoch = 1;
  
          KRaftControlRecordStateMachine partitionState = buildPartitionListener(log, staticVoterSet, raftMetrics, externalMetrics);
  
          assertEquals(3, getNumberOfVoters(metrics).metricValue());
  
          // Append the kraft.version control record
          KRaftVersion kraftVersion = KRaftVersion.KRAFT_VERSION_1;
          log.appendAsLeader(
              MemoryRecords.withKRaftVersionRecord(
                  log.endOffset().offset(),
                  0,
                  epoch,
                  bufferSupplier.get(300),
                  new KRaftVersionRecord().setKRaftVersion(kraftVersion.featureLevel())
              ),
              epoch
          );
  
          // Append the voter set control record
          VoterSet voterSet = VoterSetTest.voterSet(VoterSetTest.voterMap(IntStream.of(4, 5, 6), true));
          log.appendAsLeader(
              MemoryRecords.withVotersRecord(
                  log.endOffset().offset(),
                  0,
                  epoch,
                  bufferSupplier.get(300),
                  voterSet.toVotersRecord((short) 0)
              ),
              epoch
          );
  
          // Read the entire partition
          partitionState.updateState();
          Mockito.verify(externalMetrics, Mockito.times(1)).setIgnoredStaticVoters(true);
  
          assertEquals(voterSet, partitionState.lastVoterSet());
          assertEquals(Optional.of(voterSet), partitionState.voterSetAtOffset(log.endOffset().offset() - 1));
          assertEquals(kraftVersion, partitionState.kraftVersionAtOffset(log.endOffset().offset() - 1));
          assertEquals(3, getNumberOfVoters(metrics).metricValue());
      }
  
      @Test
      void testUpdateWithEmptySnapshot() {
          Metrics metrics = new Metrics();
          KafkaRaftMetrics raftMetrics = new KafkaRaftMetrics(metrics, "raft");
          ExternalKRaftMetrics externalMetrics = Mockito.mock(ExternalKRaftMetrics.class);
          MockLog log = buildLog();
          VoterSet staticVoterSet = VoterSetTest.voterSet(VoterSetTest.voterMap(IntStream.of(1, 2, 3), true));
          BufferSupplier bufferSupplier = BufferSupplier.NO_CACHING;
          int epoch = 1;
  
          KRaftControlRecordStateMachine partitionState = buildPartitionListener(log, staticVoterSet, raftMetrics, externalMetrics);
  
          assertEquals(3, getNumberOfVoters(metrics).metricValue());
  
          // Create a snapshot that doesn't have any kraft.version or voter set control records
          RecordsSnapshotWriter.Builder builder = new RecordsSnapshotWriter.Builder()
              .setRawSnapshotWriter(log.createNewSnapshotUnchecked(new OffsetAndEpoch(10, epoch)).get());
          try (RecordsSnapshotWriter<?> writer = builder.build(STRING_SERDE)) {
              writer.freeze();
          }
          log.truncateToLatestSnapshot();
  
          // Append the kraft.version control record
          KRaftVersion kraftVersion = KRaftVersion.KRAFT_VERSION_1;
          log.appendAsLeader(
              MemoryRecords.withKRaftVersionRecord(
                  log.endOffset().offset(),
                  0,
                  epoch,
                  bufferSupplier.get(300),
                  new KRaftVersionRecord().setKRaftVersion(kraftVersion.featureLevel())
              ),
              epoch
          );
  
          // Append the voter set control record
          VoterSet voterSet = VoterSetTest.voterSet(VoterSetTest.voterMap(IntStream.of(4, 5, 6), true));
          log.appendAsLeader(
              MemoryRecords.withVotersRecord(
                  log.endOffset().offset(),
                  0,
                  epoch,
                  bufferSupplier.get(300),
                  voterSet.toVotersRecord((short) 0)
              ),
              epoch
          );
  
          // Read the entire partition
          partitionState.updateState();
          Mockito.verify(externalMetrics, Mockito.times(1)).setIgnoredStaticVoters(true);
  
          assertEquals(voterSet, partitionState.lastVoterSet());
          assertEquals(Optional.of(voterSet), partitionState.voterSetAtOffset(log.endOffset().offset() - 1));
          assertEquals(kraftVersion, partitionState.kraftVersionAtOffset(log.endOffset().offset() - 1));
          assertEquals(3, getNumberOfVoters(metrics).metricValue());
      }
  
      @Test
      void testUpdateWithSnapshot() {
          Metrics metrics = new Metrics();
          KafkaRaftMetrics raftMetrics = new KafkaRaftMetrics(metrics, "raft");
          ExternalKRaftMetrics externalMetrics = Mockito.mock(ExternalKRaftMetrics.class);
          MockLog log = buildLog();
          VoterSet staticVoterSet = VoterSetTest.voterSet(VoterSetTest.voterMap(IntStream.of(1, 2, 3), true));
          int epoch = 1;
  
          KRaftControlRecordStateMachine partitionState = buildPartitionListener(log, staticVoterSet, raftMetrics, externalMetrics);
  
          assertEquals(3, getNumberOfVoters(metrics).metricValue());
  
          // Create a snapshot that has kraft.version and voter set control records
          KRaftVersion kraftVersion = KRaftVersion.KRAFT_VERSION_1;
          VoterSet voterSet = VoterSetTest.voterSet(VoterSetTest.voterMap(IntStream.of(4, 5, 6), true));
  
          RecordsSnapshotWriter.Builder builder = new RecordsSnapshotWriter.Builder()
              .setRawSnapshotWriter(log.createNewSnapshotUnchecked(new OffsetAndEpoch(10, epoch)).get())
              .setKraftVersion(kraftVersion)
              .setVoterSet(Optional.of(voterSet));
          try (RecordsSnapshotWriter<?> writer = builder.build(STRING_SERDE)) {
              writer.freeze();
          }
          log.truncateToLatestSnapshot();
  
          // Read the entire partition
          partitionState.updateState();
          Mockito.verify(externalMetrics, Mockito.times(1)).setIgnoredStaticVoters(true);
  
          assertEquals(voterSet, partitionState.lastVoterSet());
          assertEquals(Optional.of(voterSet), partitionState.voterSetAtOffset(log.endOffset().offset() - 1));
          assertEquals(kraftVersion, partitionState.kraftVersionAtOffset(log.endOffset().offset() - 1));
          assertEquals(3, getNumberOfVoters(metrics).metricValue());
      }
  
      @Test
      void testUpdateWithSnapshotAndLogOverride() {
          Metrics metrics = new Metrics();
          KafkaRaftMetrics raftMetrics = new KafkaRaftMetrics(metrics, "raft");
          ExternalKRaftMetrics externalMetrics = Mockito.mock(ExternalKRaftMetrics.class);
          MockLog log = buildLog();
          VoterSet staticVoterSet = VoterSetTest.voterSet(VoterSetTest.voterMap(IntStream.of(1, 2, 3), true));
          BufferSupplier bufferSupplier = BufferSupplier.NO_CACHING;
          int epoch = 1;
  
          KRaftControlRecordStateMachine partitionState = buildPartitionListener(log, staticVoterSet, raftMetrics, externalMetrics);
  
          assertEquals(3, getNumberOfVoters(metrics).metricValue());
  
          // Create a snapshot that has kraft.version and voter set control records
          KRaftVersion kraftVersion = KRaftVersion.KRAFT_VERSION_1;
          VoterSet snapshotVoterSet = VoterSetTest.voterSet(VoterSetTest.voterMap(IntStream.of(4, 5, 6), true));
  
          OffsetAndEpoch snapshotId = new OffsetAndEpoch(10, epoch);
          RecordsSnapshotWriter.Builder builder = new RecordsSnapshotWriter.Builder()
              .setRawSnapshotWriter(log.createNewSnapshotUnchecked(snapshotId).get())
              .setKraftVersion(kraftVersion)
              .setVoterSet(Optional.of(snapshotVoterSet));
          try (RecordsSnapshotWriter<?> writer = builder.build(STRING_SERDE)) {
              writer.freeze();
          }
          log.truncateToLatestSnapshot();
  
          // Append the voter set control record
          VoterSet voterSet = snapshotVoterSet.addVoter(VoterSetTest.voterNode(7, true)).get();
          log.appendAsLeader(
              MemoryRecords.withVotersRecord(
                  log.endOffset().offset(),
                  0,
                  epoch,
                  bufferSupplier.get(300),
                  voterSet.toVotersRecord((short) 0)
              ),
              epoch
          );
  
          // Read the entire partition
          partitionState.updateState();
          Mockito.verify(externalMetrics, Mockito.times(2)).setIgnoredStaticVoters(true);
  
          assertEquals(voterSet, partitionState.lastVoterSet());
          assertEquals(Optional.of(voterSet), partitionState.voterSetAtOffset(log.endOffset().offset() - 1));
          assertEquals(kraftVersion, partitionState.kraftVersionAtOffset(log.endOffset().offset() - 1));
          assertEquals(4, getNumberOfVoters(metrics).metricValue());
  
          // Check the voter set at the snapshot
          assertEquals(Optional.of(snapshotVoterSet), partitionState.voterSetAtOffset(snapshotId.offset() - 1));
      }
  
      @Test
      void testTruncateTo() {
          Metrics metrics = new Metrics();
          KafkaRaftMetrics raftMetrics = new KafkaRaftMetrics(metrics, "raft");
          ExternalKRaftMetrics externalMetrics = Mockito.mock(ExternalKRaftMetrics.class);
          MockLog log = buildLog();
          VoterSet staticVoterSet = VoterSetTest.voterSet(VoterSetTest.voterMap(IntStream.of(1, 2, 3), true));
          BufferSupplier bufferSupplier = BufferSupplier.NO_CACHING;
          int epoch = 1;
  
          KRaftControlRecordStateMachine partitionState = buildPartitionListener(log, staticVoterSet, raftMetrics, externalMetrics);
  
          assertEquals(3, getNumberOfVoters(metrics).metricValue());
  
          // Append the kraft.version control record
          KRaftVersion kraftVersion = KRaftVersion.KRAFT_VERSION_1;
          log.appendAsLeader(
              MemoryRecords.withKRaftVersionRecord(
                  log.endOffset().offset(),
                  0,
                  epoch,
                  bufferSupplier.get(300),
                  new KRaftVersionRecord().setKRaftVersion(kraftVersion.featureLevel())
              ),
              epoch
          );
  
          // Append the voter set control record
          long firstVoterSetOffset = log.endOffset().offset();
          VoterSet firstVoterSet = VoterSetTest.voterSet(VoterSetTest.voterMap(IntStream.of(4, 5, 6), true));
          log.appendAsLeader(
              MemoryRecords.withVotersRecord(
                  firstVoterSetOffset,
                  0,
                  epoch,
                  bufferSupplier.get(300),
                  firstVoterSet.toVotersRecord((short) 0)
              ),
              epoch
          );
  
          // Append another voter set control record
          long voterSetOffset = log.endOffset().offset();
          VoterSet voterSet = firstVoterSet.addVoter(VoterSetTest.voterNode(7, true)).get();
          log.appendAsLeader(
              MemoryRecords.withVotersRecord(
                  voterSetOffset,
                  0,
                  epoch,
                  bufferSupplier.get(300),
                  voterSet.toVotersRecord((short) 0)
              ),
              epoch
          );
  
          // Read the entire partition
          partitionState.updateState();
          Mockito.verify(externalMetrics, Mockito.times(2)).setIgnoredStaticVoters(true);
  
          assertEquals(voterSet, partitionState.lastVoterSet());
          assertEquals(4, getNumberOfVoters(metrics).metricValue());
  
          // Truncate log and listener
          log.truncateTo(voterSetOffset);
          partitionState.truncateNewEntries(voterSetOffset);
  
          assertEquals(firstVoterSet, partitionState.lastVoterSet());
          assertEquals(3, getNumberOfVoters(metrics).metricValue());
  
          // Truncate the entire log
          log.truncateTo(0);
          partitionState.truncateNewEntries(0);
          Mockito.verify(externalMetrics, Mockito.times(1)).setIgnoredStaticVoters(false);
  
          assertEquals(staticVoterSet, partitionState.lastVoterSet());
          assertEquals(3, getNumberOfVoters(metrics).metricValue());
      }
  
      @Test
      void testTrimPrefixTo() {
          Metrics metrics = new Metrics();
          KafkaRaftMetrics raftMetrics = new KafkaRaftMetrics(metrics, "raft");
          ExternalKRaftMetrics externalMetrics = Mockito.mock(ExternalKRaftMetrics.class);
          MockLog log = buildLog();
          VoterSet staticVoterSet = VoterSetTest.voterSet(VoterSetTest.voterMap(IntStream.of(1, 2, 3), true));
          BufferSupplier bufferSupplier = BufferSupplier.NO_CACHING;
          int epoch = 1;
  
          KRaftControlRecordStateMachine partitionState = buildPartitionListener(log, staticVoterSet, raftMetrics, externalMetrics);
  
          assertEquals(3, getNumberOfVoters(metrics).metricValue());
  
          // Append the kraft.version control record
          long kraftVersionOffset = log.endOffset().offset();
          KRaftVersion kraftVersion = KRaftVersion.KRAFT_VERSION_1;
          log.appendAsLeader(
              MemoryRecords.withKRaftVersionRecord(
                  kraftVersionOffset,
                  0,
                  epoch,
                  bufferSupplier.get(300),
                  new KRaftVersionRecord().setKRaftVersion(kraftVersion.featureLevel())
              ),
              epoch
          );
  
          // Append the voter set control record
          long firstVoterSetOffset = log.endOffset().offset();
          VoterSet firstVoterSet = VoterSetTest.voterSet(VoterSetTest.voterMap(IntStream.of(4, 5, 6), true));
          log.appendAsLeader(
              MemoryRecords.withVotersRecord(
                  firstVoterSetOffset,
                  0,
                  epoch,
                  bufferSupplier.get(300),
                  firstVoterSet.toVotersRecord((short) 0)
              ),
              epoch
          );
  
          // Append another voter set control record
          long voterSetOffset = log.endOffset().offset();
          VoterSet voterSet = firstVoterSet.addVoter(VoterSetTest.voterNode(7, true)).get();
          log.appendAsLeader(
              MemoryRecords.withVotersRecord(
                  voterSetOffset,
                  0,
                  epoch,
                  bufferSupplier.get(300),
                  voterSet.toVotersRecord((short) 0)
              ),
              epoch
          );
  
          // Read the entire partition
          partitionState.updateState();
          Mockito.verify(externalMetrics, Mockito.times(2)).setIgnoredStaticVoters(true);
  
          assertEquals(voterSet, partitionState.lastVoterSet());
          assertEquals(kraftVersion, partitionState.kraftVersionAtOffset(kraftVersionOffset));
          assertEquals(4, getNumberOfVoters(metrics).metricValue());
  
          // Trim the prefix for the partition listener up to the kraft.version
          partitionState.truncateOldEntries(kraftVersionOffset);
          assertEquals(kraftVersion, partitionState.kraftVersionAtOffset(kraftVersionOffset));
  
          // Trim the prefix for the partition listener up to the first voter set
          partitionState.truncateOldEntries(firstVoterSetOffset);
          assertEquals(kraftVersion, partitionState.kraftVersionAtOffset(kraftVersionOffset));
          assertEquals(Optional.of(firstVoterSet), partitionState.voterSetAtOffset(firstVoterSetOffset));
  
          // Trim the prefix for the partition listener up to the second voter set
          partitionState.truncateOldEntries(voterSetOffset);
          assertEquals(kraftVersion, partitionState.kraftVersionAtOffset(kraftVersionOffset));
          assertEquals(Optional.empty(), partitionState.voterSetAtOffset(firstVoterSetOffset));
          assertEquals(Optional.of(voterSet), partitionState.voterSetAtOffset(voterSetOffset));
          assertEquals(4, getNumberOfVoters(metrics).metricValue());
      }
  }
  