/*
 * 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.DistributionStat;
import org.weakref.jmx.Managed;
import org.weakref.jmx.Nested;

public class SqlQueryManagerStats
{
    private final CounterStat startedQueries = new CounterStat();
    private final CounterStat completedQueries = new CounterStat();
    private final CounterStat failedQueries = 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 DistributionStat executionTime = new DistributionStat();

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

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

        if (info.getErrorCode() != null) {
            switch (StandardErrorCode.toErrorType(info.getErrorCode().getCode())) {
                case USER_ERROR:
                    userErrorFailures.update(1);
                    break;
                case INTERNAL:
                    internalFailures.update(1);
                    break;
                case INSUFFICIENT_RESOURCES:
                    insufficientResourcesFailures.update(1);
                    break;
                case EXTERNAL:
                    externalFailures.update(1);
                    break;
            }
            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 DistributionStat getExecutionTime()
    {
        return executionTime;
    }

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

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

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

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

@@ -1015,32 +1015,167 @@
  CounterStat();%0A
+    private final CounterStat abandonedQueries = new CounterStat();%0A    private final CounterStat canceledQueries = new CounterStat();%0A
     private fina
@@ -1219,32 +1219,32 @@
  CounterStat();%0A
-
     private fina
@@ -2402,24 +2402,363 @@
            %7D
+%0A%0A            if (info.getErrorCode().getCode() == StandardErrorCode.ABANDONED_QUERY.toErrorCode().getCode()) %7B%0A                abandonedQueries.update(1);%0A            %7D%0A            else if (info.getErrorCode().getCode() == StandardErrorCode.USER_CANCELED.toErrorCode().getCode()) %7B%0A                canceledQueries.update(1);%0A            %7D
 %0A           
@@ -3596,32 +3596,32 @@
 ailures()%0A    %7B%0A
-
         return i
@@ -3636,32 +3636,262 @@
 ailures;%0A    %7D%0A%0A
+    @Managed%0A    @Nested%0A    public CounterStat getAbandonedQueries()%0A    %7B%0A        return abandonedQueries;%0A    %7D%0A%0A    @Managed%0A    @Nested%0A    public CounterStat getCanceledQueries()%0A    %7B%0A        return canceledQueries;%0A    %7D%0A%0A
     @Managed%0A   

Add jmx stats for abandoned and canceled queries