alluxio-master-2.x/table/server/master/src/main/java/alluxio/master/table/transform/TransformManager.java
/*
 * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
 * (the "License"). You may not use this work except in compliance with the License, which is
 * available at www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied, as more fully set forth in the License.
 *
 * See the NOTICE file distributed with this work for information regarding copyright ownership.
 */

package alluxio.master.table.transform;

import alluxio.client.job.JobMasterClient;
import alluxio.collections.Pair;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.exception.ExceptionMessage;
import alluxio.exception.status.NotFoundException;
import alluxio.exception.status.UnavailableException;
import alluxio.heartbeat.FixedIntervalSupplier;
import alluxio.heartbeat.HeartbeatContext;
import alluxio.heartbeat.HeartbeatExecutor;
import alluxio.heartbeat.HeartbeatThread;
import alluxio.job.JobConfig;
import alluxio.job.wire.JobInfo;
import alluxio.job.wire.Status;
import alluxio.job.workflow.composite.CompositeConfig;
import alluxio.master.journal.DelegatingJournaled;
import alluxio.master.journal.JournalContext;
import alluxio.master.journal.Journaled;
import alluxio.master.journal.checkpoint.CheckpointName;
import alluxio.master.table.AlluxioCatalog;
import alluxio.master.table.Partition;
import alluxio.proto.journal.Journal;
import alluxio.proto.journal.Journal.JournalEntry;
import alluxio.proto.journal.Table.AddTransformJobInfoEntry;
import alluxio.proto.journal.Table.RemoveTransformJobInfoEntry;
import alluxio.resource.CloseableIterator;
import alluxio.security.user.UserState;
import alluxio.table.common.Layout;
import alluxio.table.common.transform.TransformDefinition;
import alluxio.table.common.transform.TransformPlan;

