/*
 * Copyright 2016 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:
 *
 *   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.linecorp.armeria.it.metrics;

import static org.assertj.core.api.Assertions.assertThat;
import static org.junit.Assert.assertEquals;

import java.util.concurrent.CountDownLatch;

import org.junit.Test;

import com.codahale.metrics.MetricRegistry;

import com.linecorp.armeria.client.Client;
import com.linecorp.armeria.client.ClientBuilder;
import com.linecorp.armeria.client.ClientRequestContext;
import com.linecorp.armeria.client.DecoratingClient;
import com.linecorp.armeria.client.logging.DropwizardMetricCollectingClient;
import com.linecorp.armeria.common.RpcRequest;
import com.linecorp.armeria.common.RpcResponse;
import com.linecorp.armeria.common.logging.RequestLogAvailability;
import com.linecorp.armeria.server.ServerBuilder;
import com.linecorp.armeria.server.logging.DropwizardMetricCollectingService;
import com.linecorp.armeria.server.thrift.THttpService;
import com.linecorp.armeria.service.test.thrift.main.HelloService.Iface;
import com.linecorp.armeria.test.AbstractServerTest;

public class DropwizardMetricsIntegrationTest extends AbstractServerTest {

    private final MetricRegistry metricRegistry = new MetricRegistry();

    @Override
    protected void configureServer(ServerBuilder sb) throws Exception {
        sb.serviceAt("/helloservice", THttpService.of((Iface) name -> {
            if ("world".equals(name)) {
                return "success";
            }
            throw new IllegalArgumentException("bad argument");
        }).decorate(DropwizardMetricCollectingService.newDecorator(
                metricRegistry, MetricRegistry.name("services"))));
    }

    @Test(timeout = 10000L)
    public void normal() throws Exception {
        final CountDownLatch latch = new CountDownLatch(7);
        makeRequest("world", latch);
        makeRequest("world", latch);
        makeRequest("space", latch);
        makeRequest("world", latch);
        makeRequest("space", latch);
        makeRequest("space", latch);
        makeRequest("world", latch);

        // Wait until all RequestLogs are collected.
        latch.await();

        assertEquals(3, metricRegistry.getMeters()
                                      .get(clientMetricName("hello", "failures"))
                                      .getCount());
        assertEquals(3, metricRegistry.getMeters()
                                      .get(serverMetricName("hello", "failures"))
                                      .getCount());
        assertEquals(4, metricRegistry.getMeters()
                                      .get(clientMetricName("hello", "successes"))
                                      .getCount());
        assertEquals(4, metricRegistry.getMeters()
                                      .get(serverMetricName("hello", "successes"))
                                      .getCount());
        assertEquals(7, metricRegistry.getTimers()
                                      .get(clientMetricName("hello", "requests"))
                                      .getCount());
        assertEquals(7, metricRegistry.getTimers()
                                      .get(serverMetricName("hello", "requests"))
                                      .getCount());
        assertEquals(210, metricRegistry.getMeters()
                                      .get(clientMetricName("hello", "requestBytes"))
                                      .getCount());
        assertEquals(210, metricRegistry.getMeters()
                                      .get(serverMetricName("hello", "requestBytes"))
                                      .getCount());

        // Can't assert with exact byte count because the failure responses contain stack traces.
        assertThat(metricRegistry.getMeters()
                                 .get(clientMetricName("hello", "responseBytes"))
                                 .getCount()).isGreaterThan(0);
        assertThat(metricRegistry.getMeters()
                                 .get(serverMetricName("hello", "responseBytes"))
                                 .getCount()).isGreaterThan(0);
    }

    private static String serverMetricName(String method, String property) {
        return MetricRegistry.name("services", "/helloservice", method, property);
    }

    private static String clientMetricName(String method, String property) {
        return MetricRegistry.name("clients", "HelloService", method, property);
    }

    private void makeRequest(String name, CountDownLatch latch) {
        Iface client = new ClientBuilder("tbinary+" + uri("/helloservice"))
                .decorator(RpcRequest.class, RpcResponse.class,
                           DropwizardMetricCollectingClient.newDecorator(
                                   metricRegistry, MetricRegistry.name("clients", "HelloService")))
                .decorator(RpcRequest.class, RpcResponse.class,
                           delegate -> new CountDownClient(delegate, latch))
                .build(Iface.class);
        try {
            client.hello(name);
        } catch (Throwable t) {
            // Ignore, we will count these up
        }
    }

    private static class CountDownClient
            extends DecoratingClient<RpcRequest, RpcResponse, RpcRequest, RpcResponse> {

        private final CountDownLatch latch;

        CountDownClient(Client<? super RpcRequest, ? extends RpcResponse> delegate, CountDownLatch latch) {
            super(delegate);
            this.latch = latch;
        }

        @Override
        public RpcResponse execute(ClientRequestContext ctx, RpcRequest req) throws Exception {
            ctx.log().addListener(unused -> latch.countDown(), RequestLogAvailability.COMPLETE);
            return delegate().execute(ctx, req);
        }
    }
}

@@ -1785,24 +1785,59 @@
 egistry();%0A%0A
+    private CountDownLatch latch;%0A%0A
     @Overrid
@@ -2134,16 +2134,235 @@
 ment%22);%0A
+        %7D).decorate((delegate, ctx, req) -%3E %7B%0A            ctx.log().addListener(log -%3E latch.countDown(),%0A                                  RequestLogAvailability.COMPLETE);%0A            return delegate.serve(ctx, req);%0A
         
@@ -2576,37 +2576,16 @@
         
-final CountDownLatch 
 latch = 
@@ -2607,12 +2607,14 @@
 tch(
-7
+14
 );%0A
+%0A
     
@@ -2628,39 +2628,32 @@
 eRequest(%22world%22
-, latch
 );%0A        makeR
@@ -2658,39 +2658,32 @@
 eRequest(%22world%22
-, latch
 );%0A        makeR
@@ -2688,39 +2688,32 @@
 eRequest(%22space%22
-, latch
 );%0A        makeR
@@ -2726,23 +2726,16 @@
 (%22world%22
-, latch
 );%0A     
@@ -2748,39 +2748,32 @@
 eRequest(%22space%22
-, latch
 );%0A        makeR
@@ -2786,23 +2786,16 @@
 (%22space%22
-, latch
 );%0A     
@@ -2816,23 +2816,16 @@
 (%22world%22
-, latch
 );%0A%0A    
@@ -5257,34 +5257,12 @@
 name
-, CountDownLatch latch
 ) %7B%0A
+
     
@@ -5638,84 +5638,29 @@
 ass,
-%0A                           delegate -%3E new CountDownClient(delegate, latch)
+ CountDownClient::new
 )%0A  
@@ -5847,22 +5847,21 @@
 private 
-static
+final
  class C
@@ -5969,53 +5969,8 @@
  %7B%0A%0A
-        private final CountDownLatch latch;%0A%0A
     
@@ -6051,30 +6051,8 @@
 gate
-, CountDownLatch latch
 ) %7B%0A
@@ -6084,40 +6084,8 @@
 e);%0A
-            this.latch = latch;%0A
     

Fix flaky DropwizardMetricsIntegrationTest (#470)