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.controller.metrics;
  
  import org.apache.kafka.common.Uuid;
  import org.apache.kafka.image.AclsImage;
  import org.apache.kafka.image.ClientQuotasImage;
  import org.apache.kafka.image.ClusterImage;
  import org.apache.kafka.image.ConfigurationsImage;
  import org.apache.kafka.image.DelegationTokenImage;
  import org.apache.kafka.image.FeaturesImage;
  import org.apache.kafka.image.MetadataDelta;
  import org.apache.kafka.image.MetadataImage;
  import org.apache.kafka.image.MetadataProvenance;
  import org.apache.kafka.image.ProducerIdsImage;
  import org.apache.kafka.image.ScramImage;
  import org.apache.kafka.image.TopicsImage;
  import org.apache.kafka.image.loader.LoaderManifest;
  import org.apache.kafka.image.loader.LogDeltaManifest;
  import org.apache.kafka.image.loader.SnapshotManifest;
  import org.apache.kafka.image.writer.ImageReWriter;
  import org.apache.kafka.image.writer.ImageWriterOptions;
  import org.apache.kafka.raft.LeaderAndEpoch;
  import org.apache.kafka.server.common.MetadataVersion;
  import org.apache.kafka.server.fault.MockFaultHandler;
  
  import org.junit.jupiter.api.Test;
  
  import java.util.Optional;
  
  import static org.apache.kafka.controller.metrics.ControllerMetricsTestUtils.FakePartitionRegistrationType.NON_PREFERRED_LEADER;
  import static org.apache.kafka.controller.metrics.ControllerMetricsTestUtils.FakePartitionRegistrationType.NORMAL;
  import static org.apache.kafka.controller.metrics.ControllerMetricsTestUtils.FakePartitionRegistrationType.OFFLINE;
  import static org.apache.kafka.controller.metrics.ControllerMetricsTestUtils.fakePartitionRegistration;
  import static org.apache.kafka.controller.metrics.ControllerMetricsTestUtils.fakeTopicImage;
  import static org.apache.kafka.controller.metrics.ControllerMetricsTestUtils.fakeTopicsImage;
  import static org.junit.jupiter.api.Assertions.assertEquals;
  
  public class ControllerMetadataMetricsPublisherTest {
      static class TestEnv implements AutoCloseable {
          MockFaultHandler faultHandler =
                  new MockFaultHandler("ControllerMetadataMetricsPublisher");
          ControllerMetadataMetrics metrics =
                  new ControllerMetadataMetrics(Optional.empty());
          ControllerMetadataMetricsPublisher publisher;
  
          TestEnv() {
              this(false);
          }
  
          TestEnv(boolean isDisklessTopics) {
              publisher = new ControllerMetadataMetricsPublisher(metrics, faultHandler, topicName -> isDisklessTopics);
          }
  
          @Override
          public void close() {
              publisher.close();
              faultHandler.maybeRethrowFirstException();
          }
      }
  
      @Test
      public void testMetricsBeforePublishing() {
          try (TestEnv env = new TestEnv()) {
              assertEquals(0, env.metrics.activeBrokerCount());
              assertEquals(0, env.metrics.globalTopicCount());
              assertEquals(0, env.metrics.globalPartitionCount());
              assertEquals(0, env.metrics.offlinePartitionCount());
              assertEquals(0, env.metrics.preferredReplicaImbalanceCount());
              assertEquals(0, env.metrics.metadataErrorCount());
          }
      }
  
      static MetadataImage fakeImageFromTopicsImage(TopicsImage topicsImage) {
          return new MetadataImage(
              MetadataProvenance.EMPTY,
              FeaturesImage.EMPTY,
              ClusterImage.EMPTY,
              topicsImage,
              ConfigurationsImage.EMPTY,
              ClientQuotasImage.EMPTY,
              ProducerIdsImage.EMPTY,
              AclsImage.EMPTY,
              ScramImage.EMPTY,
              DelegationTokenImage.EMPTY);
      }
  
      static final TopicsImage TOPICS_IMAGE1;
  
      static final MetadataImage IMAGE1;
  
      static {
          TOPICS_IMAGE1 = fakeTopicsImage(
              fakeTopicImage("foo",
                  Uuid.fromString("JKNp6fQaT-icHxh654ok-w"),
                      fakePartitionRegistration(NORMAL)),
              fakeTopicImage("bar",
                  Uuid.fromString("pEMSdUVWTXaFQUzLTznFSw"),
                      fakePartitionRegistration(NORMAL),
                      fakePartitionRegistration(NORMAL),
                      fakePartitionRegistration(NON_PREFERRED_LEADER)),
              fakeTopicImage("quux",
                      Uuid.fromString("zkUT4lyyRke6VIaTw6RQWg"),
                      fakePartitionRegistration(OFFLINE),
                      fakePartitionRegistration(OFFLINE),
                      fakePartitionRegistration(OFFLINE))
          );
          IMAGE1 = fakeImageFromTopicsImage(TOPICS_IMAGE1);
      }
  
      @Test
      public void testPublish() {
          try (TestEnv env = new TestEnv()) {
              assertEquals(0, env.metrics.activeBrokerCount());
              assertEquals(0, env.metrics.globalTopicCount());
              assertEquals(0, env.metrics.globalPartitionCount());
              assertEquals(0, env.metrics.offlinePartitionCount());
              assertEquals(0, env.metrics.preferredReplicaImbalanceCount());
              assertEquals(0, env.metrics.metadataErrorCount());
          }
      }
  
      static LoaderManifest fakeManifest(boolean isSnapshot) {
          if (isSnapshot) {
              return new SnapshotManifest(MetadataProvenance.EMPTY, 0);
          } else {
              return LogDeltaManifest.newBuilder()
                  .provenance(MetadataProvenance.EMPTY)
                  .leaderAndEpoch(LeaderAndEpoch.UNKNOWN)
                  .numBatches(0)
                  .elapsedNs(0)
                  .numBytes(0).build();
          }
      }
  
      @Test
      public void testLoadSnapshot() {
          try (TestEnv env = new TestEnv()) {
              MetadataDelta delta = new MetadataDelta(MetadataImage.EMPTY);
              ImageReWriter writer = new ImageReWriter(delta);
              IMAGE1.write(writer, new ImageWriterOptions.Builder(MetadataVersion.MINIMUM_VERSION).build());
              env.publisher.onMetadataUpdate(delta, IMAGE1, fakeManifest(true));
              assertEquals(0, env.metrics.activeBrokerCount());
              assertEquals(3, env.metrics.globalTopicCount());
              assertEquals(7, env.metrics.globalPartitionCount());
              assertEquals(3, env.metrics.offlinePartitionCount());
              assertEquals(4, env.metrics.preferredReplicaImbalanceCount());
              assertEquals(0, env.metrics.metadataErrorCount());
          }
      }
  
      @Test
      public void testLoadSnapshotWithDisklessTopics() {
          try (TestEnv env = new TestEnv(true)) {
              MetadataDelta delta = new MetadataDelta(MetadataImage.EMPTY);
              ImageReWriter writer = new ImageReWriter(delta);
              IMAGE1.write(writer, new ImageWriterOptions.Builder(MetadataVersion.MINIMUM_VERSION).build());
              env.publisher.onMetadataUpdate(delta, IMAGE1, fakeManifest(true));
              assertEquals(0, env.metrics.activeBrokerCount());
              assertEquals(3, env.metrics.globalTopicCount());
              assertEquals(7, env.metrics.globalPartitionCount());
              // Diskless topics are not counted in the metrics
              assertEquals(0, env.metrics.offlinePartitionCount());
              assertEquals(0, env.metrics.preferredReplicaImbalanceCount());
              assertEquals(0, env.metrics.metadataErrorCount());
          }
      }
  }
  