<commit_before>package org.batfish.main;

import static org.batfish.common.plugin.PluginConsumer.DEFAULT_HEADER_LENGTH_BYTES;
import static org.batfish.common.plugin.PluginConsumer.detectFormat;

import com.google.common.base.Throwables;
import com.google.common.io.Closer;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.OutputStream;
import java.io.PushbackInputStream;
import java.io.Serializable;
import java.nio.file.DirectoryStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Map;
import java.util.Map.Entry;
import java.util.SortedMap;
import java.util.TreeMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.BiFunction;
import java.util.stream.Collectors;
import java.util.zip.GZIPInputStream;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import net.jpountz.lz4.LZ4FrameInputStream;
import net.jpountz.lz4.LZ4FrameOutputStream;
import org.batfish.common.BatfishException;
import org.batfish.common.BatfishLogger;
import org.batfish.common.BfConsts;
import org.batfish.common.Version;
import org.batfish.common.plugin.PluginConsumer.Format;
import org.batfish.common.topology.Layer1Topology;
import org.batfish.common.util.BatfishObjectMapper;
import org.batfish.common.util.CommonUtil;
import org.batfish.datamodel.Configuration;
import org.batfish.datamodel.Topology;
import org.batfish.datamodel.answers.ConvertConfigurationAnswerElement;

/** A utility class that abstracts the underlying file system storage used by {@link Batfish}. */
final class BatfishStorage {
  private final BatfishLogger _logger;
  private final Path _containerDir;
  private final BiFunction<String, Integer, AtomicInteger> _newBatch;

  /** Create a new {@link BatfishStorage} instance that uses the given root path as a container. */
  public BatfishStorage(
      Path containerDir,
      BatfishLogger logger,
      BiFunction<String, Integer, AtomicInteger> newBatch) {
    _containerDir = containerDir;
    _logger = logger;
    _newBatch = newBatch;
  }

  /**
   * Returns the compressed configuration files for the given testrig. If a serialized copy of these
   * configurations is not already present, then this function returns {@code null}.
   */
  @Nullable
  public SortedMap<String, Configuration> loadCompressedConfigurations(String testrig) {
    Path testrigDir = getTestrigDir(testrig);
    Path indepDir = testrigDir.resolve(BfConsts.RELPATH_COMPRESSED_CONFIG_DIR);
    return loadConfigurations(testrig, indepDir);
  }

  /**
   * Returns the configuration files for the given testrig. If a serialized copy of these
   * configurations is not already present, then this function returns {@code null}.
   */
  @Nullable
  public SortedMap<String, Configuration> loadConfigurations(String testrig) {
    Path testrigDir = getTestrigDir(testrig);
    Path indepDir = testrigDir.resolve(BfConsts.RELPATH_VENDOR_INDEPENDENT_CONFIG_DIR);
    return loadConfigurations(testrig, indepDir);
  }

  private SortedMap<String, Configuration> loadConfigurations(String testrig, Path indepDir) {
    // If the directory that would contain these configs does not even exist, no cache exists.
    if (!Files.exists(indepDir)) {
      _logger.debugf("Unable to load configs for %s from disk: no cache directory", testrig);
      return null;
    }

    // If the directory exists, then likely the configs exist and are useful. Still, we need to
    // confirm that they were serialized with a compatible version of Batfish first.
    if (!cachedConfigsAreCompatible(testrig)) {
      _logger.debugf(
          "Unable to load configs for %s from disk: error or incompatible version", testrig);
      return null;
    }

    _logger.info("\n*** DESERIALIZING VENDOR-INDEPENDENT CONFIGURATION STRUCTURES ***\n");
    Map<Path, String> namesByPath = new TreeMap<>();
    try (DirectoryStream<Path> stream = Files.newDirectoryStream(indepDir)) {
      for (Path serializedConfig : stream) {
        String name = serializedConfig.getFileName().toString();
        namesByPath.put(serializedConfig, name);
      }
    } catch (IOException e) {
      throw new BatfishException(
          "Error reading vendor-independent configs directory: '" + indepDir + "'", e);
    }
    try {
      return deserializeObjects(namesByPath, Configuration.class);
    } catch (BatfishException e) {
      return null;
    }
  }

  @Nullable
  public ConvertConfigurationAnswerElement loadConvertConfigurationAnswerElement(String testrig) {
    Path ccaePath = getTestrigDir(testrig).resolve(BfConsts.RELPATH_CONVERT_ANSWER_PATH);
    if (!Files.exists(ccaePath)) {
      return null;
    }
    try {
      return deserializeObject(ccaePath, ConvertConfigurationAnswerElement.class);
    } catch (BatfishException e) {
      _logger.errorf(
          "Failed to deserialize ConvertConfigurationAnswerElement: %s",
          Throwables.getStackTraceAsString(e));
      return null;
    }
  }

