Java 类com.codahale.metrics.health.HealthCheckRegistry 实例源码

private ServletContextHandler getMetricsHandler() {
    MetricRegistry registry = Metrics.getRegistry();
    HealthCheckRegistry healthCheckRegistry = Metrics.getHealthCheckRegistry();
    healthCheckRegistry.register("rotation", new Rotation(configuration.getRotationStatusFilePath()));

    registry.registerAll(new GarbageCollectorMetricSet());
    registry.registerAll(new MemoryUsageGaugeSet());
    registry.registerAll(new ThreadStatesGaugeSet());
    registry.registerAll(new JvmAttributeGaugeSet());

    ServletContextHandler servletContextHandler = new ServletContextHandler();
    servletContextHandler.setContextPath("/__metrics");
    servletContextHandler.setAttribute(MetricsServlet.class.getCanonicalName() + ".registry", registry);
    servletContextHandler.setAttribute(HealthCheckServlet.class.getCanonicalName() + ".registry", healthCheckRegistry);
    servletContextHandler.addServlet(new ServletHolder(new AdminServlet()), "/*");

    return servletContextHandler;
}

public static Handler healthCheckHandler(
    HealthCheckRegistry healthCheckRegistry, ObjectMapper mapper) {
  Preconditions.checkState(healthCheckRegistry != null);
  Preconditions.checkState(mapper != null);

  SortedMap<String, HealthCheck.Result> healthChecks = healthCheckRegistry.runHealthChecks();

  return xrpcRequest ->
      Recipes.newResponseOk(
          xrpcRequest
              .getAlloc()
              .directBuffer()
              .writeBytes(
                  mapper.writerWithDefaultPrettyPrinter().writeValueAsBytes(healthChecks)),
          Recipes.ContentType.Application_Json);
}

public static HikariDataSource getDataSourceFromConfig(
    Config conf
    , MetricRegistry metricRegistry
    , HealthCheckRegistry healthCheckRegistry) {

    HikariConfig jdbcConfig = new HikariConfig();
    jdbcConfig.setPoolName(conf.getString("poolName"));
    jdbcConfig.setMaximumPoolSize(conf.getInt("maximumPoolSize"));
    jdbcConfig.setMinimumIdle(conf.getInt("minimumIdle"));
    jdbcConfig.setJdbcUrl(conf.getString("jdbcUrl"));
    jdbcConfig.setUsername(conf.getString("username"));
    jdbcConfig.setPassword(conf.getString("password"));

    jdbcConfig.addDataSourceProperty("cachePrepStmts", conf.getBoolean("cachePrepStmts"));
    jdbcConfig.addDataSourceProperty("prepStmtCacheSize", conf.getInt("prepStmtCacheSize"));
    jdbcConfig.addDataSourceProperty("prepStmtCacheSqlLimit", conf.getInt("prepStmtCacheSqlLimit"));
    jdbcConfig.addDataSourceProperty("useServerPrepStmts", conf.getBoolean("useServerPrepStmts"));

    // Add HealthCheck
    jdbcConfig.setHealthCheckRegistry(healthCheckRegistry);

    // Add Metrics
    jdbcConfig.setMetricRegistry(metricRegistry);
    return new HikariDataSource(jdbcConfig);
}

public static HikariDataSource getDataSourceFromConfig(
    Config conf
    , MetricRegistry metricRegistry
    , HealthCheckRegistry healthCheckRegistry) {

    HikariConfig jdbcConfig = new HikariConfig();
    jdbcConfig.setPoolName(conf.getString("poolName"));
    jdbcConfig.setMaximumPoolSize(conf.getInt("maximumPoolSize"));
    jdbcConfig.setMinimumIdle(conf.getInt("minimumIdle"));
    jdbcConfig.setJdbcUrl(conf.getString("jdbcUrl"));
    jdbcConfig.setUsername(conf.getString("username"));
    jdbcConfig.setPassword(conf.getString("password"));

    jdbcConfig.addDataSourceProperty("cachePrepStmts", conf.getBoolean("cachePrepStmts"));
    jdbcConfig.addDataSourceProperty("prepStmtCacheSize", conf.getInt("prepStmtCacheSize"));
    jdbcConfig.addDataSourceProperty("prepStmtCacheSqlLimit", conf.getInt("prepStmtCacheSqlLimit"));
    jdbcConfig.addDataSourceProperty("useServerPrepStmts", conf.getBoolean("useServerPrepStmts"));

    // Add HealthCheck
    jdbcConfig.setHealthCheckRegistry(healthCheckRegistry);

    // Add Metrics
    jdbcConfig.setMetricRegistry(metricRegistry);
    return new HikariDataSource(jdbcConfig);
}

