package com.atlassian.bitbucket.mesh.test.context;

import com.atlassian.bitbucket.mesh.grpc.GrpcServer;
import com.atlassian.bitbucket.mesh.grpc.GrpcStatistics;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.Conventions;
import org.springframework.test.context.TestContext;
import org.springframework.test.context.TestExecutionListener;
import org.springframework.test.context.support.AbstractTestExecutionListener;

import javax.annotation.Nonnull;
import java.util.concurrent.TimeUnit;

import static org.junit.Assert.fail;

//There's a version of this listener in Bitbucket Server's Git SCM, but its implementation relies on ClientInterceptor
//since it can't integrate with the GrpcServer. If you make changes here, you may need to make changes there too.
/**
 * A {@link TestExecutionListener} which uses {@link GrpcStatistics} to track how many calls each test method
 * starts and ensure they all complete to verify no calls are "leaked".
 */
public class GrpcCallTestExecutionListener extends AbstractTestExecutionListener {

    private static final String ATTR_REQUEST_BASELINE = Conventions.getQualifiedAttributeName(
            GrpcCallTestExecutionListener.class, "requestBaseline");
    private static final String ATTR_REQUEST_OFFSET = Conventions.getQualifiedAttributeName(
            GrpcCallTestExecutionListener.class, "requestOffset");

    private static final Logger log = LoggerFactory.getLogger(GrpcCallTestExecutionListener.class);

    private final GrpcServer server;

    public GrpcCallTestExecutionListener(GrpcServer server) {
        this.server = server;
    }

    @Override
    public void afterTestMethod(@Nonnull TestContext testContext) {
        //If baseline and/or offset was set in beforeTestMethod, we want to ensure both are always cleared
        long baseline = popAttribute(testContext, ATTR_REQUEST_BASELINE);
        long offset = popAttribute(testContext, ATTR_REQUEST_OFFSET);

        GrpcStatistics statistics = server.getStatistics();
        long newlyStarted = statistics.getTotalRequests() - baseline;
        if (newlyStarted < 1L) {
            //The test didn't start any new RPCs, so there's no reason to check for incomplete ones
            return;
        }

        //"offset" is inherently racy, because incomplete RPCs left over from a previous test could complete
        //at any moment. That means the calculated "incomplete" count here could be negative, if all the new
        //RPCs for the test completed and one or more "offset" ones did too. That also means we could get 0
        //when a newly-started RPC leaked, if a previously-leaked RPC completed. These are edge cases though,
        //and the check is still sufficiently reliable to be a useful indicator for leaks
        long incomplete = statistics.getRunningRequests() - offset;
        if (incomplete < 1L) {
            return;
        }

        long start = System.currentTimeMillis();
        long timeout = start + TimeUnit.SECONDS.toMillis(5L);

        //Because when the future the request thread blocks on is released and when the server detects the RPC
        //as complete as inherently async and disconnected, a minimal (typically milliseconds) delay after the
        //test completes before the server detects the RPCs are complete isn't indicative as a problem. This
        //loop allows RPCs a little grace period to complete before failing the test
        long stillIncomplete;
        do {
            try {
                Thread.sleep(0L, 100);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }

            statistics = server.getStatistics();
            stillIncomplete = statistics.getRunningRequests() - offset;
            if (stillIncomplete == 0L) {
                //Because the client is almost always released before the server, some small delay between test
                //completion and RPC completion isn't unlikely. We apply a short "ultra-grace" period here where
                //we don't log, to reduce reported false positives
                long delay = System.currentTimeMillis() - start;
                if (delay > 10L) {
                    log.warn("{}.{} started {} gRPC request(s), of which {} had not completed when the " +
                                    "test method ended. The trailing request(s) completed {}ms later",
                            testContext.getTestClass().getSimpleName(), testContext.getTestMethod().getName(),
                            newlyStarted, incomplete, delay);
                }

                return;
            }
        } while (System.currentTimeMillis() < timeout);

        //If we make it here, at least 1 RPC still hasn't completed after 5 seconds. That's a strong indication
        //that there's a leak in the processing
        fail(testContext.getTestClass().getSimpleName() + "." + testContext.getTestMethod().getName() +
                " started " + newlyStarted + " gRPC request(s), of which " + incomplete +
                " had not completed 5 seconds after the test method ended. It may be leaking requests.");
    }

    @Override
    public void beforeTestMethod(@Nonnull TestContext testContext) {
        GrpcStatistics statistics = server.getStatistics();

        //If a previous test has started a request that hasn't finished, don't let that make the
        //check fail for every subsequent test. Before each test starts, capture the offset from
        //how many requests have started vs. how many have finished so we can compare it later
        testContext.setAttribute(ATTR_REQUEST_BASELINE, statistics.getTotalRequests());
        testContext.setAttribute(ATTR_REQUEST_OFFSET, statistics.getRunningRequests());
    }

    private static long popAttribute(TestContext context, String name) {
        Long value = (Long) context.removeAttribute(name);

        return value == null ? 0L : value;
    }
}
