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.image.loader;
  
  import org.apache.kafka.common.Uuid;
  import org.apache.kafka.common.message.SnapshotHeaderRecord;
  import org.apache.kafka.common.metadata.AbortTransactionRecord;
  import org.apache.kafka.common.metadata.BeginTransactionRecord;
  import org.apache.kafka.common.metadata.EndTransactionRecord;
  import org.apache.kafka.common.metadata.FeatureLevelRecord;
  import org.apache.kafka.common.metadata.PartitionRecord;
  import org.apache.kafka.common.metadata.RemoveTopicRecord;
  import org.apache.kafka.common.metadata.TopicRecord;
  import org.apache.kafka.common.utils.MockTime;
  import org.apache.kafka.image.MetadataDelta;
  import org.apache.kafka.image.MetadataImage;
  import org.apache.kafka.image.MetadataProvenance;
  import org.apache.kafka.image.publisher.MetadataPublisher;
  import org.apache.kafka.raft.Batch;
  import org.apache.kafka.raft.BatchReader;
  import org.apache.kafka.raft.ControlRecord;
  import org.apache.kafka.raft.LeaderAndEpoch;
  import org.apache.kafka.server.common.ApiMessageAndVersion;
  import org.apache.kafka.server.common.MetadataVersion;
  import org.apache.kafka.server.common.OffsetAndEpoch;
  import org.apache.kafka.server.fault.MockFaultHandler;
  import org.apache.kafka.snapshot.SnapshotReader;
  import org.apache.kafka.test.TestUtils;
  
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.api.Timeout;
  import org.junit.jupiter.params.ParameterizedTest;
  import org.junit.jupiter.params.provider.CsvSource;
  import org.junit.jupiter.params.provider.ValueSource;
  
  import java.util.ArrayList;
  import java.util.Iterator;
  import java.util.List;
  import java.util.Optional;
  import java.util.OptionalLong;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.RejectedExecutionException;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicReference;
  
  import static org.apache.kafka.server.common.MetadataVersion.IBP_3_5_IV0;
  import static org.apache.kafka.server.common.MetadataVersion.MINIMUM_VERSION;
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertFalse;
  import static org.junit.jupiter.api.Assertions.assertNotNull;
  import static org.junit.jupiter.api.Assertions.assertNull;
  import static org.junit.jupiter.api.Assertions.assertThrows;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  
  
  @Timeout(value = 40)
  public class MetadataLoaderTest {
      @Test
      public void testCreateAndClose() throws Exception {
          MockFaultHandler faultHandler = new MockFaultHandler("testCreateAndClose");
          try (MetadataLoader loader = new MetadataLoader.Builder().
                  setFaultHandler(faultHandler).
                  setHighWaterMarkAccessor(OptionalLong::empty).
                  build()) {
              assertEquals(-1L, loader.lastAppliedOffset());
          }
          faultHandler.maybeRethrowFirstException();
      }
  
      static class MockPublisher implements MetadataPublisher {
          final CompletableFuture<Void> firstPublish = new CompletableFuture<>();
          private final String name;
          volatile MetadataDelta latestDelta = null;
          volatile MetadataImage latestImage = null;
          volatile LogDeltaManifest latestLogDeltaManifest = null;
          volatile SnapshotManifest latestSnapshotManifest = null;
          volatile boolean closed = false;
  
          MockPublisher() {
              this("MockPublisher");
          }
  
          MockPublisher(String name) {
              this.name = name;
          }
  
          @Override
          public String name() {
              return name;
          }
  
          @Override
          public void onMetadataUpdate(
              MetadataDelta delta,
              MetadataImage newImage,
              LoaderManifest manifest
          ) {
              latestDelta = delta;
              latestImage = newImage;
              switch (manifest.type()) {
                  case LOG_DELTA:
                      latestLogDeltaManifest = (LogDeltaManifest) manifest;
                      break;
                  case SNAPSHOT:
                      latestSnapshotManifest = (SnapshotManifest) manifest;
                      break;
                  default:
                      throw new RuntimeException("Invalid manifest type " + manifest.type());
              }
              firstPublish.complete(null);
          }
  
          @Override
          public void close() throws Exception {
              firstPublish.completeExceptionally(new RejectedExecutionException());
              closed = true;
          }
      }
  
      /**
       * Install 2 publishers and check that the publishers that were installed are closed when the
       * loader is closed.
       */
      @Test
      public void testInstallPublishers() throws Exception {
          MockFaultHandler faultHandler = new MockFaultHandler("testInstallPublishers");
          List<MockPublisher> publishers = List.of(new MockPublisher("a"),
                  new MockPublisher("b"),
                  new MockPublisher("c"));
          try (MetadataLoader loader = new MetadataLoader.Builder().
                  setFaultHandler(faultHandler).
                  setHighWaterMarkAccessor(OptionalLong::empty).
                  build()) {
              loader.installPublishers(publishers.subList(0, 2)).get();
          }
          assertTrue(publishers.get(0).closed);
          assertNull(publishers.get(0).latestImage);
          assertTrue(publishers.get(1).closed);
          assertNull(publishers.get(1).latestImage);
          assertFalse(publishers.get(2).closed);
          assertNull(publishers.get(2).latestImage);
          faultHandler.maybeRethrowFirstException();
      }
  
      static class MockSnapshotReader implements SnapshotReader<ApiMessageAndVersion> {
          private final MetadataProvenance provenance;
          private final Iterator<Batch<ApiMessageAndVersion>> iterator;
          private MockTime time = null;
          boolean closed = false;
  
          static MockSnapshotReader fromRecordLists(
              MetadataProvenance provenance,
              List<List<ApiMessageAndVersion>> lists
          ) {
              List<Batch<ApiMessageAndVersion>> batches = lists
                  .stream()
                  .map(records -> Batch.data(0, 0, 0, 0, records))
                  .toList();
  
              return new MockSnapshotReader(provenance, batches);
          }
  
          MockSnapshotReader(
              MetadataProvenance provenance,
              List<Batch<ApiMessageAndVersion>> batches
          ) {
              this.provenance = provenance;
              this.iterator = batches.iterator();
          }
  
          MockSnapshotReader setTime(MockTime time) {
              this.time = time;
              return this;
          }
  
          @Override
          public OffsetAndEpoch snapshotId() {
              return provenance.snapshotId();
          }
  
          @Override
          public long lastContainedLogOffset() {
              return provenance.lastContainedOffset();
          }
  
          @Override
          public int lastContainedLogEpoch() {
              return provenance.lastContainedEpoch();
          }
  
          @Override
          public long lastContainedLogTimestamp() {
              return provenance.lastContainedLogTimeMs();
          }
  
          @Override
          public void close() {
              closed = true;
          }
  
          @Override
          public boolean hasNext() {
              if (time != null) time.sleep(1);
              return iterator.hasNext();
          }
  
          @Override
          public Batch<ApiMessageAndVersion> next() {
              if (time != null) time.sleep(1);
              return iterator.next();
          }
      }
  
      /**
       * Test that a publisher cannot be installed more than once.
       */
      @ParameterizedTest
      @CsvSource(value = {"false,false", "false,true", "true,false", "true,true"})
      public void testPublisherCannotBeInstalledMoreThanOnce(
          boolean loadSnapshot,
          boolean sameObject
      ) throws Exception {
          MockFaultHandler faultHandler = new MockFaultHandler("testPublisherCannotBeInstalledMoreThanOnce");
          MockPublisher publisher = new MockPublisher();
          try (MetadataLoader loader = new MetadataLoader.Builder().
                  setFaultHandler(faultHandler).
                  setHighWaterMarkAccessor(() -> OptionalLong.of(0L)).
                  build()) {
              loader.installPublishers(List.of(publisher)).get();
              if (loadSnapshot) {
                  MockSnapshotReader snapshotReader = new MockSnapshotReader(
                      new MetadataProvenance(200, 100, 4000, true),
                      List.of(
                          Batch.control(
                              200,
                              100,
                              4000,
                              10,
                              List.of(ControlRecord.of(new SnapshotHeaderRecord()))
                          ),
                          Batch.data(0, 0, 0, 0,
                              List.of(new ApiMessageAndVersion(new FeatureLevelRecord().
                                  setName(MetadataVersion.FEATURE_NAME).
                                  setFeatureLevel(MINIMUM_VERSION.featureLevel()), (short) 0))
                          )
                      )
                  );
                  loader.handleLoadSnapshot(snapshotReader);
                  TestUtils.retryOnExceptionWithTimeout(30_000, () ->
                      assertEquals(1L, loader.metrics().handleLoadSnapshotCount())
                  );
              } else {
                  TestUtils.retryOnExceptionWithTimeout(30_000, () ->
                      assertEquals(0L, loader.metrics().handleLoadSnapshotCount())
                  );
              }
              loader.waitForAllEventsToBeHandled();
              if (sameObject) {
                  assertEquals("testPublisherCannotBeInstalledMoreThanOnce: Attempted to install " +
                      "publisher MockPublisher, which is already installed.",
                          assertThrows(ExecutionException.class,
                                  () -> loader.installPublishers(List.of(publisher)).get()).
                                  getCause().getMessage());
              } else {
                  assertEquals("testPublisherCannotBeInstalledMoreThanOnce: Attempted to install " +
                      "a new publisher named MockPublisher, but there is already a publisher with that name.",
                          assertThrows(ExecutionException.class,
                                  () -> loader.installPublishers(List.of(new MockPublisher())).get()).
                                  getCause().getMessage());
              }
          }
      }
  
      /**
       * Install 2 publishers and remove one.
       */
      @Test
      public void testRemovePublisher() throws Exception {
          MockFaultHandler faultHandler = new MockFaultHandler("testRemovePublisher");
          List<MockPublisher> publishers = List.of(new MockPublisher("a"),
                  new MockPublisher("b"),
                  new MockPublisher("c"));
          try (MetadataLoader loader = new MetadataLoader.Builder().
                  setFaultHandler(faultHandler).
                  setHighWaterMarkAccessor(() -> OptionalLong.of(1L)).
                  build()) {
              loader.installPublishers(publishers.subList(0, 2)).get();
              loader.removeAndClosePublisher(publishers.get(1)).get();
              MockSnapshotReader snapshotReader = MockSnapshotReader.fromRecordLists(
                  new MetadataProvenance(100, 50, 2000, true),
                  List.of(List.of(new ApiMessageAndVersion(
                      new FeatureLevelRecord().
                          setName(MetadataVersion.FEATURE_NAME).
                          setFeatureLevel(MINIMUM_VERSION.featureLevel()), (short) 0))));
              assertFalse(snapshotReader.closed);
              loader.handleLoadSnapshot(snapshotReader);
              loader.waitForAllEventsToBeHandled();
              assertTrue(snapshotReader.closed);
              publishers.get(0).firstPublish.get(1, TimeUnit.MINUTES);
              loader.removeAndClosePublisher(publishers.get(0)).get();
          }
          assertTrue(publishers.get(0).closed);
          assertEquals(Optional.of(MINIMUM_VERSION),
                  publishers.get(0).latestImage.features().metadataVersion());
          assertTrue(publishers.get(1).closed);
          assertNull(publishers.get(1).latestImage);
          assertFalse(publishers.get(2).closed);
          assertNull(publishers.get(2).latestImage);
          faultHandler.maybeRethrowFirstException();
      }
  
      /**
       * Test loading a snapshot with 0 records.
       */
      @Test
      public void testLoadEmptySnapshot() throws Exception {
          MockFaultHandler faultHandler = new MockFaultHandler("testLoadEmptySnapshot");
          MockTime time = new MockTime();
          List<MockPublisher> publishers = List.of(new MockPublisher());
          try (MetadataLoader loader = new MetadataLoader.Builder().
                  setFaultHandler(faultHandler).
                  setTime(time).
                  setHighWaterMarkAccessor(() -> OptionalLong.of(0L)).
                  build()) {
              loader.installPublishers(publishers).get();
              loadEmptySnapshot(loader, 200);
              loader.waitForAllEventsToBeHandled();
              assertFalse(publishers.get(0).firstPublish.isDone());
              loader.handleCommit(MockBatchReader.newSingleBatchReader(250, 50, List.of(
                  new ApiMessageAndVersion(new FeatureLevelRecord()
                      .setName(MetadataVersion.FEATURE_NAME)
                      .setFeatureLevel(MINIMUM_VERSION.featureLevel()), (short) 0)))
              );
              publishers.get(0).firstPublish.get(10, TimeUnit.SECONDS);
              assertEquals(250L, loader.lastAppliedOffset());
              loadEmptySnapshot(loader, 300);
              assertEquals(300L, loader.lastAppliedOffset());
              assertEquals(new SnapshotManifest(new MetadataProvenance(300, 100, 4000, true), 3000000L),
                  publishers.get(0).latestSnapshotManifest);
              assertEquals(MINIMUM_VERSION,
                  loader.metrics().currentMetadataVersion());
          }
          assertTrue(publishers.get(0).closed);
          assertEquals(Optional.of(MINIMUM_VERSION), publishers.get(0).latestImage.features().metadataVersion());
          var latestImage = publishers.get(0).latestImage;
          assertFalse(latestImage.isEmpty());
          assertFalse(latestImage.features().isEmpty());
          assertTrue(latestImage.features().finalizedVersions().isEmpty());
          assertTrue(latestImage.cluster().isEmpty());
          assertTrue(latestImage.topics().isEmpty());
          assertTrue(latestImage.cluster().isEmpty());
          assertTrue(latestImage.configs().isEmpty());
          assertTrue(latestImage.producerIds().isEmpty());
          assertTrue(latestImage.acls().isEmpty());
          assertTrue(latestImage.scram().isEmpty());
          assertTrue(latestImage.delegationTokens().isEmpty());
          faultHandler.maybeRethrowFirstException();
      }
  
      private void loadEmptySnapshot(
          MetadataLoader loader,
          long offset
      ) throws Exception {
          MockSnapshotReader snapshotReader = new MockSnapshotReader(
              new MetadataProvenance(offset, 100, 4000, true),
              List.of(
                  Batch.control(
                      200,
                      100,
                      4000,
                      10,
                      List.of(ControlRecord.of(new SnapshotHeaderRecord()))
                  )
              )
          );
          if (loader.time() instanceof MockTime) {
              snapshotReader.setTime((MockTime) loader.time());
          }
          loader.handleLoadSnapshot(snapshotReader);
          loader.waitForAllEventsToBeHandled();
      }
  
      static class MockBatchReader implements BatchReader<ApiMessageAndVersion> {
          private final long baseOffset;
          private final Iterator<Batch<ApiMessageAndVersion>> iterator;
          private boolean closed = false;
          private MockTime time = null;
  
          static MockBatchReader newSingleBatchReader(
              long batchBaseOffset,
              int epoch,
              List<ApiMessageAndVersion> records
          ) {
              return new MockBatchReader(batchBaseOffset,
                  List.of(newBatch(batchBaseOffset, epoch, records)));
          }
  
          static Batch<ApiMessageAndVersion> newBatch(
              long batchBaseOffset,
              int epoch,
              List<ApiMessageAndVersion> records
          ) {
              return Batch.data(batchBaseOffset, epoch, 0, 0, records);
          }
  
          MockBatchReader(
              long baseOffset,
              List<Batch<ApiMessageAndVersion>> batches
          ) {
              this.baseOffset = baseOffset;
              this.iterator = batches.iterator();
          }
  
          private MockBatchReader setTime(MockTime time) {
              this.time = time;
              return this;
          }
  
          @Override
          public long baseOffset() {
              return baseOffset;
          }
  
          @Override
          public OptionalLong lastOffset() {
              return OptionalLong.empty();
          }
  
          @Override
          public void close() {
              this.closed = true;
          }
  
          @Override
          public boolean hasNext() {
              if (time != null) time.sleep(1);
              return iterator.hasNext();
          }
  
          @Override
          public Batch<ApiMessageAndVersion> next() {
              if (time != null) time.sleep(1);
              return iterator.next();
          }
      }
  
      /**
       * Test loading a batch with 0 records.
       */
      @Test
      public void testLoadEmptyBatch() throws Exception {
          MockFaultHandler faultHandler = new MockFaultHandler("testLoadEmptyBatch");
          MockTime time = new MockTime();
          List<MockPublisher> publishers = List.of(new MockPublisher());
          try (MetadataLoader loader = new MetadataLoader.Builder().
                  setFaultHandler(faultHandler).
                  setTime(time).
                  setHighWaterMarkAccessor(() -> OptionalLong.of(1L)).
                  build()) {
              loader.installPublishers(publishers).get();
              loadTestSnapshot(loader, 200);
              publishers.get(0).firstPublish.get(10, TimeUnit.SECONDS);
              MockBatchReader batchReader = new MockBatchReader(
                  300,
                  List.of(
                      Batch.control(
                          300,
                          100,
                          4000,
                          10,
                          List.of(ControlRecord.of(new SnapshotHeaderRecord()))
                      )
                  )
              ).setTime(time);
              loader.handleCommit(batchReader);
              loader.waitForAllEventsToBeHandled();
              assertTrue(batchReader.closed);
              assertEquals(300L, loader.lastAppliedOffset());
          }
          assertTrue(publishers.get(0).closed);
          assertEquals(
              LogDeltaManifest.newBuilder()
                  .provenance(new MetadataProvenance(300, 100, 4000, true))
                  .leaderAndEpoch(LeaderAndEpoch.UNKNOWN)
                  .numBatches(1)
                  .elapsedNs(0L)
                  .numBytes(10)
                  .build(),
              publishers.get(0).latestLogDeltaManifest);
          assertEquals(Optional.of(MINIMUM_VERSION),
              publishers.get(0).latestImage.features().metadataVersion());
          faultHandler.maybeRethrowFirstException();
      }
  
      /**
       * Test that the lastAppliedOffset moves forward as expected.
       */
      @Test
      public void testLastAppliedOffset() throws Exception {
          MockFaultHandler faultHandler = new MockFaultHandler("testLastAppliedOffset");
          List<MockPublisher> publishers = List.of(new MockPublisher("a"),
                  new MockPublisher("b"));
          try (MetadataLoader loader = new MetadataLoader.Builder().
                  setFaultHandler(faultHandler).
                  setHighWaterMarkAccessor(() -> OptionalLong.of(1L)).
                  build()) {
              loader.installPublishers(publishers).get();
              loader.handleLoadSnapshot(MockSnapshotReader.fromRecordLists(
                  new MetadataProvenance(200, 100, 4000, true), List.of(
                      List.of(new ApiMessageAndVersion(new FeatureLevelRecord().
                          setName(MetadataVersion.FEATURE_NAME).
                          setFeatureLevel(MINIMUM_VERSION.featureLevel()), (short) 0)),
                      List.of(new ApiMessageAndVersion(new TopicRecord().
                          setName("foo").
                          setTopicId(Uuid.fromString("Uum7sfhHQP-obSvfywmNUA")), (short) 0))
                  )));
              for (MockPublisher publisher : publishers) {
                  publisher.firstPublish.get(1, TimeUnit.MINUTES);
              }
              loader.waitForAllEventsToBeHandled();
              assertEquals(200L, loader.lastAppliedOffset());
              loader.handleCommit(new MockBatchReader(201, List.of(
                  MockBatchReader.newBatch(201, 100, List.of(
                      new ApiMessageAndVersion(new RemoveTopicRecord().
                          setTopicId(Uuid.fromString("Uum7sfhHQP-obSvfywmNUA")), (short) 0))))));
              loader.waitForAllEventsToBeHandled();
              assertEquals(201L, loader.lastAppliedOffset());
          }
          for (int i = 0; i < 2; i++) {
              assertTrue(publishers.get(i).closed);
              assertTrue(publishers.get(i).closed);
              assertEquals(Optional.of(MINIMUM_VERSION), publishers.get(i).latestImage.features().metadataVersion());
          }
          faultHandler.maybeRethrowFirstException();
      }
  
      /**
       * Test that we do not leave the catchingUp state until we have loaded up to the high
       * watermark.
       */
      @Test
      public void testCatchingUpState() throws Exception {
          MockFaultHandler faultHandler = new MockFaultHandler("testLastAppliedOffset");
          List<MockPublisher> publishers = List.of(new MockPublisher("a"),
                  new MockPublisher("b"));
          AtomicReference<OptionalLong> highWaterMark = new AtomicReference<>(OptionalLong.empty());
          try (MetadataLoader loader = new MetadataLoader.Builder().
                  setFaultHandler(faultHandler).
                  setHighWaterMarkAccessor(highWaterMark::get).
                  build()) {
              loader.installPublishers(publishers).get();
              loadTestSnapshot(loader, 200);
  
              // We don't update lastAppliedOffset because we're still in catchingUp state due to
              // highWaterMark being OptionalLong.empty (aka unknown).
              assertEquals(-1L, loader.lastAppliedOffset());
              assertFalse(publishers.get(0).firstPublish.isDone());
  
              // This still doesn't advance lastAppliedOffset since the high watermark at 221
              // is greater than our snapshot at 210.
              highWaterMark.set(OptionalLong.of(221));
              loadTestSnapshot(loader, 210);
              assertEquals(-1L, loader.lastAppliedOffset());
  
              // Loading a test snapshot at 220 allows us to leave catchUp state.
              loadTestSnapshot(loader, 220);
              assertEquals(220L, loader.lastAppliedOffset());
              publishers.get(0).firstPublish.get(1, TimeUnit.MINUTES);
          }
          faultHandler.maybeRethrowFirstException();
      }
  
      private void loadTestSnapshot(
          MetadataLoader loader,
          long offset
      ) throws Exception {
          loader.handleLoadSnapshot(MockSnapshotReader.fromRecordLists(
                  new MetadataProvenance(offset, 100, 4000, true), List.of(
                          List.of(new ApiMessageAndVersion(new FeatureLevelRecord().
                                  setName(MetadataVersion.FEATURE_NAME).
                                  setFeatureLevel(MINIMUM_VERSION.featureLevel()), (short) 0)),
                          List.of(new ApiMessageAndVersion(new TopicRecord().
                                  setName("foo").
                                  setTopicId(Uuid.fromString("Uum7sfhHQP-obSvfywmNUA")), (short) 0))
                  )));
          loader.waitForAllEventsToBeHandled();
      }
  
      private void loadTestSnapshot2(
          MetadataLoader loader,
          long offset
      ) throws Exception {
          loader.handleLoadSnapshot(MockSnapshotReader.fromRecordLists(
                  new MetadataProvenance(offset, 100, 4000, true), List.of(
                          List.of(new ApiMessageAndVersion(new FeatureLevelRecord().
                                  setName(MetadataVersion.FEATURE_NAME).
                                  setFeatureLevel(MetadataVersion.latestProduction().featureLevel()), (short) 0)),
                          List.of(new ApiMessageAndVersion(new TopicRecord().
                                  setName("bar").
                                  setTopicId(Uuid.fromString("VcL2Mw-cT4aL6XV9VujzoQ")), (short) 0))
                  )));
          loader.waitForAllEventsToBeHandled();
      }
  
      /**
       * Test that loading a snapshot clears the previous state.
       */
      @Test
      public void testReloadSnapshot() throws Exception {
          MockFaultHandler faultHandler = new MockFaultHandler("testLastAppliedOffset");
          List<MockPublisher> publishers = List.of(new MockPublisher("a"));
          try (MetadataLoader loader = new MetadataLoader.Builder().
                  setFaultHandler(faultHandler).
                  setHighWaterMarkAccessor(() -> OptionalLong.of(0)).
                  build()) {
              loadTestSnapshot(loader, 100);
              loader.installPublishers(publishers).get();
              loader.waitForAllEventsToBeHandled();
              assertTrue(publishers.get(0).firstPublish.isDone());
              assertTrue(publishers.get(0).latestDelta.image().isEmpty());
              assertEquals(100L, publishers.get(0).latestImage.provenance().lastContainedOffset());
  
              loadTestSnapshot(loader, 200);
              assertEquals(200L, loader.lastAppliedOffset());
              assertEquals(MINIMUM_VERSION.featureLevel(),
                  loader.metrics().currentMetadataVersion().featureLevel());
              assertFalse(publishers.get(0).latestDelta.image().isEmpty());
  
              loadTestSnapshot2(loader, 400);
              assertEquals(400L, loader.lastAppliedOffset());
              assertEquals(MetadataVersion.latestProduction().featureLevel(),
                  loader.metrics().currentMetadataVersion().featureLevel());
  
              // Make sure the topic in the initial snapshot was overwritten by loading the new snapshot.
              assertFalse(publishers.get(0).latestImage.topics().topicsByName().containsKey("foo"));
              assertTrue(publishers.get(0).latestImage.topics().topicsByName().containsKey("bar"));
  
              loader.handleCommit(new MockBatchReader(500, List.of(
                  MockBatchReader.newBatch(500, 100, List.of(
                      new ApiMessageAndVersion(new FeatureLevelRecord().
                          setName(MetadataVersion.FEATURE_NAME).
                          setFeatureLevel(IBP_3_5_IV0.featureLevel()), (short) 0))))));
              loader.waitForAllEventsToBeHandled();
              assertEquals(IBP_3_5_IV0.featureLevel(),
                  loader.metrics().currentMetadataVersion().featureLevel());
          }
          faultHandler.maybeRethrowFirstException();
      }
  
      @ParameterizedTest
      @ValueSource(booleans = {true, false})
      public void testPublishTransaction(boolean abortTxn) throws Exception {
          MockFaultHandler faultHandler = new MockFaultHandler("testTransactions");
          MockPublisher publisher = new MockPublisher("testTransactions");
          List<MockPublisher> publishers = List.of(publisher);
          try (MetadataLoader loader = new MetadataLoader.Builder().
                  setFaultHandler(faultHandler).
                  setHighWaterMarkAccessor(() -> OptionalLong.of(0)).
                  build()) {
              loader.installPublishers(publishers).get();
              loader.handleCommit(
                  MockBatchReader.newSingleBatchReader(400, 50, List.of(
                      new ApiMessageAndVersion(new FeatureLevelRecord()
                          .setName(MetadataVersion.FEATURE_NAME)
                          .setFeatureLevel(MINIMUM_VERSION.featureLevel()), (short) 0)))
              );
              loader.waitForAllEventsToBeHandled();
  
              loader.handleCommit(
                  MockBatchReader.newSingleBatchReader(500, 100, List.of(
                      new ApiMessageAndVersion(new BeginTransactionRecord()
                          .setName("testTransactions"), (short) 0),
                      new ApiMessageAndVersion(new TopicRecord()
                          .setName("foo")
                          .setTopicId(Uuid.fromString("dMCqhcK4T5miGH5wEX7NsQ")), (short) 0)
              )));
              loader.waitForAllEventsToBeHandled();
              publisher.firstPublish.get(30, TimeUnit.SECONDS);
              assertNull(publisher.latestImage.topics().getTopic("foo"),
                  "Topic should not be visible since we started transaction");
  
              loader.handleCommit(
                  MockBatchReader.newSingleBatchReader(500, 100, List.of(
                      new ApiMessageAndVersion(new PartitionRecord()
                          .setTopicId(Uuid.fromString("dMCqhcK4T5miGH5wEX7NsQ"))
                          .setPartitionId(0), (short) 0),
                      new ApiMessageAndVersion(new PartitionRecord()
                          .setTopicId(Uuid.fromString("dMCqhcK4T5miGH5wEX7NsQ"))
                          .setPartitionId(1), (short) 0)
                  )));
              loader.waitForAllEventsToBeHandled();
              assertNull(publisher.latestImage.topics().getTopic("foo"),
                  "Topic should not be visible after subsequent batch");
  
              if (abortTxn) {
                  loader.handleCommit(
                      MockBatchReader.newSingleBatchReader(500, 100, List.of(
                          new ApiMessageAndVersion(new AbortTransactionRecord(), (short) 0)
                      )));
                  loader.waitForAllEventsToBeHandled();
  
                  assertNull(publisher.latestImage.topics().getTopic("foo"),
                      "Topic should not be visible since the transaction was aborted");
              } else {
                  loader.handleCommit(
                      MockBatchReader.newSingleBatchReader(500, 100, List.of(
                          new ApiMessageAndVersion(new EndTransactionRecord(), (short) 0)
                      )));
                  loader.waitForAllEventsToBeHandled();
  
                  assertNotNull(publisher.latestImage.topics().getTopic("foo"),
                      "Topic should be visible now that transaction has ended");
              }
          }
          faultHandler.maybeRethrowFirstException();
      }
  
      @Test
      public void testPublishTransactionWithinBatch() throws Exception {
          MockFaultHandler faultHandler = new MockFaultHandler("testPublishTransactionWithinBatch");
          MockPublisher publisher = new MockPublisher("testPublishTransactionWithinBatch");
          List<MockPublisher> publishers = List.of(publisher);
          try (MetadataLoader loader = new MetadataLoader.Builder().
                  setFaultHandler(faultHandler).
                  setHighWaterMarkAccessor(() -> OptionalLong.of(0)).
                  build()) {
              loader.installPublishers(publishers).get();
              loader.handleCommit(
                  MockBatchReader.newSingleBatchReader(400, 50, List.of(
                      new ApiMessageAndVersion(new FeatureLevelRecord().
                          setName(MetadataVersion.FEATURE_NAME).
                          setFeatureLevel(MINIMUM_VERSION.featureLevel()), (short) 0)))
              );
              loader.waitForAllEventsToBeHandled();
  
              loader.handleCommit(
                  MockBatchReader.newSingleBatchReader(500, 100, List.of(
                      new ApiMessageAndVersion(new BeginTransactionRecord()
                          .setName("txn-1"), (short) 0),
                      new ApiMessageAndVersion(new TopicRecord()
                          .setName("foo")
                          .setTopicId(Uuid.fromString("HQSM3ccPQISrHqYK_C8GpA")), (short) 0),
                      new ApiMessageAndVersion(new EndTransactionRecord(), (short) 0)
                  )));
              loader.waitForAllEventsToBeHandled();
  
              // After MetadataLoader is fixed to handle arbitrary transactions, we would expect "foo"
              // to be visible at this point.
              publisher.firstPublish.get(30, TimeUnit.SECONDS);
              assertNotNull(publisher.latestImage.topics().getTopic("foo"));
          }
          faultHandler.maybeRethrowFirstException();
      }
  
      @Test
      public void testSnapshotDuringTransaction() throws Exception {
          MockFaultHandler faultHandler = new MockFaultHandler("testSnapshotDuringTransaction");
          MockPublisher publisher = new MockPublisher("testSnapshotDuringTransaction");
          List<MockPublisher> publishers = List.of(publisher);
          try (MetadataLoader loader = new MetadataLoader.Builder().
                  setFaultHandler(faultHandler).
                  setHighWaterMarkAccessor(() -> OptionalLong.of(0)).
                  build()) {
              loader.installPublishers(publishers).get();
              loader.handleCommit(
                  MockBatchReader.newSingleBatchReader(400, 50, List.of(
                      new ApiMessageAndVersion(new FeatureLevelRecord().
                          setName(MetadataVersion.FEATURE_NAME).
                          setFeatureLevel(MINIMUM_VERSION.featureLevel()), (short) 0)))
              );
              loader.waitForAllEventsToBeHandled();
  
              loader.handleCommit(
                  MockBatchReader.newSingleBatchReader(500, 100, List.of(
                      new ApiMessageAndVersion(new BeginTransactionRecord()
                          .setName("txn-1"), (short) 0),
                      new ApiMessageAndVersion(new TopicRecord()
                          .setName("foo")
                          .setTopicId(Uuid.fromString("HQSM3ccPQISrHqYK_C8GpA")), (short) 0)
                  )));
              loader.waitForAllEventsToBeHandled();
              publisher.firstPublish.get(30, TimeUnit.SECONDS);
              assertNull(publisher.latestImage.topics().getTopic("foo"));
  
              // loading a snapshot discards any in-flight transaction
              loader.handleLoadSnapshot(MockSnapshotReader.fromRecordLists(
                  new MetadataProvenance(600, 101, 4000, true), List.of(
                      List.of(new ApiMessageAndVersion(new TopicRecord().
                          setName("foo").
                          setTopicId(Uuid.fromString("Uum7sfhHQP-obSvfywmNUA")), (short) 0))
                  )));
              loader.waitForAllEventsToBeHandled();
              assertEquals("Uum7sfhHQP-obSvfywmNUA",
                  publisher.latestImage.topics().getTopic("foo").id().toString());
          }
          faultHandler.maybeRethrowFirstException();
      }
  
      @Test
      public void testNoPublishEmptyImage() throws Exception {
          MockFaultHandler faultHandler = new MockFaultHandler("testNoPublishEmptyImage");
          List<MetadataImage> capturedImages = new ArrayList<>();
          CompletableFuture<Void> firstPublish = new CompletableFuture<>();
          MetadataPublisher capturingPublisher = new MetadataPublisher() {
              @Override
              public String name() {
                  return "testNoPublishEmptyImage";
              }
  
              @Override
              public void onMetadataUpdate(MetadataDelta delta, MetadataImage newImage, LoaderManifest manifest) {
                  capturedImages.add(newImage);
                  if (!firstPublish.isDone()) {
                      firstPublish.complete(null);
                  }
              }
          };
  
          try (MetadataLoader loader = new MetadataLoader.Builder().
                  setFaultHandler(faultHandler).
                  setHighWaterMarkAccessor(() -> OptionalLong.of(1)).
                  build()) {
              loader.installPublishers(List.of(capturingPublisher)).get();
              loader.handleCommit(
                  MockBatchReader.newSingleBatchReader(0, 1, List.of(
                      // Any record will work here
                      new ApiMessageAndVersion(new FeatureLevelRecord()
                          .setName(MetadataVersion.FEATURE_NAME)
                          .setFeatureLevel(MINIMUM_VERSION.featureLevel()), (short) 0)))
              );
              firstPublish.get(30, TimeUnit.SECONDS);
  
              assertFalse(capturedImages.isEmpty());
              capturedImages.forEach(metadataImage ->
                  assertFalse(metadataImage.isEmpty())
              );
  
          }
          faultHandler.maybeRethrowFirstException();
      }
  }
  