/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF 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 org.apache.beam.runners.dataflow.worker;

import static org.apache.beam.runners.dataflow.worker.counters.DataflowCounterUpdateExtractor.longToSplitInt;
import static org.apache.beam.runners.dataflow.worker.counters.DataflowCounterUpdateExtractor.splitIntToLong;
import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.hamcrest.Matchers.emptyIterable;
import static org.hamcrest.Matchers.hasItems;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertThat;

import com.google.api.services.dataflow.model.CounterMetadata;
import com.google.api.services.dataflow.model.CounterStructuredName;
import com.google.api.services.dataflow.model.CounterStructuredNameAndMetadata;
import com.google.api.services.dataflow.model.CounterUpdate;
import com.google.api.services.dataflow.model.DistributionUpdate;
import org.apache.beam.runners.core.metrics.ExecutionStateSampler;
import org.apache.beam.runners.core.metrics.ExecutionStateTracker;
import org.apache.beam.runners.core.metrics.ExecutionStateTracker.ExecutionState;
import org.apache.beam.runners.dataflow.worker.BatchModeExecutionContext.BatchModeExecutionState;
import org.apache.beam.runners.dataflow.worker.counters.NameContext;
import org.apache.beam.runners.dataflow.worker.profiler.ScopedProfiler.NoopProfileScope;
import org.apache.beam.runners.dataflow.worker.profiler.ScopedProfiler.ProfileScope;
import org.apache.beam.sdk.metrics.Counter;
import org.apache.beam.sdk.metrics.Distribution;
import org.apache.beam.sdk.metrics.MetricName;
import org.apache.beam.sdk.metrics.MetricsContainer;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.hamcrest.Matchers;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.mockito.Mockito;

/** Tests for {@link BatchModeExecutionContext}. */
@RunWith(JUnit4.class)
public class BatchModeExecutionContextTest {

  @Test
  public void extractMetricUpdatesCounter() {
    BatchModeExecutionContext executionContext =
        BatchModeExecutionContext.forTesting(PipelineOptionsFactory.create(), "testStage");
    DataflowOperationContext operationContext =
        executionContext.createOperationContext(NameContextsForTests.nameContextForTest());

    Counter counter =
        operationContext
            .metricsContainer()
            .getCounter(MetricName.named("namespace", "some-counter"));
    counter.inc(1);
    counter.inc(41);
    counter.inc(1);
    counter.inc(-1);

    final CounterUpdate expected =
        new CounterUpdate()
            .setStructuredNameAndMetadata(
                new CounterStructuredNameAndMetadata()
                    .setName(
                        new CounterStructuredName()
                            .setOrigin("USER")
                            .setOriginNamespace("namespace")
                            .setName("some-counter")
                            .setOriginalStepName("originalName"))
                    .setMetadata(new CounterMetadata().setKind("SUM")))
            .setCumulative(true)
            .setInteger(longToSplitInt(42));

    assertThat(executionContext.extractMetricUpdates(false), containsInAnyOrder(expected));

    executionContext.commitMetricUpdates();

    Counter counterUncommitted =
        operationContext
            .metricsContainer()
            .getCounter(MetricName.named("namespace", "uncommitted-counter"));
    counterUncommitted.inc(64);

    final CounterUpdate expectedUncommitted =
        new CounterUpdate()
            .setStructuredNameAndMetadata(
                new CounterStructuredNameAndMetadata()
                    .setName(
                        new CounterStructuredName()
                            .setOrigin("USER")
                            .setOriginNamespace("namespace")
                            .setName("uncommitted-counter")
                            .setOriginalStepName("originalName"))
                    .setMetadata(new CounterMetadata().setKind("SUM")))
            .setCumulative(true)
            .setInteger(longToSplitInt(64));

    // Expect to get only the uncommitted metric, unless final update.
    assertThat(
        executionContext.extractMetricUpdates(false), containsInAnyOrder(expectedUncommitted));
    assertThat(
        executionContext.extractMetricUpdates(true),
        containsInAnyOrder(expected, expectedUncommitted));

    executionContext.commitMetricUpdates();
    // All Metrics are committed, expect none unless final update.
    assertThat(executionContext.extractMetricUpdates(false), emptyIterable());
    assertThat(
        executionContext.extractMetricUpdates(true),
        containsInAnyOrder(expected, expectedUncommitted));
  }

  @Test
  public void extractMetricUpdatesDistribution() {
    BatchModeExecutionContext executionContext =
        BatchModeExecutionContext.forTesting(PipelineOptionsFactory.create(), "testStage");
    DataflowOperationContext operationContext =
        executionContext.createOperationContext(NameContextsForTests.nameContextForTest());

    Distribution distribution =
        operationContext
            .metricsContainer()
            .getDistribution(MetricName.named("namespace", "some-distribution"));
    distribution.update(2);
    distribution.update(8);

    final CounterUpdate expected =
        new CounterUpdate()
            .setStructuredNameAndMetadata(
                new CounterStructuredNameAndMetadata()
                    .setName(
                        new CounterStructuredName()
                            .setOrigin("USER")
                            .setOriginNamespace("namespace")
                            .setName("some-distribution")
                            .setOriginalStepName("originalName"))
                    .setMetadata(new CounterMetadata().setKind("DISTRIBUTION")))
            .setCumulative(true)
            .setDistribution(
                new DistributionUpdate()
                    .setCount(longToSplitInt(2))
                    .setMax(longToSplitInt(8))
                    .setMin(longToSplitInt(2))
                    .setSum(longToSplitInt(10)));

    assertThat(executionContext.extractMetricUpdates(false), containsInAnyOrder(expected));
  }

