Spring Data Repository Metrics Spring Data Rep

DefaultRepositoryTagsProvider.java

/*
 * Copyright 2020 the original author or authors.
 *
 * Licensed 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 com.example.demo;

import java.util.Collections;
import java.util.List;

import io.micrometer.core.instrument.Tag;
import io.micrometer.core.instrument.Tags;
import org.aopalliance.intercept.MethodInvocation;

/**
 * Default implementation of {@link RepositoryTagsProvider}.
 */
public class DefaultRepositoryTagsProvider implements RepositoryTagsProvider {

	private final List<RepositoryTagsContributor> contributors;

	public DefaultRepositoryTagsProvider() {
		this(Collections.emptyList());
	}

	/**
	 * Creates a new {@link DefaultRepositoryTagsProvider} that will provide tags from the
	 * given {@code contributors} in addition to its own.
	 * @param contributors the contributors that will provide additional tags
	 */
	public DefaultRepositoryTagsProvider(List<RepositoryTagsContributor> contributors) {
		this.contributors = contributors;
	}

	@Override
	public Iterable<Tag> getTags(Class<?> repositoryInterface, MethodInvocation invocation, Object result, Throwable exception) {
		Tags tags = Tags.of(RepositoryTags.repository(repositoryInterface), RepositoryTags
				.method(invocation), RepositoryTags
				.invocation(invocation), RepositoryTags
				.exception(exception), RepositoryTags.result(result));
		for (RepositoryTagsContributor contributor : this.contributors) {
			tags = tags.and(contributor
					.getTags(repositoryInterface, invocation, result, exception));
		}
		return tags;
	}

}

RepositoryMetricsConfiguration.java

/*
 * Copyright 2020 the original author or authors.
 *
 * Licensed 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 com.example.demo;

import java.lang.reflect.AnnotatedElement;
import java.lang.reflect.Method;
import java.util.Collections;
import java.util.Set;
import java.util.concurrent.CompletionStage;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import io.micrometer.core.annotation.Timed;
import io.micrometer.core.instrument.MeterRegistry;
import io.micrometer.core.instrument.Timer;
import org.aopalliance.intercept.MethodInterceptor;
import org.aopalliance.intercept.MethodInvocation;

import org.springframework.aop.framework.ProxyFactory;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.beans.factory.config.BeanPostProcessor;
import org.springframework.boot.actuate.metrics.AutoTimer;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.annotation.MergedAnnotationCollectors;
import org.springframework.core.annotation.MergedAnnotations;
import org.springframework.data.repository.core.RepositoryInformation;
import org.springframework.data.repository.core.support.RepositoryFactoryBeanSupport;
import org.springframework.data.repository.core.support.RepositoryProxyPostProcessor;
import org.springframework.data.repository.util.ReactiveWrapperConverters;
import org.springframework.data.repository.util.ReactiveWrappers;
import org.springframework.util.concurrent.ListenableFuture;
import org.springframework.util.concurrent.ListenableFutureCallback;

/**
 * @author Mark Paluch
 */
@Configuration(proxyBeanMethods = false)
public class RepositoryMetricsConfiguration {

	@Bean
	@ConditionalOnMissingBean(RepositoryTagsProvider.class)
	public DefaultRepositoryTagsProvider repositoryTagsProvider(ObjectProvider<RepositoryTagsContributor> contributors) {
		return new DefaultRepositoryTagsProvider(contributors.orderedStream()
				.collect(Collectors.toList()));
	}

	@Bean
	public RepoMetricsPostProcessor repoMetricsPostProcessor(MeterRegistry registry,
			RepositoryTagsProvider tagsProvider) {
		return new RepoMetricsPostProcessor(registry, tagsProvider);
	}

	static class RepoMetricsPostProcessor implements BeanPostProcessor {

		private final MeterRegistry registry;

		private final RepositoryTagsProvider tagsProvider;

		public RepoMetricsPostProcessor(MeterRegistry registry, RepositoryTagsProvider tagsProvider) {
			this.registry = registry;
			this.tagsProvider = tagsProvider;
		}

		@Override
		public Object postProcessBeforeInitialization(Object bean, String beanName) throws BeansException {

			if (bean instanceof RepositoryFactoryBeanSupport) {

				RepositoryFactoryBeanSupport<?, ?, ?> repositoryFactoryBean = (RepositoryFactoryBeanSupport<?, ?, ?>) bean;

				MetricsMethodInterceptor interceptor = new MetricsMethodInterceptor(this.registry, repositoryFactoryBean
						.getObjectType(), this.tagsProvider, "spring.data.repositories", AutoTimer.ENABLED);

				repositoryFactoryBean
						.addProxyPostProcessor(new RepositoryMetricsProxyPostProcessor(interceptor));
			}

			return bean;
		}
	}

	static class RepositoryMetricsProxyPostProcessor implements RepositoryProxyPostProcessor {

		private final MetricsMethodInterceptor metricsMethodInterceptor;

