geronimo-opentracing-extension/src/main/java/org/apache/geronimo/opentracing/extension/proxy/TracingProxyFactory.java

/*
 * 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.geronimo.opentracing.extension.proxy;

import static java.util.Collections.emptyMap;
import static java.util.Optional.ofNullable;

import java.io.Serializable;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.concurrent.CompletionStage;
import java.util.stream.Stream;

import io.opentracing.Scope;
import io.opentracing.Span;
import io.opentracing.Tracer;
import io.opentracing.tag.Tags;

public class TracingProxyFactory {
    public <T> T decorate(final Tracer tracer, final T instance) {
        return decorate(tracer, instance, emptyMap());
    }

    public <T> T decorate(final Tracer tracer, final T instance, final Map<String, String> tags) {
        final Class<?>[] interfaces = instance.getClass().getInterfaces();
        if (interfaces.length == 0) {
            throw new IllegalArgumentException("Can't determine the API to proxy: " + instance);
        }
        final Class<T> mainApi = (Class<T>) interfaces[0];
        final Class<?>[] otherApis = interfaces.length == 1 ?
                new Class<?>[0] : Stream.of(interfaces).skip(1).toArray(Class[]::new);
        return decorate(tracer, instance, mainApi, tags, otherApis);
    }

    public <T> T decorate(final Tracer tracer,
                          final T instance,
                          final Class<T> mainApi,
                          final Map<String, String> tags,
                          final Class<?>... otherApis) {
        return mainApi.cast(Proxy.newProxyInstance(
                ofNullable(Thread.currentThread().getContextClassLoader()).orElseGet(ClassLoader::getSystemClassLoader),
                Stream.concat(Stream.of(mainApi), Stream.of(otherApis)).toArray(Class[]::new),
                new TracingHandler(instance, tracer, tags)));
    }

    private static class TracingHandler implements InvocationHandler, Serializable {
        private final Object delegate;
        private final Tracer tracer;
        private final Map<String, String> tags;

        private TracingHandler(final Object delegate, final Tracer tracer, final Map<String, String> tags) {
            this.delegate = delegate;
            this.tracer = tracer;
            this.tags = tags;
        }

        @Override
        public Object invoke(final Object proxy, final Method method, final Object[] args) throws Throwable {
            final Tracer.SpanBuilder builder = tracer.buildSpan(method.getDeclaringClass().getName() + "." + method.getName());
            builder.withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_SERVER);
            builder.withTag(Tags.COMPONENT.getKey(), "proxy");
            tags.forEach(builder::withTag);
            ofNullable(tracer.activeSpan()).ifPresent(builder::asChildOf);

            final Scope scope = builder.startActive(false /*just handle span inheritance for async case*/);
            boolean doFinish = true;
            try {
                final Object result = method.invoke(delegate, args);
                if (CompletionStage.class.isInstance(result)) {
                    doFinish = false;
                    final CompletionStage<?> stage = CompletionStage.class.cast(result);
                    return stage.handle((r, e) -> {
                        try {
                            if (e != null) {
                                onError(scope, e);
                                return rethrow(e);
                            }
                            return r;
                        } finally {
                            scope.span().finish();
                        }
                    });
                }
                return result;
            } catch (final InvocationTargetException ite) {
                onError(scope, ite.getTargetException());
                throw ite.getTargetException();
            } finally {
                if (doFinish) {
                    scope.span().finish();
                }
                scope.close();
            }
        }

        private Object rethrow(final Throwable e) {
            if (RuntimeException.class.isInstance(e)) {
                throw RuntimeException.class.cast(e);
            }
            if (Error.class.isInstance(e)) {
                throw Error.class.cast(e);
            }
            throw new IllegalStateException(e);
        }

        private void onError(final Scope scope, final Throwable e) {
            final Span span = scope.span();
            Tags.ERROR.set(span, true);

            final Map<String, Object> logs = new LinkedHashMap<>();
            logs.put("event", Tags.ERROR.getKey());
            logs.put("error.object", e);
            span.log(logs);
        }
    }
}