  public @Nullable Topology loadLegacyTopology(@Nonnull String testrig) {
    Path path =
        getTestrigDir(testrig)
            .resolve(
                Paths.get(
                    BfConsts.RELPATH_TEST_RIG_DIR, BfConsts.RELPATH_TESTRIG_LEGACY_TOPOLOGY_PATH));
    if (!Files.exists(path)) {
      return null;
    }
    _newBatch.apply("Reading legacy topology", 0);
    String topologyFileText = CommonUtil.readFile(path);
    try {
      return BatfishObjectMapper.mapper().readValue(topologyFileText, Topology.class);
    } catch (IOException e) {
      _logger.warnf(
          "Unexpected exception caught while loading legacy testrig topology for testrig %s: %s",
          testrig, Throwables.getStackTraceAsString(e));
      return null;
    }
  }

  public @Nullable Layer1Topology loadLayer1Topology(@Nonnull String testrig) {
    Path path =
        getTestrigDir(testrig)
            .resolve(
                Paths.get(
                    BfConsts.RELPATH_TEST_RIG_DIR, BfConsts.RELPATH_TESTRIG_L1_TOPOLOGY_PATH));
    if (!Files.exists(path)) {
      return null;
    }
    _newBatch.apply("Reading layer-1 topology", 0);
    String topologyFileText = CommonUtil.readFile(path);
    try {
      return BatfishObjectMapper.mapper().readValue(topologyFileText, Layer1Topology.class);
    } catch (IOException e) {
      _logger.warnf(
          "Unexpected exception caught while loading layer-1 topology for testrig %s: %s",
          testrig, Throwables.getStackTraceAsString(e));
      return null;
    }
  }