@Override
public void loadDimensionDictionary(DimensionDictionary dimensions) {
    HealthCheckRegistry healthCheckRegistry = HealthCheckRegistryFactory.getRegistry();
    for (DimensionConfig dimension : configSource) {
        dimensions.add(new KeyValueStoreDimension(dimension));

        //Register KeyValueStore Health Check
        healthCheckRegistry.register(
            dimension.getApiName() + " keyValueStore check",
            new KeyValueStoreHealthCheck(dimension.getKeyValueStore())
        );

        //Register SearchProvider Health Check
        healthCheckRegistry.register(
            dimension.getApiName() + " searchProvider check",
            new SearchProviderHealthCheck(dimension.getSearchProvider())
        );
    }
}

@Override
public void activateService() {
    metricRegistry = new MetricRegistry();
    healthCheckRegistry = new HealthCheckRegistry();
    CodahaleMetricsDeclaration declaration = descriptor.metaInfo( CodahaleMetricsDeclaration.class );
    prefix = declaration.prefix() != null ? declaration.prefix() : app.name();
    fqcn = declaration.fqcn();
    if( declaration.jmx() )
    {
        JmxReporter jmxReporter = JmxReporter.forRegistry( metricRegistry ).build();
        jmxReporter.start();
        reporters.add( jmxReporter );
    }
    for( Function<MetricRegistry, Reporter> reporterFactory : declaration.reportersFactories())
    {
        reporters.add( reporterFactory.apply( metricRegistry ) );
    }
}

public static void main(String[] args) {
    SpringApplication.run(MonetaSpringBootApplication.class, args);

    // Find and read application configuration
    MonetaConfiguration config = new MonetaConfiguration();

    // Install all health checks
    HealthCheckRegistry registry = new HealthCheckRegistry();
    for (String checkName : MonetaEnvironment.getConfiguration()
            .getHealthChecks()
            .keySet()) {
        registry.register(checkName, MonetaEnvironment.getConfiguration()
                .getHealthChecks()
                .get(checkName));
    }
    ActuatorHealthIndicator.setHealthCheckRegistry(registry);

    // Install metrics and JMX
    MetricRegistry metricRegistry = new MetricRegistry();
    final JmxReporter jmxReporter = JmxReporter.forRegistry(metricRegistry)
            .build();
    jmxReporter.start();
}

/**
 * Shows that that when the 2 healthchecks with the same name are registered, then the latest one will replace the previous one.
 */
@Test
public void testRegisteringSameHealthcheckMultipleTimes() {
    final HealthCheckRegistry registry = new HealthCheckRegistry();
    final String healthCheckName = "health-check";
    final HealthCheck healthCheck1 = new HealthCheck() {

        @Override
        protected HealthCheck.Result check() throws Exception {
            return HealthCheck.Result.healthy();
        }
    };
    final HealthCheck healthCheck2 = new HealthCheck() {

        @Override
        protected HealthCheck.Result check() throws Exception {
            return HealthCheck.Result.healthy();
        }
    };
    registry.register(healthCheckName, healthCheck1);
    registry.register(healthCheckName, healthCheck2);

    assertThat(registry.getNames().size(), is(1));
}

