/*
 * Licensed 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.facebook.presto.server.remotetask;

import io.airlift.stats.DistributionStat;
import org.weakref.jmx.Managed;
import org.weakref.jmx.Nested;

public class RemoteTaskStats
{
    private final DistributionStat updateRoundTripMillis = new DistributionStat();
    private final DistributionStat infoRoundTripMillis = new DistributionStat();
    private final DistributionStat statusRoundTripMillis = new DistributionStat();
    private final DistributionStat responseSizeBytes = new DistributionStat();

    private long requestSuccess;
    private long requestFailure;

    public void statusRoundTripMillis(long roundTripMillis)
    {
        statusRoundTripMillis.add(roundTripMillis);
    }

    public void infoRoundTripMillis(long roundTripMillis)
    {
        infoRoundTripMillis.add(roundTripMillis);
    }

    public void updateRoundTripMillis(long roundTripMillis)
    {
        updateRoundTripMillis.add(roundTripMillis);
    }

    public void responseSize(long responseSizeBytes)
    {
        this.responseSizeBytes.add(responseSizeBytes);
    }

    public void updateSuccess()
    {
        requestSuccess++;
    }

    public void updateFailure()
    {
        requestFailure++;
    }

    @Managed
    @Nested
    public DistributionStat getResponseSizeBytes()
    {
        return responseSizeBytes;
    }

    @Managed
    @Nested
    public DistributionStat getStatusRoundTripMillis()
    {
        return statusRoundTripMillis;
    }

    @Managed
    @Nested
    public DistributionStat getUpdateRoundTripMillis()
    {
        return updateRoundTripMillis;
    }

    @Managed
    @Nested
    public DistributionStat getInfoRoundTripMillis()
    {
        return infoRoundTripMillis;
    }

    @Managed
    public long getRequestSuccess()
    {
        return requestSuccess;
    }

    @Managed
    public long getRequestFailure()
    {
        return requestFailure;
    }
}

@@ -614,41 +614,54 @@
 ort 
-io.airlift.stats.DistributionStat
+com.google.common.util.concurrent.AtomicDouble
 ;%0Aim
@@ -694,37 +694,54 @@
 ed;%0A
+%0A
 import 
-org.weakref.jmx.Nested
+javax.annotation.concurrent.ThreadSafe
 ;%0A%0Ap
@@ -788,32 +788,34 @@
 e final 
-DistributionStat
+IncrementalAverage
  updateR
@@ -835,32 +835,34 @@
 s = new 
-DistributionStat
+IncrementalAverage
 ();%0A    
@@ -875,32 +875,34 @@
 e final 
-DistributionStat
+IncrementalAverage
  infoRou
@@ -920,32 +920,34 @@
 s = new 
-DistributionStat
+IncrementalAverage
 ();%0A    
@@ -960,32 +960,34 @@
 e final 
-DistributionStat
+IncrementalAverage
  statusR
@@ -1007,32 +1007,34 @@
 s = new 
-DistributionStat
+IncrementalAverage
 ();%0A    
@@ -1047,32 +1047,34 @@
 e final 
-DistributionStat
+IncrementalAverage
  respons
@@ -1090,32 +1090,34 @@
 s = new 
-DistributionStat
+IncrementalAverage
 ();%0A%0A   
@@ -1835,43 +1835,21 @@
     
-@Nested%0A    public DistributionStat
+public double
  get
@@ -1906,16 +1906,22 @@
 izeBytes
+.get()
 ;%0A    %7D%0A
@@ -1942,43 +1942,21 @@
     
-@Nested%0A    public DistributionStat
+public double
  get
@@ -2013,32 +2013,38 @@
 sRoundTripMillis
+.get()
 ;%0A    %7D%0A%0A    @Ma
@@ -2057,43 +2057,21 @@
     
-@Nested%0A    public DistributionStat
+public double
  get
@@ -2128,32 +2128,38 @@
 eRoundTripMillis
+.get()
 ;%0A    %7D%0A%0A    @Ma
@@ -2172,43 +2172,21 @@
     
-@Nested%0A    public DistributionStat
+public double
  get
@@ -2247,16 +2247,22 @@
 ipMillis
+.get()
 ;%0A    %7D%0A
@@ -2446,11 +2446,454 @@
 e;%0A    %7D
+%0A%0A    @ThreadSafe%0A    private static class IncrementalAverage%0A    %7B%0A        private long count;%0A        private final AtomicDouble average = new AtomicDouble();%0A%0A        synchronized void add(long value)%0A        %7B%0A            count++;%0A            double oldAverage = average.get();%0A            average.set(oldAverage + ((value - oldAverage) / count));%0A        %7D%0A%0A        double get()%0A        %7B%0A            return average.get();%0A        %7D%0A    %7D
 %0A%7D%0A

Replace DistributionStat with IncrementalAverage in RemoteTaskStats