		public RepositoryMetricsProxyPostProcessor(MetricsMethodInterceptor metricsMethodInterceptor) {

			this.metricsMethodInterceptor = metricsMethodInterceptor;
		}

		@Override
		public void postProcess(ProxyFactory factory, RepositoryInformation repositoryInformation) {
			factory.addAdvice(this.metricsMethodInterceptor);
		}
	}

	static class MetricsMethodInterceptor implements MethodInterceptor {

		private final MeterRegistry registry;

		private final Class<?> repositoryInterface;

		private final RepositoryTagsProvider tagsProvider;

		private final String metricName;

		private final AutoTimer autoTimer;

		public MetricsMethodInterceptor(MeterRegistry registry, Class<?> repositoryInterface, RepositoryTagsProvider tagsProvider, String metricName, AutoTimer autoTimer) {
			this.registry = registry;
			this.repositoryInterface = repositoryInterface;
			this.tagsProvider = tagsProvider;
			this.metricName = metricName;
			this.autoTimer = autoTimer;
		}

		@Override
		public Object invoke(MethodInvocation invocation) throws Throwable {
			Timer.Sample timerSample = Timer.start(this.registry);
			try {
				Object result = invocation.proceed();
				if (ReactiveWrappers.supports(invocation.getMethod().getReturnType())) {

					Object resultToUse = ReactiveWrapperConverters
							.doOnSuccess(result, () -> record(timerSample, invocation, result, null));
					return ReactiveWrapperConverters
							.doOnError(resultToUse, (ex) -> record(timerSample, invocation, result, ex));
				}
				if (result instanceof CompletionStage) {
					return ((CompletionStage<Object>) result)
							.whenComplete((o, throwable) -> {
								record(timerSample, invocation, result, throwable);
							});
				}
				if (result instanceof ListenableFuture) {

					ListenableFuture<Object> future = (ListenableFuture<Object>) result;
					future.addCallback(new ListenableFutureCallback<Object>() {
						@Override
						public void onFailure(Throwable ex) {
							record(timerSample, invocation, null, ex);
						}

						@Override
						public void onSuccess(Object result) {
							record(timerSample, invocation, result, null);
						}
					});
					return future;
				}
				if (result instanceof Stream) {
					return ((Stream<Object>) result)
							.onClose(() -> record(timerSample, invocation, result, null));
				}
				record(timerSample, invocation, result, null);
				return result;
			}
			catch (Throwable throwable) {
				record(timerSample, invocation, null, throwable);
				throw throwable;
			}
		}

		private void record(Timer.Sample timerSample, MethodInvocation invocation, Object result,
				Throwable exception) {
			Set<Timed> annotations = getTimedAnnotations(invocation.getMethod());
			if (annotations.isEmpty()) {
				if (this.autoTimer.isEnabled()) {
					Timer.Builder builder = this.autoTimer.builder(this.metricName);
					timerSample
							.stop(getTimer(builder, invocation, result, exception));
				}
			}
			else {
				for (Timed annotation : annotations) {
					Timer.Builder builder = Timer.builder(annotation, this.metricName);
					timerSample
							.stop(getTimer(builder, invocation, result, exception));
				}
			}
		}

		private Set<Timed> getTimedAnnotations(Method method) {
			Set<Timed> methodAnnotations = findTimedAnnotations(method);
			if (!methodAnnotations.isEmpty()) {
				return methodAnnotations;
			}
			return findTimedAnnotations(method.getDeclaringClass());
		}

		private Set<Timed> findTimedAnnotations(AnnotatedElement element) {
			MergedAnnotations annotations = MergedAnnotations.from(element);
			if (!annotations.isPresent(Timed.class)) {
				return Collections.emptySet();
			}
			return annotations.stream(Timed.class)
					.collect(MergedAnnotationCollectors.toAnnotationSet());
		}

		private Timer getTimer(Timer.Builder builder, MethodInvocation invocation, Object result,
				Throwable exception) {
			return builder.tags(this.tagsProvider
					.getTags(this.repositoryInterface, invocation, result, exception))
					.register(this.registry);
		}

	}

}

RepositoryTags.java

/*
 * Copyright 2020 the original author or authors.
 *
 * Licensed 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 com.example.demo;

import java.util.Collection;
import java.util.Optional;
import java.util.concurrent.CompletionStage;
import java.util.stream.Stream;

import io.micrometer.core.instrument.Tag;
import org.aopalliance.intercept.MethodInvocation;

import org.springframework.data.repository.util.ReactiveWrappers;
import org.springframework.lang.Nullable;
import org.springframework.util.StringUtils;
import org.springframework.util.concurrent.ListenableFuture;

/**
 * Factory methods for {@link Tag Tags} associated with repository invocation that
 * is handled by Spring Data.
 */
public final class RepositoryTags {

	private static final Tag RESULT_NULL = Tag.of("result", "NULL");

	private static final Tag RESULT_EMPTY = Tag.of("result", "EMPTY");

	private static final Tag RESULT_PRESENT = Tag.of("result", "PRESENT");