@Test
public void shouldRegisterHealthCheck() throws Exception {
    //GIVEN
    AmazonSQS sqs = mock(AmazonSQS.class);

    SqsBundle spiedBundle = spy(bundle);
    doReturn(sqs).when(spiedBundle).getAmazonSQS();

    LifecycleEnvironment lifecycle = mock(LifecycleEnvironment.class);
    doNothing().when(lifecycle).manage((Managed) anyObject());
    when(environment.lifecycle()).thenReturn(lifecycle);

    HealthCheckRegistry healthChecks = mock(HealthCheckRegistry.class);
    doNothing().when(healthChecks).register(anyObject(), anyObject());
    when(environment.healthChecks()).thenReturn(healthChecks);

    //WHEN
    spiedBundle.run(configurationHolder, environment);

    //THEN
    verify(healthChecks, times(1)).register(eq("SqsBundle"), any(SqsBundleHealthCheck.class));
}

@Override
public void run() {
    StormClusterState clusterState = workerData.getZkCluster();
    String topologyId = workerData.getTopologyId();

    Map<Integer, HealthCheckRegistry> taskHealthCheckMap = JStormHealthCheck.getTaskhealthcheckmap();
    int cnt = 0;
    for (Map.Entry<Integer, HealthCheckRegistry> entry : taskHealthCheckMap.entrySet()) {
        Integer taskId = entry.getKey();
        Map<String, Result> results = entry.getValue().runHealthChecks();

        for (Map.Entry<String, Result> result : results.entrySet()) {
            if (!result.getValue().isHealthy()) {
                try {
                    clusterState.report_task_error(topologyId, taskId, result.getValue().getMessage(), null);
                    cnt++;
                } catch (Exception e) {
                    LOG.error("Failed to update health data in ZK for topo-{} task-{}.", topologyId, taskId, e);
                }
            }
        }
    }
    LOG.info("Successfully updated {} health data to ZK for topology:{}", cnt, topologyId);
}

@Override
public void onActivate( Application application )
    throws ActivationException
{
    application.plugin( JSON.class ).mapper()
        .registerModule( new MetricsModule( SECONDS, MILLISECONDS, true ) )
        .registerModule( new HealthCheckModule() );
    MetricRegistry metrics = new MetricRegistry();
    HealthCheckRegistry healthChecks = new HealthCheckRegistry();

    registerMetrics( application, metrics );
    registerMetricsReporters( application, metrics );
    registerHealthChecks( application, healthChecks );

    api = new Metrics( metrics, healthChecks );
}

public HealthChecksThreadPool(HealthCheckRegistry healthCheckRegistry) {
    super(MAX_THREAD_POOL_SIZE, MAX_THREAD_POOL_SIZE, 5L, TimeUnit.SECONDS,
            new LinkedBlockingQueue<Runnable>(), //
            new ExceptionCatchingThreadFactory(new DaemonThreadFactory(new ThreadFactory() {
                private final AtomicInteger number = new AtomicInteger();

                public Thread newThread(Runnable r) {
                    return new Thread(r, "Metrics-HealthChecks-" + number.incrementAndGet());
                }
            })), new MetricsRejectedExecutionHandler(healthCheckRegistry));
    this.allowCoreThreadTimeOut(true); // allow stopping all threads if idle
    this.healthCheckRegistry = healthCheckRegistry;
    LOGGER.log(Level.FINE,
            "Created thread pool with a max of {0} threads (plus {1} in queue) for {2} health checks",
            new Object[] { getMaximumPoolSize(),
                           Math.max(0, healthCheckRegistry.getNames().size() + 1 - getMaximumPoolSize()),
                           healthCheckRegistry.getNames().size() });
}

/**
 * Drop the oldest health check in executor queue and cancel it
 */
static void dropOldestInQueue(ThreadPoolExecutor executor, HealthCheckRegistry healthCheckRegistry) {
    // capture the state up front
    Object[] params = {
            executor, executor.getQueue().size(), healthCheckRegistry.getNames(),
            healthCheckRegistry.getNames().size(), Arrays.asList(executor.getQueue().toArray())
    };

    Runnable discarded = executor.getQueue().poll();
    // if there are two workers taking jobs concurrently, the queue will be empty by the time we get here
    if (discarded != null) {
        LOGGER.log(Level.WARNING,
                "Too many health check executions queued, dropping oldest one. This may mean some health checks "
                        + "are taking too long to execute:"
                        + " {0}, queue size={1}, health checks={2} ({3}) {4}",
                params);
        cancelQueuedHealthCheck(discarded);
    }
}