  /**
   * Stores the configurations into the compressed config path for the given testrig. Will replace
   * any previously-stored compressed configurations.
   */
  public void storeCompressedConfigurations(
      Map<String, Configuration> configurations, String testrig) {
    Path testrigDir = getTestrigDir(testrig);

    if (!testrigDir.toFile().exists() && !testrigDir.toFile().mkdirs()) {
      throw new BatfishException(
          String.format("Unable to create testrig directory '%s'", testrigDir));
    }

    Path outputDir = testrigDir.resolve(BfConsts.RELPATH_COMPRESSED_CONFIG_DIR);

    String batchName =
        String.format(
            "Serializing %s compressed configuration structures for testrig %s",
            configurations.size(), testrig);
    storeConfigurations(outputDir, batchName, configurations);<commit_msg>BatfishStorage: no empty batches (#1891)<commit_after>package org.batfish.main;

import static org.batfish.common.plugin.PluginConsumer.DEFAULT_HEADER_LENGTH_BYTES;
import static org.batfish.common.plugin.PluginConsumer.detectFormat;

import com.google.common.base.Throwables;
import com.google.common.io.Closer;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.OutputStream;
import java.io.PushbackInputStream;
import java.io.Serializable;
import java.nio.file.DirectoryStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Map;
import java.util.Map.Entry;
import java.util.SortedMap;
import java.util.TreeMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.BiFunction;
import java.util.stream.Collectors;
import java.util.zip.GZIPInputStream;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import net.jpountz.lz4.LZ4FrameInputStream;
import net.jpountz.lz4.LZ4FrameOutputStream;
import org.batfish.common.BatfishException;
import org.batfish.common.BatfishLogger;
import org.batfish.common.BfConsts;
import org.batfish.common.Version;
import org.batfish.common.plugin.PluginConsumer.Format;
import org.batfish.common.topology.Layer1Topology;
import org.batfish.common.util.BatfishObjectMapper;
import org.batfish.common.util.CommonUtil;
import org.batfish.datamodel.Configuration;
import org.batfish.datamodel.Topology;
import org.batfish.datamodel.answers.ConvertConfigurationAnswerElement;

/** A utility class that abstracts the underlying file system storage used by {@link Batfish}. */
final class BatfishStorage {
  private final BatfishLogger _logger;
  private final Path _containerDir;
  private final BiFunction<String, Integer, AtomicInteger> _newBatch;

  /** Create a new {@link BatfishStorage} instance that uses the given root path as a container. */
  public BatfishStorage(
      Path containerDir,
      BatfishLogger logger,
      BiFunction<String, Integer, AtomicInteger> newBatch) {
    _containerDir = containerDir;
    _logger = logger;
    _newBatch = newBatch;
  }

  /**
   * Returns the compressed configuration files for the given testrig. If a serialized copy of these
   * configurations is not already present, then this function returns {@code null}.
   */
  @Nullable
  public SortedMap<String, Configuration> loadCompressedConfigurations(String testrig) {
    Path testrigDir = getTestrigDir(testrig);
    Path indepDir = testrigDir.resolve(BfConsts.RELPATH_COMPRESSED_CONFIG_DIR);
    return loadConfigurations(testrig, indepDir);
  }

  /**
   * Returns the configuration files for the given testrig. If a serialized copy of these
   * configurations is not already present, then this function returns {@code null}.
   */
  @Nullable
  public SortedMap<String, Configuration> loadConfigurations(String testrig) {
    Path testrigDir = getTestrigDir(testrig);
    Path indepDir = testrigDir.resolve(BfConsts.RELPATH_VENDOR_INDEPENDENT_CONFIG_DIR);
    return loadConfigurations(testrig, indepDir);
  }

  private SortedMap<String, Configuration> loadConfigurations(String testrig, Path indepDir) {
    // If the directory that would contain these configs does not even exist, no cache exists.
    if (!Files.exists(indepDir)) {
      _logger.debugf("Unable to load configs for %s from disk: no cache directory", testrig);
      return null;
    }

    // If the directory exists, then likely the configs exist and are useful. Still, we need to
    // confirm that they were serialized with a compatible version of Batfish first.
    if (!cachedConfigsAreCompatible(testrig)) {
      _logger.debugf(
          "Unable to load configs for %s from disk: error or incompatible version", testrig);
      return null;
    }

    _logger.info("\n*** DESERIALIZING VENDOR-INDEPENDENT CONFIGURATION STRUCTURES ***\n");
    Map<Path, String> namesByPath = new TreeMap<>();
    try (DirectoryStream<Path> stream = Files.newDirectoryStream(indepDir)) {
      for (Path serializedConfig : stream) {
        String name = serializedConfig.getFileName().toString();
        namesByPath.put(serializedConfig, name);
      }
    } catch (IOException e) {
      throw new BatfishException(
          "Error reading vendor-independent configs directory: '" + indepDir + "'", e);
    }
    try {
      return deserializeObjects(namesByPath, Configuration.class);
    } catch (BatfishException e) {
      return null;
    }
  }

  @Nullable
  public ConvertConfigurationAnswerElement loadConvertConfigurationAnswerElement(String testrig) {
    Path ccaePath = getTestrigDir(testrig).resolve(BfConsts.RELPATH_CONVERT_ANSWER_PATH);
    if (!Files.exists(ccaePath)) {
      return null;
    }
    try {
      return deserializeObject(ccaePath, ConvertConfigurationAnswerElement.class);
    } catch (BatfishException e) {
      _logger.errorf(
          "Failed to deserialize ConvertConfigurationAnswerElement: %s",
          Throwables.getStackTraceAsString(e));
      return null;
    }
  }

  public @Nullable Topology loadLegacyTopology(@Nonnull String testrig) {
    Path path =
        getTestrigDir(testrig)
            .resolve(
                Paths.get(
                    BfConsts.RELPATH_TEST_RIG_DIR, BfConsts.RELPATH_TESTRIG_LEGACY_TOPOLOGY_PATH));
    if (!Files.exists(path)) {
      return null;
    }
    AtomicInteger counter = _newBatch.apply("Reading legacy topology", 1);
    String topologyFileText = CommonUtil.readFile(path);
    try {
      return BatfishObjectMapper.mapper().readValue(topologyFileText, Topology.class);
    } catch (IOException e) {
      _logger.warnf(
          "Unexpected exception caught while loading legacy testrig topology for testrig %s: %s",
          testrig, Throwables.getStackTraceAsString(e));
      return null;
    } finally {
      counter.incrementAndGet();
    }
  }

  public @Nullable Layer1Topology loadLayer1Topology(@Nonnull String testrig) {
    Path path =
        getTestrigDir(testrig)
            .resolve(
                Paths.get(
                    BfConsts.RELPATH_TEST_RIG_DIR, BfConsts.RELPATH_TESTRIG_L1_TOPOLOGY_PATH));
    if (!Files.exists(path)) {
      return null;
    }
    AtomicInteger counter = _newBatch.apply("Reading layer-1 topology", 1);
    String topologyFileText = CommonUtil.readFile(path);
    try {
      return BatfishObjectMapper.mapper().readValue(topologyFileText, Layer1Topology.class);
    } catch (IOException e) {
      _logger.warnf(
          "Unexpected exception caught while loading layer-1 topology for testrig %s: %s",
          testrig, Throwables.getStackTraceAsString(e));
      return null;
    } finally {
      counter.incrementAndGet();
    }
  }

  /**
   * Stores the configurations into the compressed config path for the given testrig. Will replace
   * any previously-stored compressed configurations.
   */
  public void storeCompressedConfigurations(
      Map<String, Configuration> configurations, String testrig) {
    Path testrigDir = getTestrigDir(testrig);

    if (!testrigDir.toFile().exists() && !testrigDir.toFile().mkdirs()) {
      throw new BatfishException(
          String.format("Unable to create testrig directory '%s'", testrigDir));
    }

    Path outputDir = testrigDir.resolve(BfConsts.RELPATH_COMPRESSED_CONFIG_DIR);

    String batchName =
        String.format(
            "Serializing %s compressed configuration structures for testrig %s",
            configurations.size(), testrig);
    storeConfigurations(outputDir, batchName, configurations);
{'fraction_non_alphanumeric': 0.09061359867330017, 'fraction_numerical': 0.001990049751243781, 'mean_word_length': 3.833600512984931, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, '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': '7392166', 'n_tokens_mistral': 4386, 'n_tokens_neox': 4279, 'n_words': 1085}