import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.google.common.collect.Iterators;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * Manages transformations.
 *
 * It executes a transformation by submitting a job to the job service.
 *
 * A background thread periodically checks whether the job succeeds, fails, or is still running,
 * the period is configurable by {@link PropertyKey#TABLE_TRANSFORM_MANAGER_JOB_MONITOR_INTERVAL}.
 *
 * It keeps information about all running transformations not only in memory,
 * but also in the journal,so even if it is restarted, the information of the previous running
 * transformations is not lost.
 *
 * It keeps history of all succeeded or failed transformations in memory for a time period which is
 * configurable by {@link PropertyKey#TABLE_TRANSFORM_MANAGER_JOB_HISTORY_RETENTION_TIME}.
 *
 * If the job succeeds, it updates the {@link Partition}'s location by
 * {@link AlluxioCatalog#completeTransformTable(JournalContext, String, String, String, Map)}.
 */
public class TransformManager implements DelegatingJournaled {
  private static final Logger LOG = LoggerFactory.getLogger(TransformManager.class);
  private static final long INVALID_JOB_ID = -1;

  /** The function used to create a {@link JournalContext}. */
  private final ThrowingSupplier<JournalContext, UnavailableException> mCreateJournalContext;

  /** The catalog. */
  private final AlluxioCatalog mCatalog;

  /** The client to talk to job master. */
  private final JobMasterClient mJobMasterClient;

  /**
   * A cache from job ID to the job's information.
   * It contains history of finished jobs, no matter whether it succeeds or fails.
   * Each history is kept for a configurable time period.
   * This is not journaled, so after restarting TableMaster, the job history is lost.
   */
  private final Cache<Long, TransformJobInfo> mJobHistory = CacheBuilder.newBuilder()
      .expireAfterWrite(Configuration.getMs(
          PropertyKey.TABLE_TRANSFORM_MANAGER_JOB_HISTORY_RETENTION_TIME),
          TimeUnit.MILLISECONDS)
      .build();

  /** Journaled state. */
  private final State mState = new State();

  /**
   * An internal job master client will be created.
   *
   * @param createJournalContext journal context creator
   * @param catalog the table catalog
   * @param jobMasterClient the job master client
   */
  public TransformManager(
      ThrowingSupplier<JournalContext, UnavailableException> createJournalContext,
      AlluxioCatalog catalog, JobMasterClient jobMasterClient) {
    mCreateJournalContext = createJournalContext;
    mCatalog = catalog;
    mJobMasterClient = jobMasterClient;
  }

  /**
   * Starts background heartbeats.
   * The heartbeats are stopped when unchecked error happens inside the heartbeats,
   * or the hearbeat threads are interrupted,
   * or the provided executor service is stopped.
   * This class will not stop the executor service.
   *
   * @param executorService the executor service for executing heartbeat threads
   * @param userState the user state for the heartbeat
   */
  public void start(ExecutorService executorService, UserState userState) {
    executorService.submit(
        new HeartbeatThread(HeartbeatContext.MASTER_TABLE_TRANSFORMATION_MONITOR, new JobMonitor(),
            () -> new FixedIntervalSupplier(
                Configuration.getMs(PropertyKey.TABLE_TRANSFORM_MANAGER_JOB_MONITOR_INTERVAL)),
            Configuration.global(), userState));
  }

  /**
   * Executes the plans for the table transformation.
   *
   * This method executes a transformation job with type{@link CompositeConfig},
   * the transformation job concurrently executes the plans,
   * each plan has a list of jobs to be executed sequentially.
   *
   * This method triggers the execution of the transformation job asynchronously without waiting
   * for it to finish. The returned job ID can be used to poll the job service for the status of
   * this transformation.
   *
   * @param dbName the database name
   * @param tableName the table name
   * @param definition the parsed transformation definition
   * @return the job ID for the transformation job
   * @throws IOException when there is an ongoing transformation on the table, or the transformation
   *    job fails to be started, or all partitions of the table have been transformed with the same
   *    definition
   */
  public long execute(String dbName, String tableName, TransformDefinition definition)
      throws IOException {
    List<TransformPlan> plans = mCatalog.getTransformPlan(dbName, tableName, definition);
    if (plans.isEmpty()) {
      throw new IOException(ExceptionMessage.TABLE_ALREADY_TRANSFORMED.getMessage(
          dbName, tableName, definition.getDefinition()));
    }
    Pair<String, String> dbTable = new Pair<>(dbName, tableName);
    // Atomically try to acquire the permit to execute the transformation job.
    // This PUT does not need to be journaled, because if this PUT succeeds and master crashes,
    // when master restarts, this temporary placeholder entry will not exist, which is correct
    // behavior.
    Long existingJobId = mState.acquireJobPermit(dbTable);
    if (existingJobId != null) {
      if (existingJobId == INVALID_JOB_ID) {
        throw new IOException("A concurrent transformation request is going to be executed");
      } else {
        throw new IOException(ExceptionMessage.TABLE_BEING_TRANSFORMED
            .getMessage(existingJobId.toString(), tableName, dbName));
      }
    }

    ArrayList<JobConfig> concurrentJobs = new ArrayList<>(plans.size());
    for (TransformPlan plan : plans) {
      concurrentJobs.add(new CompositeConfig(plan.getJobConfigs(), true));
    }
    CompositeConfig transformJob = new CompositeConfig(concurrentJobs, false);

    long jobId;
    try {
      jobId = mJobMasterClient.run(transformJob);
    } catch (IOException e) {
      // The job fails to start, clear the acquired permit for execution.
      // No need to journal this REMOVE, if master crashes, when it restarts, the permit placeholder
      // entry will not exist any more, which is correct behavior.
      mState.releaseJobPermit(dbTable);
      String error = String.format("Fails to start job to transform table %s in database %s",
          tableName, dbName);
      LOG.error(error, e);
      throw new IOException(error, e);
    }

    Map<String, Layout> transformedLayouts = new HashMap<>(plans.size());
    for (TransformPlan plan : plans) {
      transformedLayouts.put(plan.getBaseLayout().getSpec(), plan.getTransformedLayout());
    }
    AddTransformJobInfoEntry journalEntry = AddTransformJobInfoEntry.newBuilder()
        .setDbName(dbName)
        .setTableName(tableName)
        .setDefinition(definition.getDefinition())
        .setJobId(jobId)
        .putAllTransformedLayouts(Maps.transformValues(transformedLayouts, Layout::toProto))
        .build();
    try (JournalContext journalContext = mCreateJournalContext.apply()) {
      applyAndJournal(journalContext, Journal.JournalEntry.newBuilder()
          .setAddTransformJobInfo(journalEntry).build());
    }
    return jobId;
  }

  /**
   * @param jobId the job ID
   * @return the job information
   */
  public Optional<TransformJobInfo> getTransformJobInfo(long jobId) {
    TransformJobInfo job = mState.getRunningJob(jobId);
    if (job == null) {
      job = mJobHistory.getIfPresent(jobId);
    }
    return job == null ? Optional.empty() : Optional.of(job);
  }

  /**
   * @return all transformation jobs, including the running jobs and the kept finished jobs,
   *    sorted by job ID in increasing order
   */
  public List<TransformJobInfo> getAllTransformJobInfo() {
    List<TransformJobInfo> jobs = Lists.newArrayList(mJobHistory.asMap().values());
    jobs.addAll(mState.getRunningJobs());
    jobs.sort(Comparator.comparing(TransformJobInfo::getJobId));
    return jobs;
  }

  @Override
  public Journaled getDelegate() {
    return mState;
  }

  /**
   * Periodically polls the job service to monitor the status of the running transformation jobs,
   * if a transformation job succeeds, then update the Table partitions' layouts.
   */
  private final class JobMonitor implements HeartbeatExecutor {
    private void onFinish(TransformJobInfo job) {
      mJobHistory.put(job.getJobId(), job);
      RemoveTransformJobInfoEntry journalEntry = RemoveTransformJobInfoEntry.newBuilder()
          .setDbName(job.getDb())
          .setTableName(job.getTable())
          .build();
      try (JournalContext journalContext = mCreateJournalContext.apply()) {
        applyAndJournal(journalContext, Journal.JournalEntry.newBuilder()
            .setRemoveTransformJobInfo(journalEntry).build());
      } catch (UnavailableException e) {
        LOG.error("Failed to create journal for RemoveTransformJobInfo for database {} table {}",
            job.getDb(), job.getTable());
      }
    }

    /**
     * Handle the cases where a job fails, is cancelled, or job ID not found.
     *
     * @param job the transformation job
     * @param status the job status
     * @param error the job error message
     */
    private void handleJobError(TransformJobInfo job, Status status, String error) {
      job.setJobStatus(status);
      job.setJobErrorMessage(error);
      onFinish(job);
    }

    /**
     * Handle the case where a job is completed.
     *
     * @param job the transformation job
     */
    private void handleJobSuccess(TransformJobInfo job) {
      try (JournalContext journalContext = mCreateJournalContext.apply()) {
        mCatalog.completeTransformTable(journalContext, job.getDb(), job.getTable(),
            job.getDefinition(), job.getTransformedLayouts());
        job.setJobStatus(Status.COMPLETED);
      } catch (IOException e) {
        String error = String.format("Failed to update partition layouts for database %s table %s",
            job.getDb(), job.getTable());
        LOG.error(error);
        job.setJobStatus(Status.FAILED);
        job.setJobErrorMessage(error);
      }
      onFinish(job);
    }

    @Override
    public void heartbeat(long timeLimitMs) throws InterruptedException {
      for (TransformJobInfo job : mState.getRunningJobs()) {
        if (Thread.currentThread().isInterrupted()) {
          throw new InterruptedException("TransformManager's heartbeat was interrupted");
        }
        long jobId = job.getJobId();
        try {
          LOG.debug("Polling for status of transformation job {}", jobId);
          JobInfo jobInfo = mJobMasterClient.getJobStatus(jobId);
          switch (jobInfo.getStatus()) {
            case FAILED: // fall through
            case CANCELED: // fall through
              LOG.warn("Transformation job {} for database {} table {} {}: {}", jobId,
                  job.getDb(), job.getTable(),
                  jobInfo.getStatus() == Status.FAILED ? "failed" : "canceled",
                  jobInfo.getErrorMessage());
              handleJobError(job, jobInfo.getStatus(), jobInfo.getErrorMessage());
              break;
            case COMPLETED:
              LOG.info("Transformation job {} for database {} table {} succeeds", jobId,
                  job.getDb(), job.getTable());
              handleJobSuccess(job);
              break;
            case RUNNING: // fall through
            case CREATED:
              break;
            default:
              throw new IllegalStateException("Unrecognized job status: " + jobInfo.getStatus());
          }
        } catch (NotFoundException e) {
          String error = ExceptionMessage.TRANSFORM_JOB_ID_NOT_FOUND_IN_JOB_SERVICE.getMessage(
              jobId, job.getDb(), job.getTable(), e.getMessage());
          LOG.warn(error);
          handleJobError(job, Status.FAILED, error);
        } catch (IOException e) {
          LOG.error("Failed to get status for job (id={})", jobId, e);
        }
      }
    }

    @Override
    public void close() {
      // EMPTY
    }
  }

  /**
   * Journaled state.
   *
   * The internal data structure should never be exposed outside of the class,
   * all changes to the internal state should happen through applying journal entries.
   */
  private final class State implements Journaled {
    /**
     * Map from (db, table) to the ID of the running transformation job.
     * When trying to start a transformation on a table, a placeholder ID is put first.
     * When removing a job, first remove from {@link #mRunningJobs}, then remove from this map,
     * otherwise, there might be concurrent transformations running on the same table.
     */
    private final ConcurrentHashMap<Pair<String, String>, Long> mRunningJobIds =
        new ConcurrentHashMap<>();
    /**
     * Map from job ID to job info.
     */
    private final ConcurrentHashMap<Long, TransformJobInfo> mRunningJobs =
        new ConcurrentHashMap<>();

    /**
     * @return all running jobs
     */
    public Collection<TransformJobInfo> getRunningJobs() {
      return mRunningJobs.values();
    }

    /**
     * @param jobId the job ID
     * @return corresponding job or null if not exists
     */
    public TransformJobInfo getRunningJob(long jobId) {
      return mRunningJobs.get(jobId);
    }

    /**
     * Acquires a permit for transforming a table.
     *
     * @param dbTable a pair of database and table name
     * @return the ID of the existing transformation on the table, or null
     */
    public Long acquireJobPermit(Pair<String, String> dbTable) {
      return mRunningJobIds.putIfAbsent(dbTable, INVALID_JOB_ID);
    }

    /**
     * Releases the previously acquired permit for transforming a table.
     *
     * @param dbTable a pair of database and table name
     */
    public void releaseJobPermit(Pair<String, String> dbTable) {
      mRunningJobIds.remove(dbTable);
    }

    @Override
    public boolean processJournalEntry(JournalEntry entry) {
      if (entry.hasAddTransformJobInfo()) {
        applyAddTransformJobInfoEntry(entry.getAddTransformJobInfo());
      } else if (entry.hasRemoveTransformJobInfo()) {
        applyRemoveTransformJobInfoEntry(entry.getRemoveTransformJobInfo());
      } else {
        return false;
      }
      return true;
    }

    private void applyAddTransformJobInfoEntry(AddTransformJobInfoEntry entry) {
      Map<String, alluxio.grpc.table.Layout> layouts = entry.getTransformedLayoutsMap();
      Map<String, Layout> transformedLayouts = Maps.transformValues(layouts,
          layout -> mCatalog.getLayoutRegistry().create(layout));
      TransformJobInfo job = new TransformJobInfo(entry.getDbName(), entry.getTableName(),
          entry.getDefinition(), entry.getJobId(), transformedLayouts);
      mRunningJobIds.put(job.getDbTable(), job.getJobId());
      mRunningJobs.put(job.getJobId(), job);
    }

    private void applyRemoveTransformJobInfoEntry(RemoveTransformJobInfoEntry entry) {
      Pair<String, String> dbTable = new Pair<>(entry.getDbName(), entry.getTableName());
      long jobId = mRunningJobIds.get(dbTable);
      mRunningJobs.remove(jobId);
      mRunningJobIds.remove(dbTable);
    }

    @Override
    public void resetState() {
      mRunningJobs.clear();
      mRunningJobIds.clear();
      mJobHistory.invalidateAll();
      mJobHistory.cleanUp();
    }

    @Override
    public CloseableIterator<JournalEntry> getJournalEntryIterator() {
      return CloseableIterator.noopCloseable(
          Iterators.transform(mRunningJobs.values().iterator(), job -> {
            AddTransformJobInfoEntry journal = AddTransformJobInfoEntry.newBuilder()
                .setDbName(job.getDb())
                .setTableName(job.getTable())
                .setDefinition(job.getDefinition())
                .setJobId(job.getJobId())
                .putAllTransformedLayouts(Maps.transformValues(
                    job.getTransformedLayouts(), Layout::toProto))
                .build();
            return JournalEntry.newBuilder().setAddTransformJobInfo(journal).build();
          }));
    }

    @Override
    public CheckpointName getCheckpointName() {
      return CheckpointName.TABLE_MASTER_TRANSFORM_MANAGER;
    }
  }

  /**
   * A supplier with return type R that might throw exception E.
   *
   * @param <R> the return type
   * @param <E> the exception type
   */
  @FunctionalInterface
  public interface ThrowingSupplier<R, E extends Throwable> {
    /**
     * @return the result
     */
    R apply() throws E;
  }
}