@Override
public void handle(final Request req, final Response rsp) throws Throwable {
  HealthCheckRegistry registry = req.require(HealthCheckRegistry.class);
  SortedMap<String, Result> checks = req.param("name").toOptional().map(name -> {
    SortedMap<String, Result> set = ImmutableSortedMap.of(name, registry.runHealthCheck(name));
    return set;
  }).orElseGet(() -> registry.runHealthChecks());

  final Status status;
  if (checks.isEmpty()) {
    status = Status.NOT_IMPLEMENTED;
  } else {
    status = checks.values().stream()
        .filter(it -> !it.isHealthy())
        .findFirst()
        .map(it -> Status.SERVER_ERROR)
        .orElse(Status.OK);
  }
  rsp.status(status)
      .header("Cache-Control", "must-revalidate,no-cache,no-store")
      .send(checks);
}

public StatisticsController(final DelegatingAuditTrailManager auditTrailManager,
                            final CentralAuthenticationService centralAuthenticationService,
                            final MetricRegistry metricsRegistry,
                            final HealthCheckRegistry healthCheckRegistry,
                            final CasConfigurationProperties casProperties) {
    super("casstats", "/stats", casProperties.getMonitor().getEndpoints().getStatistics(), casProperties);
    this.auditTrailManager = auditTrailManager;
    this.centralAuthenticationService = centralAuthenticationService;
    this.metricsRegistry = metricsRegistry;
    this.healthCheckRegistry = healthCheckRegistry;
    this.casProperties = casProperties;
}

@Inject
public HealthCheckResource(Set<HealthCheck> healthChecks) {
  this.healthCheckRegistry = new HealthCheckRegistry();
  for (HealthCheck healthCheck : healthChecks) {
    healthCheckRegistry.register(healthCheck.getClass().getSimpleName(), healthCheck);
  }
}

public static String format(HealthCheckRegistry healthCheckRegistry) {
    return StringTableFormatter.format(
        100,
        "Healthcheck status:",
        extractHealthCheckData(healthCheckRegistry)
    );
}

@Test
public void test() throws SQLException {
    Config config = ConfigFactory.empty()
        .withValue("poolName", ConfigValueFactory.fromAnyRef("test pool"))
        .withValue("jdbcUrl", ConfigValueFactory.fromAnyRef("jdbc:hsqldb:mem:testdb"))
        .withValue("maximumPoolSize", ConfigValueFactory.fromAnyRef(10))
        .withValue("minimumIdle", ConfigValueFactory.fromAnyRef(2))
        .withValue("username", ConfigValueFactory.fromAnyRef("SA"))
        .withValue("password", ConfigValueFactory.fromAnyRef(""))
        .withValue("cachePrepStmts", ConfigValueFactory.fromAnyRef(true))
        .withValue("prepStmtCacheSize", ConfigValueFactory.fromAnyRef(256))
        .withValue("prepStmtCacheSqlLimit", ConfigValueFactory.fromAnyRef(2048))
        .withValue("useServerPrepStmts", ConfigValueFactory.fromAnyRef(true))
    ;
    MetricRegistry metricRegistry = new MetricRegistry();
    HealthCheckRegistry healthCheckRegistry = new HealthCheckRegistry();
    try (HikariDataSource ds = ConnectionPool.getDataSourceFromConfig(config, metricRegistry, healthCheckRegistry)) {
        assertTrue(ds.getPoolName().equals("test pool"));
        assertTrue(ds.getMaximumPoolSize() == 10);
        assertTrue(ds.getMinimumIdle() == 2);
        assertTrue(ds.getUsername().equals("SA"));
        assertTrue(ds.getPassword().equals(""));
        Properties dsp = ds.getDataSourceProperties();
        assertTrue(((boolean)dsp.get("cachePrepStmts")) == true);
        assertTrue(((int)dsp.get("prepStmtCacheSize")) == 256);
        assertTrue(((int)dsp.get("prepStmtCacheSqlLimit")) == 2048);
        assertTrue(((boolean)dsp.get("useServerPrepStmts")) == true);

        // Using identity equals on purpose
        assertTrue(ds.getHealthCheckRegistry() == healthCheckRegistry);
        assertTrue(ds.getMetricRegistry() == metricRegistry);

        try (Connection conn = ds.getConnection()) {
             assertTrue(conn.isValid(1000));
        }
    }
}

