<commit_before>// Copyright 2016 Foursquare Labs Inc. All Rights Reserved.

package io.fsq.hfile.reader.concrete.mmap

import com.twitter.ostrich.stats.Stats
import io.fsq.common.logging.Logger
import io.fsq.common.scala.Identity._
import io.fsq.hfile.reader.service.{HFileReader, HFileScanner, MsgAndSize}
import java.io.{File, FileInputStream, IOException}
import java.nio.ByteBuffer
import java.nio.channels.FileChannel.MapMode.READ_ONLY
import org.apache.hadoop.conf.Configuration
import org.apache.hadoop.fs.Path
import scala.annotation.tailrec
import scala.collection.JavaConverters._

trait BlockReader extends Function2[Long, Long, ByteBuffer] {
  override def apply(start: Long, length: Long): ByteBuffer
  def close(): Unit
}

case class FileChunk(startOffset: Long, buffer: ByteBuffer) {
  override def toString: String = startOffset + "," + buffer
}

object ConcreteHFileReader {
  @volatile var mmapedBytes: Long = 0
  val MB = 1024 * 1024

  def fromFile(file: File, hfilename: String, maximumMMappedBytes: Option[Long], trackCacheStats: Boolean): ConcreteHFileReader = {

    maximumMMappedBytes.map(m => Stats.addGauge("HFileService.TotalCollectionSizeLimit") { m.toDouble })

    val filePath = file.getPath
    val stream = new FileInputStream(file)
    val channel = stream.getChannel

    val lastOffset = file.length // ???

    val DefaultBlockSize = 1024 * MB // 1GB blocks
    val DefaultBlockOverlap = 8 * MB

    val chunks: List[FileChunk] = {
      val chunksBuilder = List.newBuilder[FileChunk]
      var offset = 0L
      while (offset < lastOffset) {
        val blockSize = Math.min(DefaultBlockSize, lastOffset - offset)
        val startOffset = Math.max(0, offset - DefaultBlockOverlap)
        val endOffset = offset - startOffset + blockSize
        chunksBuilder += FileChunk(startOffset, channel.map(READ_ONLY, startOffset, endOffset))
        offset += blockSize
      }
      chunksBuilder.result()
    }

    val blockReader = new BlockReader with Logger {
      override def apply(startOffset: Long, size: Long): ByteBuffer = {
        val buffer = chunks.find(c => c.startOffset <= startOffset && c.startOffset + c.buffer.limit  >= startOffset + size)
          .map(c => {
            val buffer = c.buffer.slice
            buffer.position((startOffset - c.startOffset).toInt)
            buffer.limit((startOffset - c.startOffset + size).toInt)
            buffer.slice
          })
          .getOrElse ({
            logger.info(s"offset=$startOffset, size=$size not possible for chunks: ${chunks.mkString}")
            channel.map(READ_ONLY, startOffset, size)
          })
        Bytes.copyDirectBufferToHeap(buffer)
      }

      override def close(): Unit = {
        for (chunk <- chunks) {
          Bytes.freeDirectByteBuffer(chunk.buffer)
        }
        channel.close()
      }
    }

    new ConcreteHFileReader(filePath, file.length, hfilename, blockReader, trackCacheStats, true, maximumMMappedBytes)
  }

  def fromFileChannel(file: File, hfilename: String, maximumMMappedBytes: Option[Long], trackCacheStats: Boolean): ConcreteHFileReader = {
    maximumMMappedBytes.map(m => Stats.addGauge("HFileService.TotalCollectionSizeLimit") { m.toDouble })

    val filePath = file.getPath
    val stream = new FileInputStream(file)
    val channel = stream.getChannel

    val lastOffset = file.length // ???

    val blockReader = new BlockReader with Logger {
      override def apply(startOffset: Long, size: Long): ByteBuffer = synchronized {
        channel.position(startOffset)
        val buf = ByteBuffer.allocate(size.toInt)
        var ret = 0
        while (channel.read(buf) != -1 && buf.hasRemaining) {}
        buf.flip
        buf
      }

      override def close(): Unit = {
        channel.close()
      }
    }

    new ConcreteHFileReader(filePath, file.length, hfilename, blockReader, trackCacheStats, true, maximumMMappedBytes)
  }

