/*
 * Copyright 2019 LINE Corporation
 *
 * LINE Corporation 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:
 *
 *   https://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.linecorp.armeria.server.brave;

import static com.linecorp.armeria.common.HttpStatus.BAD_REQUEST;
import static com.linecorp.armeria.common.HttpStatus.OK;
import static org.assertj.core.api.Assertions.assertThat;

import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;

import javax.annotation.Nullable;

import org.junit.After;
import org.junit.AssumptionViolatedException;
import org.junit.Test;

import com.linecorp.armeria.common.CommonPools;
import com.linecorp.armeria.common.HttpMethod;
import com.linecorp.armeria.common.HttpResponse;
import com.linecorp.armeria.common.HttpStatus;
import com.linecorp.armeria.common.MediaType;
import com.linecorp.armeria.common.SessionProtocol;
import com.linecorp.armeria.common.brave.RequestContextCurrentTraceContext;
import com.linecorp.armeria.server.Server;
import com.linecorp.armeria.server.ServerBuilder;

import brave.propagation.CurrentTraceContext;
import brave.test.http.ITHttpServer;

public class BraveServiceIntegrationTest extends ITHttpServer {

    @Nullable
    private Server server;

    @Override
    protected CurrentTraceContext.Builder currentTraceContextBuilder() {
        return RequestContextCurrentTraceContext.builder();
    }

    @Override
    protected void init() {
        final ServerBuilder sb = Server.builder();
        sb.service("/", (ctx, req) -> {
            if (req.method() == HttpMethod.OPTIONS) {
                return HttpResponse.of(OK, MediaType.PLAIN_TEXT_UTF_8, "");
            }
            return HttpResponse.of(HttpStatus.NOT_FOUND);
        });
        sb.service("/foo", (ctx, req) -> HttpResponse.of(OK, MediaType.PLAIN_TEXT_UTF_8, "bar"));
        sb.service("/async", (ctx, req) -> asyncResponse(future -> {
            future.complete(HttpResponse.of(OK, MediaType.PLAIN_TEXT_UTF_8, "bar"));
        }));

        sb.service("/exception", (ctx, req) -> {
            throw new IllegalStateException("not ready");
        });
        sb.service("/exceptionAsync", (ctx, req) -> asyncResponse(future -> {
            future.completeExceptionally(new IllegalStateException("not ready"));
        }));
        sb.service("/badrequest", (ctx, req) -> HttpResponse.of(BAD_REQUEST));

        sb.service("/items/:itemId",
                   (ctx, req) -> HttpResponse.of(OK, MediaType.PLAIN_TEXT_UTF_8,
                                                 String.valueOf(ctx.pathParam("itemId"))));
        sb.service("/async_items/:itemId", (ctx, req) -> asyncResponse(future -> {
            future.complete(HttpResponse.of(OK, MediaType.PLAIN_TEXT_UTF_8,
                                            String.valueOf(ctx.pathParam("itemId"))));
        }));
        // "/nested" left out as there's no sub-routing feature at the moment.

        sb.service("/child", (ctx, req) -> {
            tracing.tracer().nextSpan().name("child").start().finish();
            return HttpResponse.of(OK, MediaType.PLAIN_TEXT_UTF_8, "happy");
        });
        sb.service("/baggage", (ctx, req) -> {
            final String value = String.valueOf(BAGGAGE_FIELD.getValue());
            return HttpResponse.of(OK, MediaType.PLAIN_TEXT_UTF_8, value);
        });

        sb.service("/badrequest", (ctx, req) -> HttpResponse.of(BAD_REQUEST));

        sb.decorator(BraveService.newDecorator(httpTracing));

        server = sb.build();
        server.start().join();
    }

    HttpResponse asyncResponse(Consumer<CompletableFuture<HttpResponse>> completeResponse) {
        final CompletableFuture<HttpResponse> responseFuture = new CompletableFuture<>();
        final HttpResponse res = HttpResponse.from(responseFuture);
        CommonPools.workerGroup().next().submit(
                () -> completeResponse.accept(responseFuture));
        return res;
    }

    @Test
    @Override
    public void notFound() {
        throw new AssumptionViolatedException(
                "Armeria yields 'get /*' as a span name for a non-existent mapping.");
    }

    @Test
    @Override
    public void httpStatusCodeSettable_onUncaughtException() {
        throw new AssumptionViolatedException(
            "Can't currently control the HTTP status code on uncaught exception. #2656");
    }

    @Test
    @Override
    public void httpStatusCodeSettable_onUncaughtException_async() {
        throw new AssumptionViolatedException(
            "Can't currently control the HTTP status code on uncaught exception. #2656");
    }

    @After
    public void stopServer() {
        if (server != null) {
            server.stop();
        }
    }

    @Override
    protected String url(String path) {
        assertThat(server).isNotNull();
        return "http://127.0.0.1:" + server.activeLocalPort(SessionProtocol.HTTP) + path;
    }
}

@@ -850,24 +850,51 @@
 ssertThat;%0A%0A
+import java.time.Duration;%0A
 import java.
@@ -1455,24 +1455,76 @@
 aceContext;%0A
+import com.linecorp.armeria.common.util.SystemInfo;%0A
 import com.l
@@ -3533,99 +3533,924 @@
 %3E %7B%0A
-            tracing.tracer().nextSpan().name(%22child%22).start().finish();%0A            return 
+%0A            // The %60Span.start()%60 and %60Span.finish()%60 estimate the current time using%0A            // %60baseEpochMicros + (System.nanoTime() - baseTickNanos)%60.%0A            // However, Armeria achieves the current time using %60System.currentTimeMillis() * 1000%60 in Java8.%0A            // The two different way of getting the current time could cause an error of a millisecond.%0A            // Sometimes the tolerance makes the server appear to finish earlier than child in this test case.%0A            // This is not a problem in the Brave%7BService,Client%7D because they only use Armeria's timestamp.%0A            tracing.tracer().nextSpan().name(%22child%22).start().finish();%0A            final Duration delay;%0A            if (SystemInfo.javaVersion() %3E 8) %7B%0A                delay = Duration.ofNanos(1000);%0A            %7D else %7B%0A                delay = Duration.ofMillis(1);%0A            %7D%0A            return HttpResponse.delayed(
 Http
@@ -4501,16 +4501,24 @@
 %22happy%22)
+, delay)
 ;%0A      

Fix flaky test `BraveServiceIntegrationTest.childCompletesBeforeServer` (#2867)