@Autowired
public MetricsPresenter(CommandService commandService, MetricRegistry metricRegistry,
                        HealthCheckRegistry healthCheckRegistry) {
    this.commandService = commandService;
    this.metricRegistry = metricRegistry;
    this.healthCheckRegistry = healthCheckRegistry;
}

@Autowired
public MetricsService(MetricRegistry metricRegistry, HealthCheckRegistry healthCheckRegistry,
                      GameServerService gameServerService, GameServerRepository gameServerRepository, DiscordService discordService, IncidentService incidentService) {
    this.metricRegistry = metricRegistry;
    this.healthCheckRegistry = healthCheckRegistry;
    this.gameServerService = gameServerService;
    this.gameServerRepository = gameServerRepository;
    this.discordService = discordService;
    this.incidentService = incidentService;
}

@SuppressWarnings("unchecked")
@Override
protected void configure() {
    bind(environment).to(Environment.class);
    bind(environment.healthChecks()).to(HealthCheckRegistry.class);
    bind(environment.lifecycle()).to(LifecycleEnvironment.class);
    bind(environment.metrics()).to(MetricRegistry.class);
    bind(environment.getValidator()).to(Validator.class);
    bind(configuration).to(bootstrap.getApplication().getConfigurationClass()).to(Configuration.class);
    bind(environment.getObjectMapper()).to(ObjectMapper.class);
    bind(bootstrap.getApplication()).to((Class) bootstrap.getApplication().getClass()).to(Application.class);
}

@Test
public void testHealthChecks() throws Exception
{
   MetricRegistry metricRegistry = new MetricRegistry();
   HealthCheckRegistry healthRegistry = new HealthCheckRegistry();

   HikariConfig config = new HikariConfig();
   config.setMaximumPoolSize(10);
   config.setMetricRegistry(metricRegistry);
   config.setHealthCheckRegistry(healthRegistry);
   config.setPoolName("test");
   config.setDataSourceClassName("com.zaxxer.hikari.mocks.StubDataSource");
   config.addHealthCheckProperty("connectivityCheckTimeoutMs", "1000");
   config.addHealthCheckProperty("expected99thPercentileMs", "10");

   HikariDataSource ds = new HikariDataSource(config);
   try {
      UtilityElf.quietlySleep(TimeUnit.SECONDS.toMillis(2));

      Connection connection = ds.getConnection();
      connection.close();

      connection = ds.getConnection();
      connection.close();

      SortedMap<String, Result> healthChecks = healthRegistry.runHealthChecks();

      Result connectivityResult = healthChecks.get("test.pool.ConnectivityCheck");
      Assert.assertTrue(connectivityResult.isHealthy());

      Result slaResult = healthChecks.get("test.pool.Connection99Percent");
      Assert.assertTrue(slaResult.isHealthy());
   }
   finally {
      ds.close();
   }
}