	private static final Tag RESULT_STREAM = Tag.of("result", "STREAM");

	private static final Tag EXECUTION_SYNC = Tag.of("invocation", "SYNC");

	private static final Tag EXECUTION_ASYNC = Tag.of("invocation", "ASYNC");

	private static final Tag EXECUTION_REACTIVE = Tag.of("invocation", "REACTIVE");

	private static final Tag EXCEPTION_NONE = Tag.of("exception", "None");

	private RepositoryTags() {
	}

	/**
	 * Creates a {@code repository} tag based on the repository interface class of {@link Class}.
	 * @param repositoryInterface the invocation
	 * @return the repository tag whose value is the simple class name
	 */
	public static Tag repository(Class<?> repositoryInterface) {
		return Tag.of("repository", repositoryInterface
				.getSimpleName());
	}

	/**
	 * Creates a {@code method} tag based on the {@link MethodInvocation#getMethod()
	 * method} of the given {@code invocation}.
	 * @param invocation the invocation
	 * @return the method tag whose value is the method name
	 */
	public static Tag method(MethodInvocation invocation) {
		return Tag.of("method", invocation.getMethod().getName());
	}

	/**
	 * Creates a {@code invocation} tag based on the {@link MethodInvocation#getMethod()
	 * method return type} of the given {@code invocation}.
	 * @param invocation the invocation
	 * @return the execution tag whose value is sync, async or reactive
	 */
	public static Tag invocation(MethodInvocation invocation) {
		Class<?> returnType = invocation.getMethod().getReturnType();
		if (CompletionStage.class.isAssignableFrom(returnType) || ListenableFuture.class
				.isAssignableFrom(returnType)) {
			return EXECUTION_ASYNC;
		}
		if (ReactiveWrappers.supports(returnType)) {
			return EXECUTION_REACTIVE;
		}
		return EXECUTION_SYNC;
	}

	/**
	 * Creates an {@code exception} tag based on the {@link Class#getSimpleName() simple
	 * name} of the class of the given {@code exception}.
	 * @param exception the exception, may be {@code null}
	 * @return the exception tag derived from the exception
	 */
	public static Tag exception(Throwable exception) {
		if (exception != null) {
			String simpleName = exception.getClass().getSimpleName();
			return Tag.of("exception", StringUtils
					.hasText(simpleName) ? simpleName : exception.getClass().getName());
		}
		return EXCEPTION_NONE;
	}

	/**
	 * Creates an {@code result} tag based on the status of the given {@code result}.
	 * @param result the method response
	 * @return the outcome tag derived from the status of the response
	 */
	public static Tag result(@Nullable Object result) {
		if (result == null) {
			return RESULT_NULL;
		}
		if (result instanceof Optional) {
			return ((Optional<?>) result).isPresent() ? RESULT_PRESENT : RESULT_EMPTY;
		}
		if (result instanceof Collection) {
			return ((Collection<?>) result).isEmpty() ? RESULT_EMPTY : RESULT_PRESENT;
		}
		if (result instanceof Stream) {
			return RESULT_STREAM;
		}
		return RESULT_PRESENT;
	}

}

RepositoryTagsContributor.java

/*
 * Copyright 2020 the original author or authors.
 *
 * Licensed 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 com.example.demo;

import io.micrometer.core.instrument.Tag;
import org.aopalliance.intercept.MethodInvocation;

/**
 * Provides {@link Tag Tags} for Spring Data Repository invocations.
 */
public interface RepositoryTagsContributor {

	/**
	 * Provides tags to be associated with metrics for the given {@code repositoryInterface} and
	 * {@code invocation}.
	 * @param repositoryInterface the repository interface type
	 * @param invocation the invoked method
	 * @param result the result of the method invocation, can be {@code null}
	 * @param exception the current exception, if any
	 * @return tags to associate with metrics for the method invocation
	 */
	Iterable<Tag> getTags(Class<?> repositoryInterface, MethodInvocation invocation, Object result,
			Throwable exception);

}

RepositoryTagsProvider.java

/*
 * Copyright 2020 the original author or authors.
 *
 * Licensed 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 com.example.demo;

import io.micrometer.core.instrument.Tag;
import org.aopalliance.intercept.MethodInvocation;

/**
 * A contributor of {@link Tag Tags} for Spring Data repositories. Typically used
 * by a {@link RepositoryTagsProvider} to provide tags in addition to its defaults.
 */
public interface RepositoryTagsProvider {

	/**
	 * Provides tags to be associated with metrics for the given {@code repositoryInterface} and
	 * {@code invocation}.
	 * @param repositoryInterface the repository interface type
	 * @param invocation the invoked method
	 * @param result the result of the method invocation, can be {@code null}
	 * @param exception the current exception, if any
	 * @return tags to associate with metrics for the method invocation
	 */
	Iterable<Tag> getTags(Class<?> repositoryInterface, MethodInvocation invocation, Object result,
			Throwable exception);

}