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.publisher;
  
  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.loader.LogDeltaManifest;
  import org.apache.kafka.metadata.RecordTestUtils;
  import org.apache.kafka.raft.LeaderAndEpoch;
  import org.apache.kafka.server.fault.FaultHandlerException;
  import org.apache.kafka.server.fault.MockFaultHandler;
  import org.apache.kafka.test.TestUtils;
  
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.api.Timeout;
  
  import java.util.ArrayList;
  import java.util.List;
  import java.util.concurrent.CopyOnWriteArrayList;
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicReference;
  
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertNotNull;
  
  
  @Timeout(value = 40)
  public class SnapshotGeneratorTest {
      static class MockEmitter implements SnapshotGenerator.Emitter {
          private final CountDownLatch latch = new CountDownLatch(1);
          private final List<MetadataImage> images = new CopyOnWriteArrayList<>();
          private volatile RuntimeException problem = null;
  
          MockEmitter setReady() {
              latch.countDown();
              return this;
          }
  
          synchronized MockEmitter setProblem(RuntimeException problem) {
              this.problem = problem;
              return this;
          }
  
          @Override
          public synchronized void maybeEmit(MetadataImage image) {
              RuntimeException currentProblem = problem;
              if (currentProblem != null) {
                  throw currentProblem;
              }
              try {
                  latch.await(30, TimeUnit.SECONDS);
              } catch (Throwable e) {
                  throw new RuntimeException(e);
              }
              images.add(image);
          }
  
          List<MetadataImage> images() {
              return new ArrayList<>(images);
          }
      }
  
      static LogDeltaManifest.Builder logDeltaManifestBuilder() {
          return LogDeltaManifest.newBuilder()
              .provenance(new MetadataProvenance(-1L, -1, -1L, true))
              .leaderAndEpoch(LeaderAndEpoch.UNKNOWN)
              .numBatches(1)
              .elapsedNs(100)
              .numBytes(100);
      }
  
      static LogDeltaManifest.Builder notBatchAlignedLogDeltaManifestBuilder() {
          return LogDeltaManifest.newBuilder()
                  .provenance(MetadataProvenance.EMPTY)
                  .leaderAndEpoch(LeaderAndEpoch.UNKNOWN)
                  .numBatches(1)
                  .elapsedNs(100)
                  .numBytes(100);
      }
  
      private static final MetadataDelta TEST_DELTA;
  
      static {
          TEST_DELTA = new MetadataDelta.Builder().
                  setImage(MetadataImage.EMPTY).
                  build();
          TEST_DELTA.replay(RecordTestUtils.testRecord(0).message());
      }
  
      private static final MetadataImage TEST_IMAGE = TEST_DELTA.apply(MetadataProvenance.EMPTY);
  
      @Test
      public void testCreateSnapshot() throws Exception {
          MockFaultHandler faultHandler = new MockFaultHandler("SnapshotGenerator");
          MockEmitter emitter = new MockEmitter();
          try (SnapshotGenerator generator = new SnapshotGenerator.Builder(emitter).
                  setFaultHandler(faultHandler).
                  setMaxBytesSinceLastSnapshot(200).
                  setMaxTimeSinceLastSnapshotNs(TimeUnit.DAYS.toNanos(10)).
                  build()) {
              // Publish a log delta batch. This one will not trigger a snapshot yet.
              generator.publishLogDelta(TEST_IMAGE, logDeltaManifestBuilder().build());
              // Publish a log delta batch. This will trigger a snapshot.
              generator.publishLogDelta(TEST_IMAGE, logDeltaManifestBuilder().build());
              // Publish a log delta batch. This one will be ignored because there are other images
              // queued for writing.
              generator.publishLogDelta(TEST_IMAGE, logDeltaManifestBuilder().numBytes(2000).build());
              assertEquals(List.of(), emitter.images());
              emitter.setReady();
          }
          assertEquals(List.of(TEST_IMAGE), emitter.images());
          faultHandler.maybeRethrowFirstException();
      }
  
      @Test
      public void testNoSnapshotCreatedWhenLastOffsetIsNotBatchAligned() throws Exception {
          MockFaultHandler faultHandler = new MockFaultHandler("SnapshotGenerator");
          MockEmitter emitter = new MockEmitter();
          try (SnapshotGenerator generator = new SnapshotGenerator.Builder(emitter).
                  setFaultHandler(faultHandler).
                  setMaxBytesSinceLastSnapshot(200).
                  setMaxTimeSinceLastSnapshotNs(TimeUnit.DAYS.toNanos(10)).
                  build()) {
              // None of these log delta batches should trigger a snapshot since their offset is not batch aligned.
              generator.publishLogDelta(TEST_IMAGE, notBatchAlignedLogDeltaManifestBuilder().build());
              generator.publishLogDelta(TEST_IMAGE, notBatchAlignedLogDeltaManifestBuilder().build());
              generator.publishLogDelta(TEST_IMAGE, notBatchAlignedLogDeltaManifestBuilder().build());
              assertEquals(List.of(), emitter.images());
              emitter.setReady();
          }
          assertEquals(List.of(), emitter.images());
          faultHandler.maybeRethrowFirstException();
      }
  
      @Test
      public void testSnapshotsCreatedAgainWhenLastOffsetIsAligned() throws Exception {
          MockFaultHandler faultHandler = new MockFaultHandler("SnapshotGenerator");
          MockEmitter emitter = new MockEmitter();
          MetadataImage batchAlignedImage = TEST_DELTA.apply(
                  new MetadataProvenance(-1L, -1, -1L, true));
          try (SnapshotGenerator generator = new SnapshotGenerator.Builder(emitter).
                  setFaultHandler(faultHandler).
                  setMaxBytesSinceLastSnapshot(100).
                  setMaxTimeSinceLastSnapshotNs(TimeUnit.DAYS.toNanos(10)).
                  build()) {
              // These should not be published despite meeting the max bytes threshold since they are not batch aligned.
              generator.publishLogDelta(TEST_IMAGE, notBatchAlignedLogDeltaManifestBuilder().build());
              generator.publishLogDelta(TEST_IMAGE, notBatchAlignedLogDeltaManifestBuilder().build());
              // This snapshot should get published since it is batch aligned.
              generator.publishLogDelta(batchAlignedImage, logDeltaManifestBuilder().build());
              assertEquals(List.of(), emitter.images());
              emitter.setReady();
          }
          assertEquals(List.of(batchAlignedImage), emitter.images());
          faultHandler.maybeRethrowFirstException();
      }
  
      @Test
      public void testSnapshotsDisabled() throws Exception {
          MockFaultHandler faultHandler = new MockFaultHandler("SnapshotGenerator");
          MockEmitter emitter = new MockEmitter().setReady();
          AtomicReference<String> disabledReason = new AtomicReference<>();
          try (SnapshotGenerator generator = new SnapshotGenerator.Builder(emitter).
                  setFaultHandler(faultHandler).
                  setMaxBytesSinceLastSnapshot(1).
                  setMaxTimeSinceLastSnapshotNs(0).
                  setDisabledReason(disabledReason).
                  build()) {
              disabledReason.compareAndSet(null, "we are testing disable()");
              // No snapshots are generated because snapshots are disabled.
              generator.publishLogDelta(TEST_IMAGE, logDeltaManifestBuilder().build());
          }
          assertEquals(List.of(), emitter.images());
          faultHandler.maybeRethrowFirstException();
      }
  
      @Test
      public void testTimeBasedSnapshots() throws Exception {
          MockFaultHandler faultHandler = new MockFaultHandler("SnapshotGenerator");
          MockEmitter emitter = new MockEmitter().setReady();
          MockTime mockTime = new MockTime();
          try (SnapshotGenerator generator = new SnapshotGenerator.Builder(emitter).
                  setTime(mockTime).
                  setFaultHandler(faultHandler).
                  setMaxBytesSinceLastSnapshot(200).
                  setMaxTimeSinceLastSnapshotNs(TimeUnit.MINUTES.toNanos(30)).
                  build()) {
              // This image isn't published yet.
              generator.publishLogDelta(TEST_IMAGE, logDeltaManifestBuilder().numBytes(50).build());
              assertEquals(List.of(), emitter.images());
              mockTime.sleep(TimeUnit.MINUTES.toNanos(40));
              // Next image is published because of the time delay.
              generator.publishLogDelta(TEST_IMAGE, logDeltaManifestBuilder().numBytes(50).build());
              TestUtils.waitForCondition(() -> emitter.images().size() == 1, "images.size == 1");
              // bytesSinceLastSnapshot was reset to 0 by the previous snapshot,
              // so this does not trigger a new snapshot.
              generator.publishLogDelta(TEST_IMAGE, logDeltaManifestBuilder().numBytes(150).build());
          }
          assertEquals(List.of(TEST_IMAGE), emitter.images());
          faultHandler.maybeRethrowFirstException();
      }
  
      @Test
      public void testEmitterProblem() throws Exception {
          MockFaultHandler faultHandler = new MockFaultHandler("SnapshotGenerator");
          MockEmitter emitter = new MockEmitter().setProblem(new RuntimeException("oops"));
          try (SnapshotGenerator generator = new SnapshotGenerator.Builder(emitter).
                  setFaultHandler(faultHandler).
                  setMaxBytesSinceLastSnapshot(200).
                  build()) {
              for (int i = 0; i < 2; i++) {
                  generator.publishLogDelta(TEST_IMAGE,
                      logDeltaManifestBuilder().elapsedNs(10000).numBytes(50000).build());
              }
          }
          assertEquals(List.of(), emitter.images());
          assertNotNull(faultHandler.firstException());
          assertEquals(FaultHandlerException.class, faultHandler.firstException().getClass());
          assertEquals("SnapshotGenerator: KRaft snapshot file generation error: oops",
                  faultHandler.firstException().getMessage());
      }
  }
  