@Test
public void testSetters1() throws Exception
{
   HikariDataSource ds = new HikariDataSource();
   ds.setMaximumPoolSize(1);
   ds.setDataSourceClassName("com.zaxxer.hikari.mocks.StubDataSource");

   MetricRegistry metricRegistry = new MetricRegistry();
   HealthCheckRegistry healthRegistry = new HealthCheckRegistry();

   try {
      Connection connection = ds.getConnection();
      connection.close();

      // After the pool as started, we can only set them once...
      ds.setMetricRegistry(metricRegistry);
      ds.setHealthCheckRegistry(healthRegistry);

      // and never again...
      ds.setMetricRegistry(metricRegistry);
      Assert.fail("Should not have been allowed to set registry after pool started");
   }
   catch (IllegalStateException ise) {
      // pass
      try {
         ds.setHealthCheckRegistry(healthRegistry);
         Assert.fail("Should not have been allowed to set registry after pool started");
      }
      catch (IllegalStateException ise2) {
         // pass
      }
   }
   finally {
      ds.close();
   }
}

@Test
public void testSetters2() throws Exception
{
   HikariDataSource ds = new HikariDataSource();
   ds.setMaximumPoolSize(1);
   ds.setDataSourceClassName("com.zaxxer.hikari.mocks.StubDataSource");

   MetricRegistry metricRegistry = new MetricRegistry();
   HealthCheckRegistry healthRegistry = new HealthCheckRegistry();

   ds.setMetricRegistry(metricRegistry);
   ds.setHealthCheckRegistry(healthRegistry);

   // before the pool is started, we can set it any number of times...
   ds.setMetricRegistry(metricRegistry);
   ds.setHealthCheckRegistry(healthRegistry);

   try {
      Connection connection = ds.getConnection();
      connection.close();

      // after the pool is started, we cannot set it any more
      ds.setMetricRegistry(metricRegistry);
      Assert.fail("Should not have been allowed to set registry after pool started");
   }
   catch (IllegalStateException ise) {

   }
   finally {
      ds.close();
   }
}

@Bean
@Autowired
public ServletRegistrationBean servletHealthRegistryBean(HealthCheckRegistry healthCheckRegistry) {
    HealthCheckServlet hc = new HealthCheckServlet(healthCheckRegistry);
    ServletRegistrationBean srb = new ServletRegistrationBean(hc, "/health/*");
    srb.setLoadOnStartup(2);
    return srb;
}

private InstrumentedAnnotations(
        MetricRegistry metricRegistry,
        HealthCheckRegistry healthCheckRegistry,
        Predicate<Throwable> exceptionFilter
) {
    this.metricRegistry = metricRegistry;
    this.healthCheckRegistry = healthCheckRegistry;
    this.exceptionFilter = exceptionFilter;
}

public Instrumentor(
        MetricRegistry metricRegistry,
        HealthCheckRegistry healthCheckRegistry,
        Predicate<Throwable> exceptionFilter) {
    this.metricRegistry = metricRegistry;
    this.healthCheckRegistry = healthCheckRegistry;
    this.exceptionFilter = exceptionFilter;
}

@Test
public void test() throws SQLException {
    Config config = ConfigFactory.empty()
        .withValue("poolName", ConfigValueFactory.fromAnyRef("test pool"))
        .withValue("jdbcUrl", ConfigValueFactory.fromAnyRef("jdbc:hsqldb:mem:testdb"))
        .withValue("maximumPoolSize", ConfigValueFactory.fromAnyRef(10))
        .withValue("minimumIdle", ConfigValueFactory.fromAnyRef(2))
        .withValue("username", ConfigValueFactory.fromAnyRef("SA"))
        .withValue("password", ConfigValueFactory.fromAnyRef(""))
        .withValue("cachePrepStmts", ConfigValueFactory.fromAnyRef(true))
        .withValue("prepStmtCacheSize", ConfigValueFactory.fromAnyRef(256))
        .withValue("prepStmtCacheSqlLimit", ConfigValueFactory.fromAnyRef(2048))
        .withValue("useServerPrepStmts", ConfigValueFactory.fromAnyRef(true))
    ;
    MetricRegistry metricRegistry = new MetricRegistry();
    HealthCheckRegistry healthCheckRegistry = new HealthCheckRegistry();
    try (HikariDataSource ds = ConnectionPool.getDataSourceFromConfig(config, metricRegistry, healthCheckRegistry)) {
        assertTrue(ds.getPoolName().equals("test pool"));
        assertTrue(ds.getMaximumPoolSize() == 10);
        assertTrue(ds.getMinimumIdle() == 2);
        assertTrue(ds.getUsername().equals("SA"));
        assertTrue(ds.getPassword().equals(""));
        Properties dsp = ds.getDataSourceProperties();
        assertTrue(((boolean)dsp.get("cachePrepStmts")) == true);
        assertTrue(((int)dsp.get("prepStmtCacheSize")) == 256);
        assertTrue(((int)dsp.get("prepStmtCacheSqlLimit")) == 2048);
        assertTrue(((boolean)dsp.get("useServerPrepStmts")) == true);

        // Using identity equals on purpose
        assertTrue(ds.getHealthCheckRegistry() == healthCheckRegistry);
        assertTrue(ds.getMetricRegistry() == metricRegistry);

        try (Connection conn = ds.getConnection()) {
             assertTrue(conn.isValid(1000));
        }
    }
}

