/*
 * 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.gobblin.runtime.mapreduce;

import com.google.common.collect.Maps;
import org.apache.gobblin.configuration.State;
import org.apache.gobblin.configuration.WorkUnitState;
import org.apache.gobblin.metrics.event.EventSubmitter;
import org.apache.gobblin.runtime.TaskContext;
import org.apache.gobblin.runtime.task.BaseAbstractTask;
import java.io.IOException;
import java.util.Map;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang.exception.ExceptionUtils;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.mapreduce.Job;


/**
 * A task that runs an MR job.
 *
 * Usage:
 * TaskUtils.setTaskFactoryClass(workUnit, MRTaskFactory.class);
 * MRTask.serializeJobToState(workUnit, myJob);
 *
 * Subclasses can override {@link #createJob()} to customize the way the MR job is prepared.
 */
@Slf4j
public class MRTask extends BaseAbstractTask {

  private static final String JOB_CONFIGURATION_PREFIX = "MRTask.jobConfiguration.";

  public static class Events {
    public static final String MR_JOB_STARTED_EVENT = "MRJobStarted";
    public static final String MR_JOB_SUCCESSFUL = "MRJobSuccessful";
    public static final String MR_JOB_FAILED = "MRJobFailed";

    public static final String JOB_URL = "jobTrackingUrl";
    public static final String FAILURE_CONTEXT = "failureContext";
  }

  public static void serializeJobToState(State state, Job job) {
    for (Map.Entry<String, String> entry : job.getConfiguration()) {
      state.setProp(JOB_CONFIGURATION_PREFIX + entry.getKey(), entry.getValue());
    }
  }

  private final TaskContext taskContext;
  private final EventSubmitter eventSubmitter;
  protected Job mrJob;

  public MRTask(TaskContext taskContext) {
    super(taskContext);
    this.taskContext = taskContext;
    this.eventSubmitter = new EventSubmitter.Builder(this.metricContext, "gobblin.MRTask")
        .addMetadata(additionalEventMetadata()).build();
  }

  public void onMRTaskComplete (boolean isSuccess, Throwable t) throws RuntimeException {
    if (isSuccess) {
      this.workingState = WorkUnitState.WorkingState.SUCCESSFUL;
    } else if (t == null) {
      this.workingState = WorkUnitState.WorkingState.FAILED;
    } else {
      log.error ("Failed to run MR job with exception {}", ExceptionUtils.getStackTrace(t));
      this.workingState = WorkUnitState.WorkingState.FAILED;
    }
  }

  @Override
  public void commit() {
    log.debug ("State is set to {} inside onMRTaskComplete.", this.workingState);
  }

  @Override
  public void run() {

    try {
      Job job = createJob();

      job.submit();
      this.eventSubmitter.submit(Events.MR_JOB_STARTED_EVENT, Events.JOB_URL, job.getTrackingURL());
      job.waitForCompletion(false);
      this.mrJob = job;

      if (job.isSuccessful()) {
        this.eventSubmitter.submit(Events.MR_JOB_SUCCESSFUL, Events.JOB_URL, job.getTrackingURL());
        this.onMRTaskComplete(true, null);
      } else {
        this.eventSubmitter.submit(Events.MR_JOB_FAILED, Events.JOB_URL, job.getTrackingURL());
        this.onMRTaskComplete (false, new IOException("MR Job is not successful"));
      }
    } catch (Throwable t) {
      log.error("Failed to run MR job.", t);
      this.eventSubmitter.submit(Events.MR_JOB_FAILED, Events.FAILURE_CONTEXT, t.getMessage());
      this.onMRTaskComplete (false, t);
    }
  }

  protected Map<String, String> additionalEventMetadata() {
    return Maps.newHashMap();
  }

  protected Job createJob() throws IOException {
    Job job = Job.getInstance(new Configuration());
    for (Map.Entry<Object, Object> entry : this.taskContext.getTaskState().getProperties().entrySet()) {
      if (entry.getKey() instanceof String && ((String) entry.getKey()).startsWith(JOB_CONFIGURATION_PREFIX)) {
        String actualKey = ((String) entry.getKey()).substring(JOB_CONFIGURATION_PREFIX.length());
        job.getConfiguration().set(actualKey, (String) entry.getValue());
      }
    }
    return job;
  }

}

@@ -2005,16 +2005,80 @@
 Failed%22;
+%0A    public static final String MR_JOB_SKIPPED = %22MRJobSkipped%22;
 %0A%0A    pu
@@ -3439,16 +3439,268 @@
 Job();%0A%0A
+      if (job == null) %7B%0A        log.info(%22No MR job created. Skipping.%22);%0A        this.workingState = WorkUnitState.WorkingState.SUCCESSFUL;%0A        this.eventSubmitter.submit(Events.MR_JOB_SKIPPED);%0A        onSkippedMRJob();%0A        return;%0A      %7D%0A%0A
       jo
@@ -4562,24 +4562,213 @@
 Map();%0A  %7D%0A%0A
+  /**%0A   * Create the %7B@link Job%7D to run in this task.%0A   * @return the %7B@link Job%7D to run. If this method returns null, no job will be run and the task will be marked as successful.%0A   */%0A
   protected 
@@ -5256,24 +5256,24 @@
     %7D%0A    %7D%0A
-
     return j
@@ -5281,10 +5281,161 @@
 b;%0A  %7D%0A%0A
+  /**%0A   * Called when a job is skipped (because %7B@link #createJob()%7D returned null).%0A   */%0A  protected void onSkippedMRJob() %7B%0A    // do nothing%0A  %7D%0A%0A
 %7D%0A

Allow skipping execution of MR job in MR tasks.