/*
 * 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.execution;

import com.facebook.presto.spi.StandardErrorCode;
import io.airlift.stats.CounterStat;
import io.airlift.stats.TimeStat;
import org.weakref.jmx.Managed;
import org.weakref.jmx.Nested;

import static com.facebook.presto.spi.StandardErrorCode.ABANDONED_QUERY;
import static com.facebook.presto.spi.StandardErrorCode.USER_CANCELED;
import static java.util.concurrent.TimeUnit.MILLISECONDS;

public class SqlQueryManagerStats
{
    private final CounterStat startedQueries = new CounterStat();
    private final CounterStat completedQueries = new CounterStat();
    private final CounterStat failedQueries = new CounterStat();
    private final CounterStat abandonedQueries = new CounterStat();
    private final CounterStat canceledQueries = new CounterStat();
    private final CounterStat userErrorFailures = new CounterStat();
    private final CounterStat internalFailures = new CounterStat();
    private final CounterStat externalFailures = new CounterStat();
    private final CounterStat insufficientResourcesFailures = new CounterStat();
    private final TimeStat executionTime = new TimeStat(MILLISECONDS);

    public void queryStarted()
    {
        startedQueries.update(1);
    }

    public void queryFinished(QueryInfo info)
    {
        completedQueries.update(1);
        executionTime.add(info.getQueryStats().getEndTime().getMillis() - info.getQueryStats().getCreateTime().getMillis(), MILLISECONDS);

        if (info.getErrorCode() != null) {
            switch (StandardErrorCode.toErrorType(info.getErrorCode().getCode())) {
                case USER_ERROR:
                    userErrorFailures.update(1);
                    break;
                case INTERNAL_ERROR:
                    internalFailures.update(1);
                    break;
                case INSUFFICIENT_RESOURCES:
                    insufficientResourcesFailures.update(1);
                    break;
                case EXTERNAL:
                    externalFailures.update(1);
                    break;
            }

            if (info.getErrorCode().getCode() == ABANDONED_QUERY.toErrorCode().getCode()) {
                abandonedQueries.update(1);
            }
            else if (info.getErrorCode().getCode() == USER_CANCELED.toErrorCode().getCode()) {
                canceledQueries.update(1);
            }
            failedQueries.update(1);
        }
    }

    @Managed
    public long getRunningQueries()
    {
        return Math.max(0, startedQueries.getTotalCount() - completedQueries.getTotalCount());
    }

    @Managed
    @Nested
    public CounterStat getStartedQueries()
    {
        return startedQueries;
    }

    @Managed
    @Nested
    public CounterStat getCompletedQueries()
    {
        return completedQueries;
    }

    @Managed
    @Nested
    public CounterStat getFailedQueries()
    {
        return failedQueries;
    }

    @Managed
    @Nested
    public TimeStat getExecutionTime()
    {
        return executionTime;
    }

    @Managed
    @Nested
    public CounterStat getUserErrorFailures()
    {
        return userErrorFailures;
    }

    @Managed
    @Nested
    public CounterStat getInternalFailures()
    {
        return internalFailures;
    }

    @Managed
    @Nested
    public CounterStat getAbandonedQueries()
    {
        return abandonedQueries;
    }

    @Managed
    @Nested
    public CounterStat getCanceledQueries()
    {
        return canceledQueries;
    }

    @Managed
    @Nested
    public CounterStat getExternalFailures()
    {
        return externalFailures;
    }

    @Managed
    @Nested
    public CounterStat getInsufficientResourcesFailures()
    {
        return insufficientResourcesFailures;
    }
}

@@ -678,24 +678,66 @@
 ounterStat;%0A
+import io.airlift.stats.DistributionStat;%0A
 import io.ai
@@ -754,24 +754,24 @@
 s.TimeStat;%0A
-
 import org.w
@@ -1747,24 +1747,182 @@
 LLISECONDS);
+%0A    private final DistributionStat wallInputBytesRate = new DistributionStat();%0A    private final DistributionStat cpuInputByteRate = new DistributionStat();
 %0A%0A    public
@@ -2077,24 +2077,25 @@
 ate(1);%0A
+%0A
         
 executio
@@ -2090,138 +2090,618 @@
     
-executionTime.add(info.getQueryStats().getEndTime().getMillis() - info.getQueryStats().getCreateTime().getMillis(), MILLISECONDS);
+long rawInputBytes = info.getQueryStats().getRawInputDataSize().toBytes();%0A%0A        long executionWallMillis = info.getQueryStats().getEndTime().getMillis() - info.getQueryStats().getCreateTime().getMillis();%0A        executionTime.add(executionWallMillis, MILLISECONDS);%0A        if (executionWallMillis %3E 0) %7B%0A            wallInputBytesRate.add(rawInputBytes * 1000 / executionWallMillis);%0A        %7D%0A%0A        long executionCpuMillis = info.getQueryStats().getTotalCpuTime().toMillis();%0A        if (executionCpuMillis %3E 0) %7B%0A            cpuInputByteRate.add(rawInputBytes * 1000 / executionCpuMillis);%0A        %7D
 %0A%0A  
@@ -4919,32 +4919,32 @@
 ailures()%0A    %7B%0A
-
         return i
@@ -4978,11 +4978,382 @@
 s;%0A    %7D
+%0A%0A    @Managed(description = %22Distribution of query input data rates (wall)%22)%0A    @Nested%0A    public DistributionStat getWallInputBytesRate()%0A    %7B%0A        return wallInputBytesRate;%0A    %7D%0A%0A    @Managed(description = %22Distribution of query input data rates (cpu)%22)%0A    @Nested%0A    public DistributionStat getCpuInputByteRate()%0A    %7B%0A        return cpuInputByteRate;%0A    %7D
 %0A%7D%0A

Add JMX counters for query bytes rate distribution