<commit_before>/**
 * 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.tez.dag.app.dag;

import static org.apache.tez.dag.app.dag.VertexState.FAILED;

import java.io.IOException;
import java.lang.reflect.UndeclaredThrowableException;
import java.security.PrivilegedExceptionAction;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;

import javax.annotation.Nullable;

import org.apache.commons.lang.exception.ExceptionUtils;
import org.apache.hadoop.classification.InterfaceAudience;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.yarn.event.EventHandler;
import org.apache.tez.common.Preconditions;
import org.apache.tez.common.ReflectionUtils;
import org.apache.tez.common.TezUtilsInternal;
import org.apache.tez.dag.api.InputDescriptor;
import org.apache.tez.dag.api.InputInitializerDescriptor;
import org.apache.tez.dag.api.RootInputLeafOutput;
import org.apache.tez.dag.api.TezException;
import org.apache.tez.dag.api.TezUncheckedException;
import org.apache.tez.dag.api.event.VertexState;
import org.apache.tez.dag.api.event.VertexStateUpdate;
import org.apache.tez.dag.api.oldrecords.TaskState;
import org.apache.tez.dag.app.AppContext;
import org.apache.tez.dag.app.dag.event.VertexEventRootInputFailed;
import org.apache.tez.dag.app.dag.event.VertexEventRootInputInitialized;
import org.apache.tez.dag.app.dag.impl.AMUserCodeException;
import org.apache.tez.dag.app.dag.impl.AMUserCodeException.Source;
import org.apache.tez.dag.app.dag.impl.TezRootInputInitializerContextImpl;
import org.apache.tez.dag.app.dag.impl.VertexImpl;
import org.apache.tez.dag.records.TezTaskID;
import org.apache.tez.dag.records.TezVertexID;
import org.apache.tez.runtime.api.Event;
import org.apache.tez.runtime.api.InputInitializer;
import org.apache.tez.runtime.api.InputInitializerContext;
import org.apache.tez.runtime.api.events.InputInitializerEvent;
import org.apache.tez.runtime.api.impl.TezEvent;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.LinkedListMultimap;
import com.google.common.collect.ListMultimap;
import com.google.common.collect.Lists;
import com.google.common.util.concurrent.ListeningExecutorService;

public class RootInputInitializerManager {

  private static final Logger LOG = LoggerFactory.getLogger(RootInputInitializerManager.class);

  @VisibleForTesting
  protected ListeningExecutorService executor;
  @SuppressWarnings("rawtypes")
  private final EventHandler eventHandler;
  private volatile boolean isStopped = false;
  private final UserGroupInformation dagUgi;
  private final StateChangeNotifier entityStateTracker;

  private final Vertex vertex;
  private final AppContext appContext;

  @VisibleForTesting
  final Map<String, InitializerWrapper> initializerMap = new ConcurrentHashMap<>();

  public RootInputInitializerManager(Vertex vertex, AppContext appContext,
                                     UserGroupInformation dagUgi, StateChangeNotifier stateTracker) {
    this.appContext = appContext;
    this.vertex = vertex;
    this.eventHandler = appContext.getEventHandler();
    this.executor = appContext.getExecService();
    this.dagUgi = dagUgi;
    this.entityStateTracker = stateTracker;
  }


  public void runInputInitializers(
          List<RootInputLeafOutput<InputDescriptor, InputInitializerDescriptor>> inputs, List<TezEvent> pendingInitializerEvents) {

    executor.submit(() -> createAndStartInitializing(inputs, pendingInitializerEvents));
  }

  private void createAndStartInitializing(List<RootInputLeafOutput<InputDescriptor, InputInitializerDescriptor>> inputs, List<TezEvent> pendingInitializerEvents) {
    String current = null;
    try {
      List<InitializerWrapper> result = new ArrayList<>();
      for (RootInputLeafOutput<InputDescriptor, InputInitializerDescriptor> each : inputs) {
        current = each.getName();
        InitializerWrapper initializer = createInitializerWrapper(each);
        initializerMap.put(each.getName(), initializer);
        registerPendingVertex(each, initializer);
        result.add(initializer);
      }
      handleInitializerEvents(pendingInitializerEvents);
      pendingInitializerEvents.clear();
      for (InitializerWrapper inputWrapper : result) {
        executor.submit(() -> runInitializerAndProcessResult(inputWrapper));
      }
    } catch (Throwable t) {
      VertexImpl vertexImpl = (VertexImpl) vertex;
      String msg = "Fail to create InputInitializerManager, " + ExceptionUtils.getStackTrace(t);
      LOG.info(msg);
      vertexImpl.finished(FAILED, VertexTerminationCause.INIT_FAILURE, msg);
      eventHandler.handle(new VertexEventRootInputFailed(vertex.getVertexId(), current,
              new AMUserCodeException(AMUserCodeException.Source.InputInitializer, t)));

    }
  }

  private void runInitializerAndProcessResult(InitializerWrapper initializer) {<commit_msg>TEZ-4204: Data race in RootInputInitializerManager (Must<NAME> via <NAME>)<commit_after>/**
 * 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.tez.dag.app.dag;

import static org.apache.tez.dag.app.dag.VertexState.FAILED;

import java.io.IOException;
import java.lang.reflect.UndeclaredThrowableException;
import java.security.PrivilegedExceptionAction;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicBoolean;

import javax.annotation.Nullable;

import org.apache.commons.lang.exception.ExceptionUtils;
import org.apache.hadoop.classification.InterfaceAudience;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.yarn.event.EventHandler;
import org.apache.tez.common.Preconditions;
import org.apache.tez.common.ReflectionUtils;
import org.apache.tez.common.TezUtilsInternal;
import org.apache.tez.dag.api.InputDescriptor;
import org.apache.tez.dag.api.InputInitializerDescriptor;
import org.apache.tez.dag.api.RootInputLeafOutput;
import org.apache.tez.dag.api.TezException;
import org.apache.tez.dag.api.TezUncheckedException;
import org.apache.tez.dag.api.event.VertexState;
import org.apache.tez.dag.api.event.VertexStateUpdate;
import org.apache.tez.dag.api.oldrecords.TaskState;
import org.apache.tez.dag.app.AppContext;
import org.apache.tez.dag.app.dag.event.VertexEventRootInputFailed;
import org.apache.tez.dag.app.dag.event.VertexEventRootInputInitialized;
import org.apache.tez.dag.app.dag.impl.AMUserCodeException;
import org.apache.tez.dag.app.dag.impl.AMUserCodeException.Source;
import org.apache.tez.dag.app.dag.impl.TezRootInputInitializerContextImpl;
import org.apache.tez.dag.app.dag.impl.VertexImpl;
import org.apache.tez.dag.records.TezTaskID;
import org.apache.tez.dag.records.TezVertexID;
import org.apache.tez.runtime.api.Event;
import org.apache.tez.runtime.api.InputInitializer;
import org.apache.tez.runtime.api.InputInitializerContext;
import org.apache.tez.runtime.api.events.InputInitializerEvent;
import org.apache.tez.runtime.api.impl.TezEvent;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.LinkedListMultimap;
import com.google.common.collect.ListMultimap;
import com.google.common.collect.Lists;
import com.google.common.util.concurrent.ListeningExecutorService;

public class RootInputInitializerManager {

  private static final Logger LOG = LoggerFactory.getLogger(RootInputInitializerManager.class);

  @VisibleForTesting
  protected ListeningExecutorService executor;
  @SuppressWarnings("rawtypes")
  private final EventHandler eventHandler;
  private volatile boolean isStopped = false;
  private final UserGroupInformation dagUgi;
  private final StateChangeNotifier entityStateTracker;

  private final Vertex vertex;
  private final AppContext appContext;

  @VisibleForTesting
  final Map<String, InitializerWrapper> initializerMap = new ConcurrentHashMap<>();

  public RootInputInitializerManager(Vertex vertex, AppContext appContext,
                                     UserGroupInformation dagUgi, StateChangeNotifier stateTracker) {
    this.appContext = appContext;
    this.vertex = vertex;
    this.eventHandler = appContext.getEventHandler();
    this.executor = appContext.getExecService();
    this.dagUgi = dagUgi;
    this.entityStateTracker = stateTracker;
  }


  public void runInputInitializers(List<RootInputLeafOutput<InputDescriptor, InputInitializerDescriptor>> inputs,
                                   List<TezEvent> pendingInitializerEvents) {
    List<InitializerWrapper> initWrappers = createInitializerWrappers(inputs);
    if (!initWrappers.isEmpty()) {
      executor.submit(() -> createAndStartInitializing(pendingInitializerEvents, initWrappers));
    }
  }

  /**
   * Create input wrappers for all inputs in parallel.
   *
   * @param inputs
   * @return
   */
  private List<InitializerWrapper> createInitializerWrappers(
          List<RootInputLeafOutput<InputDescriptor, InputInitializerDescriptor>> inputs) {
    String current = null;
    final List<InitializerWrapper> result = Collections.synchronizedList(new ArrayList<>());
    try {
      final List<Future<Void>> fResults = new ArrayList<>();
      for (RootInputLeafOutput<InputDescriptor, InputInitializerDescriptor> each : inputs) {
        current = each.getName();
        fResults.add(executor.submit(() -> {
          InitializerWrapper initializer = createInitializerWrapper(each);
          initializerMap.put(each.getName(), initializer);
          registerPendingVertex(each, initializer);
          result.add(initializer);
          return null;
        }));
      }
      for(Future<Void> f : fResults) {
        f.get();
      }
    } catch (InterruptedException | ExecutionException t) {
      failVertex(t, current);
    }
    return result;
  }

  void failVertex(Throwable t, String inputName) {
    VertexImpl vertexImpl = (VertexImpl) vertex;
    String msg = "Fail to create InputInitializerManager, " + ExceptionUtils.getStackTrace(t);
    LOG.info(msg);
    vertexImpl.finished(FAILED, VertexTerminationCause.INIT_FAILURE, msg);
    eventHandler.handle(new VertexEventRootInputFailed(vertex.getVertexId(), inputName,
        new AMUserCodeException(AMUserCodeException.Source.InputInitializer, t)));
  }

  /**
   * Start initializers in parallel.
   *
   * @param pendingInitializerEvents
   * @param result
   */
  private void createAndStartInitializing(List<TezEvent> pendingInitializerEvents, List<InitializerWrapper> result) {
    handleInitializerEvents(pendingInitializerEvents);
    pendingInitializerEvents.clear();
    for (InitializerWrapper inputWrapper : result) {
      executor.submit(() -> runInitializerAndProcessResult(inputWrapper));
    }
  }

  private void runInitializerAndProcessResult(InitializerWrapper initializer) {
{'fraction_non_alphanumeric': 0.09254696643055128, 'fraction_numerical': 0.0012319063751154912, 'mean_word_length': 4.824663677130045, 'pattern_counts': {'":': 0, '<': 34, '<?xml version=': 0, '>': 39, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2147440', 'n_tokens_mistral': 3648, 'n_tokens_neox': 3641, 'n_words': 792}