com.spring.boxes.dollar.support.graphql.instrument.TimeInstrumentation Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of spring-visual Show documentation
The newest version!
package com.spring.boxes.dollar.support.graphql.instrument;

import graphql.ExecutionResult;
import graphql.PublicApi;
import graphql.execution.instrumentation.InstrumentationContext;
import graphql.execution.instrumentation.InstrumentationState;
import graphql.execution.instrumentation.SimpleInstrumentation;
import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters;
import graphql.execution.instrumentation.parameters.InstrumentationFieldFetchParameters;
import graphql.schema.DataFetcher;
import graphql.schema.GraphQLNonNull;
import graphql.schema.GraphQLObjectType;
import graphql.schema.GraphQLOutputType;
import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.CompletableFuture;

@Slf4j
@PublicApi
public class TimeInstrumentation extends SimpleInstrumentation {

    @Override
    public InstrumentationState createState() {
        return new TimeState();
    }

    @Override
    public InstrumentationContext beginExecution(InstrumentationExecutionParameters parameters) {
        TimeState tracingState = parameters.getInstrumentationState();
        tracingState.startTime = System.currentTimeMillis();
        return super.beginExecution(parameters);
    }

    @Override
    public DataFetcher instrumentDataFetcher(DataFetcher dataFetcher, InstrumentationFieldFetchParameters parameters) {
        // We only care about user code
        if (parameters.isTrivialDataFetcher()) {
            return dataFetcher;
        }
        return environment -> {
            long startTime = System.currentTimeMillis();
            Object result = dataFetcher.get(environment);
            if (result instanceof CompletableFuture) {
                ((CompletableFuture) result).whenComplete((r, ex) -> {
                    long totalTime = System.currentTimeMillis() - startTime;
                    log.info("[TimeInstrumentation] Async dataFetcher [{}] took {}ms", findDataFetcherTag(parameters), totalTime);
                });
            } else {
                long totalTime = System.currentTimeMillis() - startTime;
                log.info("[TimeInstrumentation] DataFetcher [{}] took {}ms", findDataFetcherTag(parameters), totalTime);
            }

            return result;
        };
    }

    @Override
    public CompletableFuture instrumentExecutionResult(ExecutionResult executionResult,
                                                                        InstrumentationExecutionParameters parameters) {
        TimeState timeState = parameters.getInstrumentationState();
        long totalTime = System.currentTimeMillis() - timeState.startTime;
        log.info("[TimeInstrumentation] Total execution took {}ms", totalTime);
        return super.instrumentExecutionResult(executionResult, parameters);
    }

    private String findDataFetcherTag(InstrumentationFieldFetchParameters parameters) {
        GraphQLOutputType type = parameters.getExecutionStepInfo().getParent().getType();
        GraphQLObjectType parent;
        if (type instanceof GraphQLNonNull) {
            parent = (GraphQLObjectType) ((GraphQLNonNull) type).getWrappedType();
        } else {
            parent = (GraphQLObjectType) type;
        }

        return parent.getName() + "." + parameters.getExecutionStepInfo().getPath().getSegmentName();
    }

    static class TimeState implements InstrumentationState {
        long startTime;
    }
}