/**
 * Register required core health checks.
 *
 * @param healthCheckRegistry  The health check registry
 * @param dimensionDictionary  The container with all dimensions
 */
protected final void setupHealthChecks(
        HealthCheckRegistry healthCheckRegistry,
        DimensionDictionary dimensionDictionary
) {
    // Health checks are registered here since they should be configured at startup.
    HealthCheck allDimsLoaded = new AllDimensionsLoadedHealthCheck(dimensionDictionary);
    healthCheckRegistry.register(HEALTH_CHECK_NAME_DIMENSION, allDimsLoaded);
    healthCheckRegistry.register(HEALTH_CHECK_VERSION, new VersionHealthCheck());
}

/**
 * Schedule a datasource metadata loader and register its health check.
 *
 * @param healthCheckRegistry  The health check registry to register partial data health checks.
 * @param dataSourceMetadataLoader  The datasource metadata loader to use.
 */
protected final void setupDataSourceMetaData(
        HealthCheckRegistry healthCheckRegistry,
        DataSourceMetadataLoadTask dataSourceMetadataLoader
) {
    scheduleLoader(dataSourceMetadataLoader);

    // Register Segment metadata loader health check
    HealthCheck dataSourceMetadataLoaderHealthCheck = new DataSourceMetadataLoaderHealthCheck(
            dataSourceMetadataLoader,
            SEG_LOADER_HC_LAST_RUN_PERIOD_MILLIS
    );
    healthCheckRegistry.register(HEALTH_CHECK_NAME_DATASOURCE_METADATA, dataSourceMetadataLoaderHealthCheck);
}

/**
 * Schedule DimensionValueLoadTask and register its health check.
 *
 * @param healthCheckRegistry  The health check registry to register Dimension lookup health checks
 * @param dataDimensionLoader  The DruidDimensionLoader used for monitoring and health checks
 */
protected final void setupDruidDimensionsLoader(
        HealthCheckRegistry healthCheckRegistry,
        DimensionValueLoadTask dataDimensionLoader
) {
    scheduleLoader(dataDimensionLoader);

    // Register DimensionValueLoadTask health check
    HealthCheck druidDimensionsLoaderHealthCheck = new DruidDimensionsLoaderHealthCheck(
            dataDimensionLoader,
            DRUID_DIM_LOADER_HC_LAST_RUN_PERIOD_MILLIS
    );
    healthCheckRegistry.register(HEALTH_CHECK_NAME_DRUID_DIM_LOADER, druidDimensionsLoaderHealthCheck);
}

public ServletReporter(MetricRegistry metricRegistry,
                       HealthCheckRegistry healthCheckRegistry,
                       int port) {
    this.metricRegistry = metricRegistry;
    this.healthCheckRegistry = healthCheckRegistry;
    this.port = port;
    this.jettyServer = new Server(port);
}