  def fromHadoopPath(path: Path, conf: Configuration): ConcreteHFileReader = {
    val remoteFS = path.getFileSystem(conf)
    val len = remoteFS.getFileStatus(path).getLen
    val stream = remoteFS.open(path)

    val blockReader = new BlockReader {
      override def apply(start: Long, length: Long): ByteBuffer = {
        if (length > Int.MaxValue) {
          throw new IllegalArgumentException("Can't map a ByteBuffer that big: %d, %d".format(start, length))
        }
        val arr = new Array[Byte](length.toInt)
        val bytesRead = stream.read(start, arr, 0, length.toInt)
        ByteBuffer.wrap(arr, 0, bytesRead)
      }
      override def close(): Unit = {
        stream.close()
      }
    }
    new ConcreteHFileReader(path.toString, len, "hadoop", blockReader, false, false, None)
  }

  def fromBytes(name: String, input: Array[Byte]): ConcreteHFileReader = {
    val blockReader = new BlockReader {
      override def apply(start: Long, length: Long): ByteBuffer = {
        val arr = input.slice(start.toInt, (start + length).toInt)
        ByteBuffer.wrap(arr, 0, arr.length)
      }
      override def close(): Unit = { }
    }
    new ConcreteHFileReader(name, input.length, "bytes", blockReader, false, false, None)
  }

}

