core/camel-spring-boot/src/main/java/org/apache/camel/spring/boot/actuate/health/AsyncHealthIndicatorAutoConfiguration.java

package org.apache.camel.spring.boot.actuate.health;

import org.apache.camel.spring.boot.actuate.health.liveness.CamelLivenessStateHealthIndicator;
import org.apache.camel.spring.boot.actuate.health.readiness.CamelReadinessStateHealthIndicator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.boot.actuate.health.Health;
import org.springframework.boot.actuate.health.HealthContributorRegistry;
import org.springframework.boot.actuate.health.HealthIndicator;
import org.springframework.boot.actuate.health.NamedContributor;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.TaskScheduler;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;

import java.time.Duration;
import java.time.ZonedDateTime;

/**
 * Configuration class that replace synchronous Camel Health Checks with asynchronous ones.
 *
 * This implementation is based on https://github.com/spring-projects/spring-boot/issues/2652 that most probably
 * will be added in spring boot 3.2.x as a new feature in the future.
 *
 * TODO: To be refactored once async health contributors feature will be added in spring boot.
 *
 */
@Configuration
@ConditionalOnProperty(prefix = "camel.health", name = "async-camel-health-check", havingValue = "true")
public class AsyncHealthIndicatorAutoConfiguration implements InitializingBean {
	private static final Logger log = LoggerFactory.getLogger(AsyncHealthIndicatorAutoConfiguration.class);

	private HealthContributorRegistry healthContributorRegistry;
	private TaskScheduler taskScheduler;
	private CamelHealthCheckConfigurationProperties config;

	public AsyncHealthIndicatorAutoConfiguration(HealthContributorRegistry healthContributorRegistry,
												 CamelHealthCheckConfigurationProperties config) {
		this.healthContributorRegistry = healthContributorRegistry;
		this.config = config;

		ThreadPoolTaskScheduler threadPoolTaskScheduler = new ThreadPoolTaskScheduler();
		threadPoolTaskScheduler.setPoolSize(config.getHealthCheckPoolSize());
		threadPoolTaskScheduler.setThreadNamePrefix(config.getHealthCheckThreadNamePrefix());
		threadPoolTaskScheduler.initialize();
		taskScheduler = threadPoolTaskScheduler;
	}

	@Override
	public void afterPropertiesSet() throws Exception {
		for (NamedContributor<?> namedContributor : healthContributorRegistry) {
			final String name = namedContributor.getName();
			final Object contributor = namedContributor.getContributor();
			if (contributor instanceof CamelHealthCheckIndicator
				|| contributor instanceof CamelLivenessStateHealthIndicator
				|| contributor instanceof CamelReadinessStateHealthIndicator) {
				HealthIndicator camelHealthCheckIndicator = (HealthIndicator) contributor;
				healthContributorRegistry.unregisterContributor(name);
				log.debug(
						"Wrapping " + contributor.getClass().getSimpleName() + " for async health scheduling");
				WrappedHealthIndicator wrappedHealthIndicator =
						new WrappedHealthIndicator(camelHealthCheckIndicator);
				healthContributorRegistry.registerContributor(name, wrappedHealthIndicator);
				taskScheduler.scheduleWithFixedDelay(
						wrappedHealthIndicator, Duration.ofSeconds(config.getHealthCheckFrequency()));
			}
		}
	}

	/**
	 * Health Check Indicator that executes Health Checks within a Task Scheduler
	 */
	private static class WrappedHealthIndicator implements HealthIndicator, Runnable {
		private static final String LAST_CHECKED_KEY = "lastChecked";
		private static final String LAST_DURATION_KEY = "lastDuration";

		private HealthIndicator wrappedHealthIndicator;

		private Health lastHealth;

		public WrappedHealthIndicator(HealthIndicator wrappedHealthIndicator) {
			this.wrappedHealthIndicator = wrappedHealthIndicator;
		}

		@Override
		public Health health() {
			Health lastHealth = getLastHealth();
			if (lastHealth == null) {
				setLastHealth(getAndWrapHealth());
				lastHealth = getLastHealth();
			}

			return lastHealth;
		}

		private Health getAndWrapHealth() {
			ZonedDateTime startTime = ZonedDateTime.now();
			Health baseHealth = getWrappedHealthIndicator().health();
			ZonedDateTime endTime = ZonedDateTime.now();
			Duration duration = Duration.between(startTime, endTime);
			return Health.status(baseHealth.getStatus())
					.withDetails(baseHealth.getDetails())
					.withDetail(LAST_CHECKED_KEY, startTime)
					.withDetail(LAST_DURATION_KEY, duration)
					.build();
		}

		@Override
		public void run() {
			setLastHealth(getAndWrapHealth());
		}

		public HealthIndicator getWrappedHealthIndicator() {
			return wrappedHealthIndicator;
		}

		public Health getLastHealth() {
			return lastHealth;
		}

		public void setLastHealth(Health lastHealth) {
			this.lastHealth = lastHealth;
		}
	}
}