lucene/benchmark/src/java/org/apache/lucene/benchmark/byTask/feeds/ContentItemsSource.java

/*
 * 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.lucene.benchmark.byTask.feeds;

import java.io.Closeable;
import java.io.IOException;
import java.nio.file.FileVisitResult;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.SimpleFileVisitor;
import java.nio.file.attribute.BasicFileAttributes;
import java.util.ArrayList;
import org.apache.lucene.benchmark.byTask.utils.Config;
import org.apache.lucene.benchmark.byTask.utils.Format;

/**
 * Base class for source of data for benchmarking
 *
 * <p>Keeps track of various statistics, such as how many data items were generated, size in bytes
 * etc.
 *
 * <p>Supports the following configuration parameters:
 *
 * <ul>
 *   <li><b>content.source.forever</b> - specifies whether to generate items forever
 *       (<b>default=true</b>).
 *   <li><b>content.source.verbose</b> - specifies whether messages should be output by the content
 *       source (<b>default=false</b>).
 *   <li><b>content.source.encoding</b> - specifies which encoding to use when reading the files of
 *       that content source. Certain implementations may define a default value if this parameter
 *       is not specified. (<b>default=null</b>).
 *   <li><b>content.source.log.step</b> - specifies for how many items a message should be logged.
 *       If set to 0 it means no logging should occur. <b>NOTE:</b> if verbose is set to false,
 *       logging should not occur even if logStep is not 0 (<b>default=0</b>).
 * </ul>
 */
public abstract class ContentItemsSource implements Closeable {

  private long bytesCount;
  private long totalBytesCount;
  private int itemCount;
  private int totalItemCount;
  private Config config;

  private int lastPrintedNumUniqueTexts = 0;
  private long lastPrintedNumUniqueBytes = 0;
  private int printNum = 0;

  protected boolean forever;
  protected int logStep;
  protected boolean verbose;
  protected String encoding;

  /** update count of bytes generated by this source */
  protected final synchronized void addBytes(long numBytes) {
    bytesCount += numBytes;
    totalBytesCount += numBytes;
  }

  /** update count of items generated by this source */
  protected final synchronized void addItem() {
    ++itemCount;
    ++totalItemCount;
  }

  /**
   * A convenience method for collecting all the files of a content source from a given directory.
   * The collected {@link Path} instances are stored in the given <code>files</code>.
   */
  protected final void collectFiles(Path dir, final ArrayList<Path> files) throws IOException {
    Files.walkFileTree(
        dir,
        new SimpleFileVisitor<Path>() {
          @Override
          public FileVisitResult visitFile(Path file, BasicFileAttributes attrs)
              throws IOException {
            if (Files.isReadable(file)) {
              files.add(file.toRealPath());
            }
            return FileVisitResult.CONTINUE;
          }
        });
  }

  /**
   * Returns true whether it's time to log a message (depending on verbose and the number of items
   * generated).
   */
  protected final boolean shouldLog() {
    return verbose && logStep > 0 && itemCount % logStep == 0;
  }

  /** Called when reading from this content source is no longer required. */
  @Override
  public abstract void close() throws IOException;

  /** Returns the number of bytes generated since last reset. */
  public final long getBytesCount() {
    return bytesCount;
  }

  /** Returns the number of generated items since last reset. */
  public final int getItemsCount() {
    return itemCount;
  }

  public final Config getConfig() {
    return config;
  }

  /** Returns the total number of bytes that were generated by this source. */
  public final long getTotalBytesCount() {
    return totalBytesCount;
  }

  /** Returns the total number of generated items. */
  public final int getTotalItemsCount() {
    return totalItemCount;
  }

  /**
   * Resets the input for this content source, so that the test would behave as if it was just
   * started, input-wise.
   *
   * <p><b>NOTE:</b> the default implementation resets the number of bytes and items generated since
   * the last reset, so it's important to call super.resetInputs in case you override this method.
   */
  public void resetInputs() throws IOException {
    bytesCount = 0;
    itemCount = 0;
  }

  /**
   * Sets the {@link Config} for this content source. If you override this method, you must call
   * super.setConfig.
   */
  public void setConfig(Config config) {
    this.config = config;
    forever = config.get("content.source.forever", true);
    logStep = config.get("content.source.log.step", 0);
    verbose = config.get("content.source.verbose", false);
    encoding = config.get("content.source.encoding", null);
  }

  public void printStatistics(String itemsName) {
    if (!verbose) {
      return;
    }
    boolean print = false;
    String col = "                  ";
    StringBuilder sb = new StringBuilder();
    String newline = System.getProperty("line.separator");
    sb.append("------------> ")
        .append(getClass().getSimpleName())
        .append(" statistics (")
        .append(printNum)
        .append("): ")
        .append(newline);
    int nut = getTotalItemsCount();
    if (nut > lastPrintedNumUniqueTexts) {
      print = true;
      sb.append("total count of ")
          .append(itemsName)
          .append(": ")
          .append(Format.format(0, nut, col))
          .append(newline);
      lastPrintedNumUniqueTexts = nut;
    }
    long nub = getTotalBytesCount();
    if (nub > lastPrintedNumUniqueBytes) {
      print = true;
      sb.append("total bytes of ")
          .append(itemsName)
          .append(": ")
          .append(Format.format(0, nub, col))
          .append(newline);
      lastPrintedNumUniqueBytes = nub;
    }
    if (getItemsCount() > 0) {
      print = true;
      sb.append("num ")
          .append(itemsName)
          .append(" added since last inputs reset:   ")
          .append(Format.format(0, getItemsCount(), col))
          .append(newline);
      sb.append("total bytes added for ")
          .append(itemsName)
          .append(" since last inputs reset: ")
          .append(Format.format(0, getBytesCount(), col))
          .append(newline);
    }
    if (print) {
      System.out.println(sb.append(newline).toString());
      printNum++;
    }
  }
}