/*
 * Copyright 2004-present Facebook. All Rights Reserved.
 */
package com.facebook.presto.operator;

import com.facebook.presto.execution.TaskOutput;
import com.google.common.base.Preconditions;
import io.airlift.units.DataSize;
import io.airlift.units.DataSize.Unit;
import io.airlift.units.Duration;

import javax.annotation.concurrent.NotThreadSafe;
import java.lang.management.ManagementFactory;
import java.lang.management.ThreadMXBean;
import java.util.concurrent.TimeUnit;

/**
 * This class is not thread safe, but the done state of the task is properly
 * propagated from the TaskOutput, which is thread safe, to this class.
 */
@NotThreadSafe
public class OperatorStats
{
    private static final ThreadMXBean THREAD_MX_BEAN = ManagementFactory.getThreadMXBean();

    private final TaskOutput taskOutput;

    private long declaredSize;
    private long declaredPositions;

    private long completedDataSize;
    private long completedPositions;

    long wallStartTime;
    long cpuStartTime;
    long userStartTime;

    private boolean finished;

    public OperatorStats()
    {
        this.taskOutput = null;
    }

    public OperatorStats(TaskOutput taskOutput)
    {
        Preconditions.checkNotNull(taskOutput, "taskOutput is null");
        this.taskOutput = taskOutput;
    }

    public boolean isDone()
    {
        return finished || (taskOutput != null && taskOutput.getState().isDone());
    }

    public void addDeclaredSize(long bytes)
    {
        if (taskOutput == null) {
            return;
        }

        taskOutput.getStats().addInputDataSize(new DataSize(bytes, Unit.BYTE));
        declaredSize += bytes;
    }

    public void addCompletedDataSize(long bytes)
    {
        if (taskOutput == null) {
            return;
        }

        taskOutput.getStats().addCompletedDataSize(new DataSize(bytes, Unit.BYTE));
        completedDataSize += bytes;

        if (completedDataSize > declaredSize) {
            taskOutput.getStats().addInputDataSize(new DataSize(completedDataSize - declaredSize, Unit.BYTE));
            declaredSize = completedDataSize;
        }
    }

    public void addCompletedPositions(long positions)
    {
        if (taskOutput == null) {
            return;
        }

        taskOutput.getStats().addCompletedPositions(positions);
        completedPositions += positions;

        if (completedPositions > declaredPositions) {
            taskOutput.getStats().addInputPositions(completedPositions - declaredPositions);
            declaredPositions = completedPositions;
        }
    }

    public void start()
    {
        if (taskOutput == null) {
            return;
        }

        taskOutput.getStats().splitStarted();
        wallStartTime = System.nanoTime();
        cpuStartTime = THREAD_MX_BEAN.getCurrentThreadCpuTime();
        userStartTime = THREAD_MX_BEAN.getCurrentThreadUserTime();
    }

    public void finish()
    {
        if (finished) {
            return;
        }
        finished = true;

        if (taskOutput == null) {
            return;
        }

        // update the timings
        taskOutput.getStats().addSplitWallTime(Duration.nanosSince(wallStartTime));
        taskOutput.getStats().addSplitCpuTime(new Duration(Math.max(0, THREAD_MX_BEAN.getCurrentThreadCpuTime() - cpuStartTime), TimeUnit.NANOSECONDS));
        taskOutput.getStats().addSplitUserTime(new Duration(Math.max(0, THREAD_MX_BEAN.getCurrentThreadUserTime() - userStartTime), TimeUnit.NANOSECONDS));
        taskOutput.getStats().splitCompleted();
    }
}

@@ -2102,32 +2102,58 @@
 aSize;%0A        %7D
+%0A%0A        updateTimings();
 %0A    %7D%0A%0A    publ
@@ -2579,24 +2579,50 @@
 s;%0A        %7D
+%0A%0A        updateTimings();
 %0A    %7D%0A%0A    
@@ -3128,29 +3128,352 @@
     
-// update the timings
+updateTimings();%0A        taskOutput.getStats().splitCompleted();%0A    %7D%0A%0A    private void updateTimings()%0A    %7B%0A        if (taskOutput == null) %7B%0A            return;%0A        %7D%0A%0A        long now = System.nanoTime();%0A        long cpuNow = THREAD_MX_BEAN.getCurrentThreadCpuTime();%0A        long userNow = THREAD_MX_BEAN.getCurrentThreadUserTime();%0A
 %0A   
@@ -3516,24 +3516,28 @@
 allTime(
+new 
 Duration
 .nanosSi
@@ -3532,33 +3532,63 @@
 tion
-.nanosSince(wallStartTime
+(Math.max(0, now - wallStartTime), TimeUnit.NANOSECONDS
 ));%0A
@@ -3662,48 +3662,14 @@
 (0, 
-THREAD_MX_BEAN.getCurrentThreadCpuTime()
+cpuNow
  - c
@@ -3782,49 +3782,15 @@
 (0, 
-THREAD_MX_BEAN.getCurrentThreadUserTime()
+userNow
  - u
@@ -3824,32 +3824,33 @@
 OSECONDS));%0A
+%0A
         
 taskOutput.g
@@ -3841,46 +3841,91 @@
     
-taskOutput.getStats().splitCompleted()
+wallStartTime = now;%0A        cpuStartTime = cpuNow;%0A        userStartTime = userNow
 ;%0A  

Update timing stats incrementally