class ConcreteHFileReader(
    val filePath: String,
    val fileLength: Long,
    val hfilename: String,
    blockReader: BlockReader,
    val trackCacheStats: Boolean,<commit_msg>S3 read fully (fixed) (#2126)<commit_after>// Copyright 2016 Foursquare Labs Inc. All Rights Reserved.

package io.fsq.hfile.reader.concrete.mmap

import com.twitter.ostrich.stats.Stats
import io.fsq.common.logging.Logger
import io.fsq.common.scala.Identity._
import io.fsq.hfile.reader.service.{HFileReader, HFileScanner, MsgAndSize}
import java.io.{EOFException, File, FileInputStream, IOException}
import java.nio.ByteBuffer
import java.nio.channels.FileChannel.MapMode.READ_ONLY
import org.apache.hadoop.conf.Configuration
import org.apache.hadoop.fs.{FSDataInputStream, Path}
import scala.annotation.tailrec
import scala.collection.JavaConverters._

trait BlockReader extends Function2[Long, Long, ByteBuffer] {
  override def apply(start: Long, length: Long): ByteBuffer
  def close(): Unit
}

case class FileChunk(startOffset: Long, buffer: ByteBuffer) {
  override def toString: String = startOffset + "," + buffer
}

object ConcreteHFileReader {
  @volatile var mmapedBytes: Long = 0
  val MB = 1024 * 1024

  def fromFile(file: File, hfilename: String, maximumMMappedBytes: Option[Long], trackCacheStats: Boolean): ConcreteHFileReader = {

    maximumMMappedBytes.map(m => Stats.addGauge("HFileService.TotalCollectionSizeLimit") { m.toDouble })

    val filePath = file.getPath
    val stream = new FileInputStream(file)
    val channel = stream.getChannel

    val lastOffset = file.length // ???

    val DefaultBlockSize = 1024 * MB // 1GB blocks
    val DefaultBlockOverlap = 8 * MB

    val chunks: List[FileChunk] = {
      val chunksBuilder = List.newBuilder[FileChunk]
      var offset = 0L
      while (offset < lastOffset) {
        val blockSize = Math.min(DefaultBlockSize, lastOffset - offset)
        val startOffset = Math.max(0, offset - DefaultBlockOverlap)
        val endOffset = offset - startOffset + blockSize
        chunksBuilder += FileChunk(startOffset, channel.map(READ_ONLY, startOffset, endOffset))
        offset += blockSize
      }
      chunksBuilder.result()
    }

    val blockReader = new BlockReader with Logger {
      override def apply(startOffset: Long, size: Long): ByteBuffer = {
        val buffer = chunks.find(c => c.startOffset <= startOffset && c.startOffset + c.buffer.limit  >= startOffset + size)
          .map(c => {
            val buffer = c.buffer.slice
            buffer.position((startOffset - c.startOffset).toInt)
            buffer.limit((startOffset - c.startOffset + size).toInt)
            buffer.slice
          })
          .getOrElse ({
            logger.info(s"offset=$startOffset, size=$size not possible for chunks: ${chunks.mkString}")
            channel.map(READ_ONLY, startOffset, size)
          })
        Bytes.copyDirectBufferToHeap(buffer)
      }

      override def close(): Unit = {
        for (chunk <- chunks) {
          Bytes.freeDirectByteBuffer(chunk.buffer)
        }
        channel.close()
      }
    }

    new ConcreteHFileReader(filePath, file.length, hfilename, blockReader, trackCacheStats, true, maximumMMappedBytes)
  }

  def fromFileChannel(file: File, hfilename: String, maximumMMappedBytes: Option[Long], trackCacheStats: Boolean): ConcreteHFileReader = {
    maximumMMappedBytes.map(m => Stats.addGauge("HFileService.TotalCollectionSizeLimit") { m.toDouble })

    val filePath = file.getPath
    val stream = new FileInputStream(file)
    val channel = stream.getChannel

    val lastOffset = file.length // ???

    val blockReader = new BlockReader with Logger {
      override def apply(startOffset: Long, size: Long): ByteBuffer = synchronized {
        channel.position(startOffset)
        val buf = ByteBuffer.allocate(size.toInt)
        var ret = 0
        while (channel.read(buf) != -1 && buf.hasRemaining) {}
        buf.flip
        buf
      }

      override def close(): Unit = {
        channel.close()
      }
    }

    new ConcreteHFileReader(filePath, file.length, hfilename, blockReader, trackCacheStats, true, maximumMMappedBytes)
  }

  def fromHadoopPath(path: Path, conf: Configuration): ConcreteHFileReader = {
    val remoteFS = path.getFileSystem(conf)
    val len = remoteFS.getFileStatus(path).getLen
    val stream = remoteFS.open(path)
    val fsName = remoteFS.getClass.getName
    //com.amazon.ws.emr.hadoop.fs.EmrFileSystem for s3 & s3n URI scheme on EMR
    //org.apache.hadoop.fs.s3a.S3AFileSystem for s3a URI scheme
    val isS3 = fsName.contains("EmrFileSystem") || fsName.contains("S3")
    val blockReader = new BlockReader {
      override def apply(start: Long, length: Long): ByteBuffer = {
        if (length > Int.MaxValue) {
          throw new IllegalArgumentException("Can't map a ByteBuffer that big: %d, %d".format(start, length))
        }
        val arr = new Array[Byte](length.toInt)
        if (isS3) {
          readFullyWorkAround(stream, start, arr, 0, length.toInt)
        } else {
          stream.readFully(start, arr, 0, length.toInt)
        }
        ByteBuffer.wrap(arr, 0, length.toInt)
      }
      override def close(): Unit = {
        stream.close()
      }
    }
    new ConcreteHFileReader(path.toString, len, "hadoop", blockReader, false, false, None)
  }

  def fromBytes(name: String, input: Array[Byte]): ConcreteHFileReader = {
    val blockReader = new BlockReader {
      override def apply(start: Long, length: Long): ByteBuffer = {
        val arr = input.slice(start.toInt, (start + length).toInt)
        ByteBuffer.wrap(arr, 0, arr.length)
      }
      override def close(): Unit = { }
    }
    new ConcreteHFileReader(name, input.length, "bytes", blockReader, false, false, None)
  }

  //Work around for S3 before 2.8.1 see https://issues.apache.org/jira/browse/HADOOP-11694
  //code inspired from https://issues.apache.org/jira/browse/HADOOP-11694
  //TODO: remove after Amazon upgrades Hadoop
  def readFullyWorkAround (stream: FSDataInputStream, position: Long, buffer: Array[Byte], offset: Int, length: Int): Unit =
    stream.synchronized {
      var nread = 0
      stream.seek(position)
      while (nread < length) {
        val nbytes = stream.read(buffer, offset + nread, length - nread)
        if (nbytes < 0) throw new EOFException("End of file reached before reading fully.")
        nread += nbytes
      }
    }
}

class ConcreteHFileReader(
    val filePath: String,
    val fileLength: Long,
    val hfilename: String,
    blockReader: BlockReader,
    val trackCacheStats: Boolean,
{'fraction_non_alphanumeric': 0.1059073523161921, 'fraction_numerical': 0.007139821504462389, 'mean_word_length': 3.0953706926557607, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 15, 'https://': 2, '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': '1918820', 'n_tokens_mistral': 3558, 'n_tokens_neox': 3437, 'n_words': 877}