<commit_before>package spark.storage

import java.io.{File, FileOutputStream, OutputStream, RandomAccessFile}
import java.nio.ByteBuffer
import java.nio.channels.FileChannel
import java.nio.channels.FileChannel.MapMode
import java.util.{Random, Date}
import java.text.SimpleDateFormat

import scala.collection.mutable.ArrayBuffer

import it.unimi.dsi.fastutil.io.FastBufferedOutputStream

import spark.Utils
import spark.executor.ExecutorExitCode
import spark.serializer.{Serializer, SerializationStream}
import spark.Logging
import spark.network.netty.ShuffleSender
import spark.network.netty.PathResolver


/**
 * Stores BlockManager blocks on disk.
 */
private class DiskStore(blockManager: BlockManager, rootDirs: String)
  extends BlockStore(blockManager) with Logging {

  class DiskBlockObjectWriter(blockId: String, serializer: Serializer, bufferSize: Int)
    extends BlockObjectWriter(blockId) {

    private val f: File = createFile(blockId /*, allowAppendExisting */)

    // The file channel, used for repositioning / truncating the file.
    private var channel: FileChannel = null
    private var bs: OutputStream = null
    private var objOut: SerializationStream = null
    private var lastValidPosition = 0L

    override def open(): DiskBlockObjectWriter = {
      val fos = new FileOutputStream(f, true)
      channel = fos.getChannel()
      bs = blockManager.wrapForCompression(blockId, new FastBufferedOutputStream(fos))
      objOut = serializer.newInstance().serializeStream(bs)
      this
    }

    override def close() {
      objOut.close()
      bs.close()
      channel = null
      bs = null
      objOut = null
      // Invoke the close callback handler.
      super.close()
    }

    override def isOpen: Boolean = objOut != null

    // Flush the partial writes, and set valid length to be the length of the entire file.
    // Return the number of bytes written for this commit.
    override def commit(): Long = {
      // NOTE: Flush the serializer first and then the compressed/buffered output stream
      objOut.flush()
      bs.flush()
      val prevPos = lastValidPosition
      lastValidPosition = channel.position()
      lastValidPosition - prevPos
    }

    override def revertPartialWrites() {
      // Discard current writes. We do this by flushing the outstanding writes and
      // truncate the file to the last valid position.
      objOut.flush()
      bs.flush()
      channel.truncate(lastValidPosition)
    }

    override def write(value: Any) {
      objOut.writeObject(value)
    }

    override def size(): Long = lastValidPosition
  }

  val MAX_DIR_CREATION_ATTEMPTS: Int = 10
  val subDirsPerLocalDir = System.getProperty("spark.diskStore.subDirectories", "64").toInt

  var shuffleSender : ShuffleSender = null
  // Create one local directory for each path mentioned in spark.local.dir; then, inside this
  // directory, create multiple subdirectories that we will hash files into, in order to avoid
  // having really large inodes at the top level.
  val localDirs = createLocalDirs()
  val subDirs = Array.fill(localDirs.length)(new Array[File](subDirsPerLocalDir))

  addShutdownHook()

  def getBlockWriter(blockId: String, serializer: Serializer, bufferSize: Int)
    : BlockObjectWriter = {
    new DiskBlockObjectWriter(blockId, serializer, bufferSize)
  }


  override def getSize(blockId: String): Long = {
    getFile(blockId).length()
  }

  override def putBytes(blockId: String, _bytes: ByteBuffer, level: StorageLevel) {
    // So that we do not modify the input offsets !
    // duplicate does not copy buffer, so inexpensive
    val bytes = _bytes.duplicate()
    logDebug("Attempting to put block " + blockId)
    val startTime = System.currentTimeMillis
    val file = createFile(blockId)
    val channel = new RandomAccessFile(file, "rw").getChannel()
    while (bytes.remaining > 0) {
      channel.write(bytes)
    }
    channel.close()
    val finishTime = System.currentTimeMillis
    logDebug("Block %s stored as %s file on disk in %d ms".format(
      blockId, Utils.memoryBytesToString(bytes.limit), (finishTime - startTime)))
  }

  private def getFileBytes(file: File): ByteBuffer = {
    val length = file.length()
    val channel = new RandomAccessFile(file, "r").getChannel()
    val buffer = try {
      channel.map(MapMode.READ_ONLY, 0, length)
    } finally {
      channel.close()
    }

    buffer
  }

  override def putValues(
      blockId: String,
      values: ArrayBuffer[Any],
      level: StorageLevel,
      returnValues: Boolean)
    : PutResult = {

    logDebug("Attempting to write values for block " + blockId)
    val startTime = System.currentTimeMillis
    val file = createFile(blockId)
    val fileOut = blockManager.wrapForCompression(blockId,
      new FastBufferedOutputStream(new FileOutputStream(file)))
    val objOut = blockManager.defaultSerializer.newInstance().serializeStream(fileOut)
    objOut.writeAll(values.iterator)
    objOut.close()
    val length = file.length()

    val timeTaken = System.currentTimeMillis - startTime
    logDebug("Block %s stored as %s file on disk in %d ms".format(
      blockId, Utils.memoryBytesToString(length), timeTaken))

    if (returnValues) {
      // Return a byte buffer for the contents of the file
      val buffer = getFileBytes(file)
      PutResult(length, Right(buffer))
    } else {
      PutResult(length, null)
    }
  }

  override def getBytes(blockId: String): Option[ByteBuffer] = {
    val file = getFile(blockId)
    val bytes = getFileBytes(file)
    Some(bytes)
  }

  override def getValues(blockId: String): Option[Iterator[Any]] = {
    getBytes(blockId).map(bytes => blockManager.dataDeserialize(blockId, bytes))
  }

  /**
   * A version of getValues that allows a custom serializer. This is used as part of the
   * shuffle short-circuit code.
   */
  def getValues(blockId: String, serializer: Serializer): Option[Iterator[Any]] = {
    getBytes(blockId).map(bytes => blockManager.dataDeserialize(blockId, bytes, serializer))
  }

  override def remove(blockId: String): Boolean = {
    val file = getFile(blockId)
    if (file.exists()) {
      file.delete()
    } else {
      false
    }
  }

  override def contains(blockId: String): Boolean = {
    getFile(blockId).exists()
  }

  private def createFile(blockId: String, allowAppendExisting: Boolean = false): File = {
    var file = getFile(blockId)
    if (!allowAppendExisting && file.exists()) {
      // NOTE(shivaram): Delete the file if it exists. This might happen if a ShuffleMap task
      // was rescheduled on the same machine as the old task ?
      logWarning("File for block " + blockId + " already exists on disk: " + file + ". Deleting")
      file.delete()
      // throw new Exception("File for block " + blockId + " already exists on disk: " + file)<commit_msg>var to val<commit_after>package spark.storage

import java.io.{File, FileOutputStream, OutputStream, RandomAccessFile}
import java.nio.ByteBuffer
import java.nio.channels.FileChannel
import java.nio.channels.FileChannel.MapMode
import java.util.{Random, Date}
import java.text.SimpleDateFormat

import scala.collection.mutable.ArrayBuffer

import it.unimi.dsi.fastutil.io.FastBufferedOutputStream

import spark.Utils
import spark.executor.ExecutorExitCode
import spark.serializer.{Serializer, SerializationStream}
import spark.Logging
import spark.network.netty.ShuffleSender
import spark.network.netty.PathResolver


/**
 * Stores BlockManager blocks on disk.
 */
private class DiskStore(blockManager: BlockManager, rootDirs: String)
  extends BlockStore(blockManager) with Logging {

  class DiskBlockObjectWriter(blockId: String, serializer: Serializer, bufferSize: Int)
    extends BlockObjectWriter(blockId) {

    private val f: File = createFile(blockId /*, allowAppendExisting */)

    // The file channel, used for repositioning / truncating the file.
    private var channel: FileChannel = null
    private var bs: OutputStream = null
    private var objOut: SerializationStream = null
    private var lastValidPosition = 0L

    override def open(): DiskBlockObjectWriter = {
      val fos = new FileOutputStream(f, true)
      channel = fos.getChannel()
      bs = blockManager.wrapForCompression(blockId, new FastBufferedOutputStream(fos))
      objOut = serializer.newInstance().serializeStream(bs)
      this
    }

    override def close() {
      objOut.close()
      bs.close()
      channel = null
      bs = null
      objOut = null
      // Invoke the close callback handler.
      super.close()
    }

    override def isOpen: Boolean = objOut != null

    // Flush the partial writes, and set valid length to be the length of the entire file.
    // Return the number of bytes written for this commit.
    override def commit(): Long = {
      // NOTE: Flush the serializer first and then the compressed/buffered output stream
      objOut.flush()
      bs.flush()
      val prevPos = lastValidPosition
      lastValidPosition = channel.position()
      lastValidPosition - prevPos
    }

    override def revertPartialWrites() {
      // Discard current writes. We do this by flushing the outstanding writes and
      // truncate the file to the last valid position.
      objOut.flush()
      bs.flush()
      channel.truncate(lastValidPosition)
    }

    override def write(value: Any) {
      objOut.writeObject(value)
    }

    override def size(): Long = lastValidPosition
  }

  val MAX_DIR_CREATION_ATTEMPTS: Int = 10
  val subDirsPerLocalDir = System.getProperty("spark.diskStore.subDirectories", "64").toInt

  var shuffleSender : ShuffleSender = null
  // Create one local directory for each path mentioned in spark.local.dir; then, inside this
  // directory, create multiple subdirectories that we will hash files into, in order to avoid
  // having really large inodes at the top level.
  val localDirs = createLocalDirs()
  val subDirs = Array.fill(localDirs.length)(new Array[File](subDirsPerLocalDir))

  addShutdownHook()

  def getBlockWriter(blockId: String, serializer: Serializer, bufferSize: Int)
    : BlockObjectWriter = {
    new DiskBlockObjectWriter(blockId, serializer, bufferSize)
  }


  override def getSize(blockId: String): Long = {
    getFile(blockId).length()
  }

  override def putBytes(blockId: String, _bytes: ByteBuffer, level: StorageLevel) {
    // So that we do not modify the input offsets !
    // duplicate does not copy buffer, so inexpensive
    val bytes = _bytes.duplicate()
    logDebug("Attempting to put block " + blockId)
    val startTime = System.currentTimeMillis
    val file = createFile(blockId)
    val channel = new RandomAccessFile(file, "rw").getChannel()
    while (bytes.remaining > 0) {
      channel.write(bytes)
    }
    channel.close()
    val finishTime = System.currentTimeMillis
    logDebug("Block %s stored as %s file on disk in %d ms".format(
      blockId, Utils.memoryBytesToString(bytes.limit), (finishTime - startTime)))
  }

  private def getFileBytes(file: File): ByteBuffer = {
    val length = file.length()
    val channel = new RandomAccessFile(file, "r").getChannel()
    val buffer = try {
      channel.map(MapMode.READ_ONLY, 0, length)
    } finally {
      channel.close()
    }

    buffer
  }

  override def putValues(
      blockId: String,
      values: ArrayBuffer[Any],
      level: StorageLevel,
      returnValues: Boolean)
    : PutResult = {

    logDebug("Attempting to write values for block " + blockId)
    val startTime = System.currentTimeMillis
    val file = createFile(blockId)
    val fileOut = blockManager.wrapForCompression(blockId,
      new FastBufferedOutputStream(new FileOutputStream(file)))
    val objOut = blockManager.defaultSerializer.newInstance().serializeStream(fileOut)
    objOut.writeAll(values.iterator)
    objOut.close()
    val length = file.length()

    val timeTaken = System.currentTimeMillis - startTime
    logDebug("Block %s stored as %s file on disk in %d ms".format(
      blockId, Utils.memoryBytesToString(length), timeTaken))

    if (returnValues) {
      // Return a byte buffer for the contents of the file
      val buffer = getFileBytes(file)
      PutResult(length, Right(buffer))
    } else {
      PutResult(length, null)
    }
  }

  override def getBytes(blockId: String): Option[ByteBuffer] = {
    val file = getFile(blockId)
    val bytes = getFileBytes(file)
    Some(bytes)
  }

  override def getValues(blockId: String): Option[Iterator[Any]] = {
    getBytes(blockId).map(bytes => blockManager.dataDeserialize(blockId, bytes))
  }

  /**
   * A version of getValues that allows a custom serializer. This is used as part of the
   * shuffle short-circuit code.
   */
  def getValues(blockId: String, serializer: Serializer): Option[Iterator[Any]] = {
    getBytes(blockId).map(bytes => blockManager.dataDeserialize(blockId, bytes, serializer))
  }

  override def remove(blockId: String): Boolean = {
    val file = getFile(blockId)
    if (file.exists()) {
      file.delete()
    } else {
      false
    }
  }

  override def contains(blockId: String): Boolean = {
    getFile(blockId).exists()
  }

  private def createFile(blockId: String, allowAppendExisting: Boolean = false): File = {
    val file = getFile(blockId)
    if (!allowAppendExisting && file.exists()) {
      // NOTE(shivaram): Delete the file if it exists. This might happen if a ShuffleMap task
      // was rescheduled on the same machine as the old task ?
      logWarning("File for block " + blockId + " already exists on disk: " + file + ". Deleting")
      file.delete()
      // throw new Exception("File for block " + blockId + " already exists on disk: " + file)
{'fraction_non_alphanumeric': 0.09392143954355936, 'fraction_numerical': 0.0010240655401945725, 'mean_word_length': 3.406058652916532, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, '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': '7582880', 'n_tokens_mistral': 3965, 'n_tokens_neox': 3867, 'n_words': 1171}