  @Test
  public void extractMsecCounters() {
    BatchModeExecutionContext executionContext =
        BatchModeExecutionContext.forTesting(PipelineOptionsFactory.create(), "testStage");

    MetricsContainer metricsContainer = Mockito.mock(MetricsContainer.class);
    ProfileScope otherScope = Mockito.mock(ProfileScope.class);
    ProfileScope profileScope = Mockito.mock(ProfileScope.class);
    ExecutionState start1 =
        executionContext.executionStateRegistry.getState(
            NameContext.create("stage", "original-1", "system-1", "user-1"),
            ExecutionStateTracker.START_STATE_NAME,
            metricsContainer,
            profileScope);
    ExecutionState process1 =
        executionContext.executionStateRegistry.getState(
            NameContext.create("stage", "original-1", "system-1", "user-1"),
            ExecutionStateTracker.PROCESS_STATE_NAME,
            metricsContainer,
            profileScope);
    ExecutionState start2 =
        executionContext.executionStateRegistry.getState(
            NameContext.create("stage", "original-2", "system-2", "user-2"),
            ExecutionStateTracker.START_STATE_NAME,
            metricsContainer,
            profileScope);

    ExecutionState other =
        executionContext.executionStateRegistry.getState(
            NameContext.forStage("stage"), "other", null, NoopProfileScope.NOOP);

    other.takeSample(120);

    start1.takeSample(100);
    process1.takeSample(500);

    assertThat(
        executionContext.extractMsecCounters(false),
        containsInAnyOrder(
            msecStage("other-msecs", "stage", 120),
            msec("start-msecs", "stage", "original-1", 100),
            msec("process-msecs", "stage", "original-1", 500)));

    process1.takeSample(200);
    start2.takeSample(200);

    assertThat(
        executionContext.extractMsecCounters(false),
        containsInAnyOrder(
            msec("process-msecs", "stage", "original-1", 500 + 200),
            msec("start-msecs", "stage", "original-2", 200)));

    process1.takeSample(300);
    assertThat(
        executionContext.extractMsecCounters(true),
        hasItems(
            msecStage("other-msecs", "stage", 120),
            msec("start-msecs", "stage", "original-1", 100),
            msec("process-msecs", "stage", "original-1", 500 + 200 + 300),
            msec("start-msecs", "stage", "original-2", 200)));
  }

  @Test
  public void extractThrottleTimeCounters() {
    BatchModeExecutionContext executionContext =
        BatchModeExecutionContext.forTesting(PipelineOptionsFactory.create(), "testStage");
    DataflowOperationContext operationContext =
        executionContext.createOperationContext(NameContextsForTests.nameContextForTest());

    Counter counter =
        operationContext
            .metricsContainer()
            .getCounter(
                MetricName.named(
                    BatchModeExecutionContext.DATASTORE_THROTTLE_TIME_NAMESPACE,
                    "cumulativeThrottlingSeconds"));
    counter.inc(12);
    counter.inc(17);
    counter.inc(1);

    assertEquals(30L, (long) executionContext.extractThrottleTime());
  }

  private CounterUpdate msecStage(String counterName, String stageName, long value) {
    return new CounterUpdate()
        .setStructuredNameAndMetadata(
            new CounterStructuredNameAndMetadata()
                .setName(
                    new CounterStructuredName()
                        .setOrigin("SYSTEM")
                        .setName(counterName)
                        .setExecutionStepName(stageName))
                .setMetadata(new CounterMetadata().setKind("SUM")))
        .setCumulative(true)
        .setInteger(longToSplitInt(value));
  }

  private CounterUpdate msec(
      String counterName, String stageName, String originalStepName, long value) {
    return new CounterUpdate()
        .setStructuredNameAndMetadata(
            new CounterStructuredNameAndMetadata()
                .setName(
                    new CounterStructuredName()
                        .setOrigin("SYSTEM")
                        .setName(counterName)
                        .setOriginalStepName(originalStepName)
                        .setExecutionStepName(stageName))
                .setMetadata(new CounterMetadata().setKind("SUM")))
        .setCumulative(true)
        .setInteger(longToSplitInt(value));
  }

  @Test(timeout = 2000)
  public void stateSamplingInBatch() {
    // Test that when writing on one thread and reading from another, updates always eventually
    // reach the reading thread.
    BatchModeExecutionState state =
        new BatchModeExecutionState(
            NameContextsForTests.nameContextForTest(),
            "testState",
            null /* requestingStepName */,
            null /* inputIndex */,
            null /* metricsContainer */,
            NoopProfileScope.NOOP);
    ExecutionStateSampler sampler = ExecutionStateSampler.newForTest();
    try {
      sampler.start();

      ExecutionStateTracker tracker = new ExecutionStateTracker(sampler);

      Thread executionThread = new Thread();
      executionThread.setName("looping-thread-for-test");

      tracker.activate(executionThread);
      tracker.enterState(state);

      // Wait for the state to be incremented 3 times
      long value = 0;
      for (int i = 0; i < 3; i++) {
        CounterUpdate update = null;
        while (update == null) {
          update = state.extractUpdate(false);
        }
        long newValue = splitIntToLong(update.getInteger());
        assertThat(newValue, Matchers.greaterThan(value));
        value = newValue;
      }
    } finally {
      sampler.stop();
    }
  }
}