<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.solr.hadoop;

import org.apache.hadoop.mapreduce.*;
import org.apache.solr.hadoop.util.HeartBeater;
import java.io.IOException;
import java.lang.invoke.MethodHandles;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Locale;
import java.util.Set;
import java.util.concurrent.TimeUnit;

import com.google.common.collect.ImmutableMap;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.filecache.DistributedCache;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.shaded.apache.solr.client.solrj.SolrServerException;
import org.shaded.apache.solr.client.solrj.embedded.EmbeddedSolrServer;
import org.shaded.apache.solr.common.SolrInputDocument;
import org.shaded.apache.solr.core.CoreContainer;
import org.shaded.apache.solr.core.CoreDescriptor;
import org.shaded.apache.solr.core.DirectoryFactory;
import org.shaded.apache.solr.core.HdfsDirectoryFactory;
import org.shaded.apache.solr.core.SolrCore;
import org.shaded.apache.solr.core.SolrResourceLoader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class SolrRecordWriter<K, V> extends RecordWriter<K, V> {

  private static final Logger LOG = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());

  public static final String DEFAULT_CORE_NAME = "core1";
  
  public static final String SOLR_HOME_DIR =  "solr_home";

  public final static List<String> allowedConfigDirectories = new ArrayList<>(
          Arrays.asList(new String[]{"conf", "lib", "solr.xml", DEFAULT_CORE_NAME}));

  public final static Set<String> requiredConfigDirectories = new HashSet<>();

  static {
    requiredConfigDirectories.add("conf");
  }

  /**
   * Return the list of directories names that may be included in the
   * configuration data passed to the tasks.
   *
   * @return an UnmodifiableList of directory names
   */
  public static List<String> getAllowedConfigDirectories() {
    return Collections.unmodifiableList(allowedConfigDirectories);
  }

  /**
   * check if the passed in directory is required to be present in the
   * configuration data set.
   *
   * @param directory The directory to check
   * @return true if the directory is required.
   */
  public static boolean isRequiredConfigDirectory(final String directory) {
    return requiredConfigDirectories.contains(directory);
  }

  /**
   * The path that the final index will be written to
   */
  /**
   * The location in a local temporary directory that the index is built in.
   */

  private final HeartBeater heartBeater;
  private final BatchWriter batchWriter;
  private final List<SolrInputDocument> batch;
  private final int batchSize;
  private long numDocsWritten = 0;
  private long nextLogTime = System.nanoTime();
  
  private Path outputShardDir;

  private static final HashMap<TaskID, Reducer<?, ?, ?, ?>.Context> contextMap = new HashMap<>();

  public SolrRecordWriter(TaskAttemptContext context, Path outputShardDir, int batchSize) throws IOException {
    this.batchSize = batchSize;
    this.batch = new ArrayList<>(batchSize);
    Configuration conf = context.getConfiguration();

    // setLogLevel("org.apache.solr.core", "WARN");
    // setLogLevel("org.apache.solr.update", "WARN");
    heartBeater = new HeartBeater(context);
    try {
      heartBeater.needHeartBeat();

      this.outputShardDir = outputShardDir;
      FileSystem fs = outputShardDir.getFileSystem(conf);
      EmbeddedSolrServer solr = createEmbeddedSolrServer(conf, outputShardDir);
      batchWriter = new BatchWriter(solr, batchSize,
              context.getTaskAttemptID().getTaskID(),
              SolrOutputFormat.getSolrWriterThreadCount(conf),
              SolrOutputFormat.getSolrWriterQueueSize(conf));

      Path outputSolrHomeDir = new Path(outputShardDir, SOLR_HOME_DIR);
      fs.copyFromLocalFile(new Path(solr.getCoreContainer().getSolrHome()), outputSolrHomeDir);

    } finally {
      heartBeater.cancelHeartBeat();
    }
  }

  public static EmbeddedSolrServer createEmbeddedSolrServerWithHome(Configuration conf, Path outputShardDir, Path solrHomeDir)
          throws IOException {
    FileSystem fs = outputShardDir.getFileSystem(conf);
    
    LOG.info("Creating embedded Solr server with solrHomeDir: " + solrHomeDir + ", fs: " + fs + ", outputShardDir: " + outputShardDir);

    Path solrDataDir = new Path(outputShardDir, "data");
    
    String dataDirStr = solrDataDir.toUri().toString();

    SolrResourceLoader loader = new SolrResourceLoader(Paths.get(solrHomeDir.toString()), null, null);

    LOG.info(String
            .format(Locale.ENGLISH,
                    "Constructed instance information solr.home %s (%s), instance dir %s, conf dir %s, writing index to solr.data.dir %s, with permdir %s",
                    solrHomeDir, solrHomeDir.toUri(), loader.getInstancePath(),
                    loader.getConfigDir(), dataDirStr, outputShardDir));

    // TODO: This is fragile and should be well documented
    System.setProperty("solr.directoryFactory", HdfsDirectoryFactory.class.getName());
    System.setProperty("solr.lock.type", DirectoryFactory.LOCK_TYPE_HDFS);
    System.setProperty("solr.hdfs.nrtcachingdirectory.enable", "false");
    System.setProperty("solr.hdfs.blockcache.enabled", "false");
    System.setProperty("solr.hdfs.blockcache.read.enabled", "false");
    System.setProperty("solr.autoCommit.maxTime", "600000");
    System.setProperty("solr.autoSoftCommit.maxTime", "-1");

    CoreContainer container = new CoreContainer(loader);
    container.load();
    SolrCore core;
    core = container.create(DEFAULT_CORE_NAME, ImmutableMap.of(CoreDescriptor.CORE_DATADIR, dataDirStr));

    if (!(core.getDirectoryFactory() instanceof HdfsDirectoryFactory)) {
      throw new UnsupportedOperationException(
              "Invalid configuration with : " + core.getDirectoryFactory().getClass().getName() + ". Currently, the only DirectoryFactory supported is "
              + HdfsDirectoryFactory.class.getSimpleName());
    }

    EmbeddedSolrServer solr = new EmbeddedSolrServer(container, DEFAULT_CORE_NAME);
    return solr;    
  }
  
  public static EmbeddedSolrServer createEmbeddedSolrServer(Configuration conf, Path outputShardDir)
          throws IOException {
    String outputId = outputShardDir.getParent().getName();
    Path solrHomeDir = SolrRecordWriter.findSolrConfig(conf, outputId);
    return createEmbeddedSolrServerWithHome(conf, outputShardDir, solrHomeDir);
  }

  public static void incrementCounter(TaskID taskId, String groupName, String counterName, long incr) {
    Reducer<?, ?, ?, ?>.Context context = contextMap.get(taskId);
    if (context != null) {<commit_msg>Get around write access error<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.solr.hadoop;

import org.apache.hadoop.fs.*;
import org.apache.hadoop.mapreduce.*;
import org.apache.solr.hadoop.util.HeartBeater;
import java.io.IOException;
import java.lang.invoke.MethodHandles;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Locale;
import java.util.Set;
import java.util.concurrent.TimeUnit;

import com.google.common.collect.ImmutableMap;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.filecache.DistributedCache;
import org.shaded.apache.solr.client.solrj.SolrServerException;
import org.shaded.apache.solr.client.solrj.embedded.EmbeddedSolrServer;
import org.shaded.apache.solr.common.SolrException;
import org.shaded.apache.solr.common.SolrInputDocument;
import org.shaded.apache.solr.core.CoreContainer;
import org.shaded.apache.solr.core.CoreDescriptor;
import org.shaded.apache.solr.core.DirectoryFactory;
import org.shaded.apache.solr.core.HdfsDirectoryFactory;
import org.shaded.apache.solr.core.SolrCore;
import org.shaded.apache.solr.core.SolrResourceLoader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class SolrRecordWriter<K, V> extends RecordWriter<K, V> {

  private static final Logger LOG = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());

  public static final String DEFAULT_CORE_NAME = "core1";
  
  public static final String SOLR_HOME_DIR =  "solr_home";

  public final static List<String> allowedConfigDirectories = new ArrayList<>(
          Arrays.asList("conf", "lib", "solr.xml", DEFAULT_CORE_NAME));

  public final static Set<String> requiredConfigDirectories = new HashSet<>();

  static {
    requiredConfigDirectories.add("conf");
  }

  /**
   * Return the list of directories names that may be included in the
   * configuration data passed to the tasks.
   *
   * @return an UnmodifiableList of directory names
   */
  public static List<String> getAllowedConfigDirectories() {
    return Collections.unmodifiableList(allowedConfigDirectories);
  }

  /**
   * check if the passed in directory is required to be present in the
   * configuration data set.
   *
   * @param directory The directory to check
   * @return true if the directory is required.
   */
  public static boolean isRequiredConfigDirectory(final String directory) {
    return requiredConfigDirectories.contains(directory);
  }

  /**
   * The path that the final index will be written to
   */
  /**
   * The location in a local temporary directory that the index is built in.
   */

  private final HeartBeater heartBeater;
  private final BatchWriter batchWriter;
  private final List<SolrInputDocument> batch;
  private final int batchSize;
  private long numDocsWritten = 0;
  private long nextLogTime = System.nanoTime();
  private final Path outputShardDir;
  private final Path solrHomeDir;

  private static final HashMap<TaskID, Reducer<?, ?, ?, ?>.Context> contextMap = new HashMap<>();

  public SolrRecordWriter(TaskAttemptContext context, Path outputShardDir, int batchSize) throws IOException {
    this.batchSize = batchSize;
    this.batch = new ArrayList<>(batchSize);
    Configuration conf = context.getConfiguration();

    // setLogLevel("org.apache.solr.core", "WARN");
    // setLogLevel("org.apache.solr.update", "WARN");
    heartBeater = new HeartBeater(context);
    try {
      heartBeater.needHeartBeat();

      this.outputShardDir = outputShardDir;
      FileSystem fs = outputShardDir.getFileSystem(conf);
      String outputId = outputShardDir.getParent().getName();
      solrHomeDir = findAndCopySolrConf(conf, outputId);
      EmbeddedSolrServer solr = createEmbeddedSolrServerWithHome(conf, outputShardDir, solrHomeDir);
      batchWriter = new BatchWriter(solr, batchSize,
              context.getTaskAttemptID().getTaskID(),
              SolrOutputFormat.getSolrWriterThreadCount(conf),
              SolrOutputFormat.getSolrWriterQueueSize(conf));

      Path outputSolrHomeDir = new Path(outputShardDir, SOLR_HOME_DIR);
      fs.copyFromLocalFile(new Path(solr.getCoreContainer().getSolrHome()), outputSolrHomeDir);

    } finally {
      heartBeater.cancelHeartBeat();
    }
  }

  public static EmbeddedSolrServer createEmbeddedSolrServerWithHome(Configuration conf, Path outputShardDir, Path solrHomeDir)
          throws IOException {
    FileSystem fs = outputShardDir.getFileSystem(conf);
    
    LOG.info("Creating embedded Solr server with solrHomeDir: " + solrHomeDir + ", fs: " + fs + ", outputShardDir: " + outputShardDir);

    debugLs(conf, solrHomeDir);

    Path solrDataDir = new Path(outputShardDir, "data");
    
    String dataDirStr = solrDataDir.toUri().toString();

    SolrResourceLoader loader = new SolrResourceLoader(Paths.get(solrHomeDir.toString()), null, null);

    LOG.info(String
            .format(Locale.ENGLISH,
                    "Constructed instance information solr.home %s (%s), instance dir %s, conf dir %s, writing index to solr.data.dir %s, with permdir %s",
                    solrHomeDir, solrHomeDir.toUri(), loader.getInstancePath(),
                    loader.getConfigDir(), dataDirStr, outputShardDir));

    // TODO: This is fragile and should be well documented
    System.setProperty("solr.directoryFactory", HdfsDirectoryFactory.class.getName());
    System.setProperty("solr.lock.type", DirectoryFactory.LOCK_TYPE_HDFS);
    System.setProperty("solr.hdfs.nrtcachingdirectory.enable", "false");
    System.setProperty("solr.hdfs.blockcache.enabled", "false");
    System.setProperty("solr.hdfs.blockcache.read.enabled", "false");
    System.setProperty("solr.autoCommit.maxTime", "600000");
    System.setProperty("solr.autoSoftCommit.maxTime", "-1");

    CoreContainer container = new CoreContainer(loader);
    container.load();
    SolrCore core;
    core = container.create(DEFAULT_CORE_NAME, ImmutableMap.of(CoreDescriptor.CORE_DATADIR, dataDirStr));

    if (!(core.getDirectoryFactory() instanceof HdfsDirectoryFactory)) {
      throw new UnsupportedOperationException(
              "Invalid configuration with : " + core.getDirectoryFactory().getClass().getName() + ". Currently, the only DirectoryFactory supported is "
              + HdfsDirectoryFactory.class.getSimpleName());
    }

    return new EmbeddedSolrServer(container, DEFAULT_CORE_NAME);
  }

  private static void debugLs(Configuration conf, Path path) throws IOException {
    LocalFileSystem fs = FileSystem.getLocal(conf);
    RemoteIterator<LocatedFileStatus> files = fs.listFiles(path, true);
    while (files.hasNext()) {
      LocatedFileStatus file = files.next();
      LOG.info("ls: {}", file.getPath());
    }
  }

  private static Path findAndCopySolrConf(Configuration conf, String outputId) throws IOException {
    LocalFileSystem fs = FileSystem.getLocal(conf);
    Path solrHomeDir = new Path("./solrHome-" + outputId);

    // DistributedCache.getLocalCacheArchives path seems write protected now
    Path solrHomeCache = SolrRecordWriter.findSolrConfig(conf, outputId);
    LOG.info("copy solr config from {} to {}", solrHomeCache, solrHomeDir);
    fs.copyFromLocalFile(false, solrHomeCache, solrHomeDir);

    return solrHomeDir;
  }

  public static void incrementCounter(TaskID taskId, String groupName, String counterName, long incr) {
    Reducer<?, ?, ?, ?>.Context context = contextMap.get(taskId);
    if (context != null) {
{'fraction_non_alphanumeric': 0.09998748592166187, 'fraction_numerical': 0.0018771117507195595, 'mean_word_length': 4.226618705035971, 'pattern_counts': {'":': 0, '<': 30, '<?xml version=': 0, '>': 30, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 2}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1690761', 'n_tokens_mistral': 4673, 'n_tokens_neox': 4527, 'n_words': 1182}