@Before
public void setUp() throws Exception {
    mockConfiguration = mock(DummyConfiguration.class);
    mockEnvironment = mock(Environment.class);
    mockLifecycle = mock(LifecycleEnvironment.class);
    mockJersey = mock(JerseyEnvironment.class);
    mockHealthchecks = mock(HealthCheckRegistry.class);
    when(mockEnvironment.lifecycle()).thenReturn(mockLifecycle);
    when(mockEnvironment.healthChecks()).thenReturn(mockHealthchecks);
}

@Override
   public void run(DuctileDBServerConfiguration configuration, Environment environment) throws Exception {
MetricRegistry metrics = environment.metrics();
Metrics.initialize(metrics);

HealthCheckRegistry healthChecks = environment.healthChecks();

JerseyEnvironment jersey = environment.jersey();
jersey.setUrlPattern("/rest");

   }

@Before
public void setUp() throws Exception {
  this.registry = mock(HealthCheckRegistry.class);
  this.environment = mock(Environment.class);
  this.keywhizConfig = mock(KeywhizConfig.class);

  when(environment.healthChecks()).thenReturn(registry);
  when(keywhizConfig.getStatusCacheExpiry()).thenReturn(Duration.ofSeconds(1));

  this.status = new StatusResource(keywhizConfig, environment);
}

@Override
public HealthCheck.Result runHealthCheck(final String registryName, final String healthCheckName) {
    checkArgument(isNotBlank(registryName));
    checkArgument(isNotBlank(healthCheckName));
    checkArgument(SharedHealthCheckRegistries.names().contains(registryName));

    final HealthCheckRegistry registry = SharedHealthCheckRegistries.getOrCreate(registryName);
    checkArgument(registry.getNames().contains(healthCheckName));
    final HealthCheck.Result result = registry.runHealthCheck(healthCheckName);
    logAppEvent(result, registryName, healthCheckName);
    return result;
}

@Test
public void shouldCorrectlyRegisterReceiver() throws Exception {
    //GIVEN
    AmazonSQS sqs = mock(AmazonSQS.class);

    String queueUrl = "https://eu-central-1/queue.amazonaws.com/123456/test-queue";
    when(sqs.getQueueUrl("test-queue")).thenReturn(new GetQueueUrlResult()
            .withQueueUrl(queueUrl));

    LifecycleEnvironment lifecycle = mock(LifecycleEnvironment.class);
    doNothing().when(lifecycle).manage((Managed) anyObject());
    when(environment.lifecycle()).thenReturn(lifecycle);

    HealthCheckRegistry healthChecks = mock(HealthCheckRegistry.class);
    doNothing().when(healthChecks).register(anyObject(), anyObject());
    when(environment.healthChecks()).thenReturn(healthChecks);

    SqsBundle spiedBundle = spy(bundle);
    doReturn(sqs).when(spiedBundle).getAmazonSQS();

    spiedBundle.run(configurationHolder, environment);

    //WHEN
    spiedBundle.registerReceiver("test-queue", (m) -> process(m));

    //THEN
    verify(spiedBundle, times(1)).internalRegisterReceiver(eq("test-queue"), any(SqsReceiverHandler.class));
}

@Override
public boolean shouldBeInDiscovery(HealthCheckRegistry registry)
{
    return Iterables.all(registry.runHealthChecks().values(), new Predicate<HealthCheck.Result>()
    {
        @Override
        public boolean apply(HealthCheck.Result result)
        {
            return result.isHealthy();
        }
    });
}

public static void unregisterTaskHealthCheck(int taskId, String name) {
    HealthCheckRegistry healthCheckRegister = taskHealthCheckMap.get(taskId);

    if (healthCheckRegister != null) {
        healthCheckRegister.unregister(name);
    }

}

private Environment getEnvironment() {
    Environment environment = mock(Environment.class);
    JerseyEnvironment jersey = mock(JerseyEnvironment.class);
    when(environment.jersey()).thenReturn(jersey);
    LifecycleEnvironment lifecycle = mock(LifecycleEnvironment.class);
    when(environment.lifecycle()).thenReturn(lifecycle);
    HealthCheckRegistry healthchecks = mock(HealthCheckRegistry.class);
    when(environment.healthChecks()).thenReturn(healthchecks);
    return environment;
}