<commit_before>package xerial.silk.cluster.framework

import xerial.silk.framework._
import xerial.core.log.Logger
import xerial.silk.core.SilkSerializer
import xerial.silk.SilkException
import xerial.core.io.IOUtil
import java.net.URL
import xerial.larray.{LArray, MMapMode}
import xerial.silk.cluster.DataServer.{RawData, ByteData, MmapData}
import java.util.UUID
import xerial.silk.cluster.{DataServerComponent, SilkClient}
import xerial.core.util.Timer

/**
 * @author <NAME>
 */
trait DistributedSliceStorage extends SliceStorageComponent with IDUtil {
  self: SilkFramework
    with DistributedCache
    with DataServerComponent
    with NodeManagerComponent
    with LocalClientAPI =>
    //with LocalClientComponent =>

  type LocalClient = SilkClient
  val sliceStorage = new SliceStorage

  class SliceStorage extends SliceStorageAPI with Timer with Logger {

    private def slicePath(op:Silk[_], index:Int) = {
      // TODO append session path: s"${session.sessionIDPrefix}/slice/${op.idPrefix}/${index}"
      s"slice/${op.idPrefix}/${index}"
    }
    private def slicePath(opid:UUID, index:Int) = {
      // TODO append session path: s"${session.sessionIDPrefix}/slice/${op.idPrefix}/${index}"
      s"slice/${opid.prefix}/${index}"
    }

    private def partitionSlicePath(opid:UUID, partition:Int, index:Int) = {
      s"slice/${opid.prefix}/${partition}:${index}"
    }

    private def stageInfoPath(op:Silk[_]) = {
      // TODO append session path: s"${session.sessionIDPrefix}/slice/${op.idPrefix}/${index}"
      s"slice/${op.idPrefix}/info"
    }

    def getStageInfo(op:Silk[_]) : Option[StageInfo] = {
      val p = stageInfoPath(op)
      cache.get(p).map(b => SilkSerializer.deserializeObj[StageInfo](b))
    }

    def setStageInfo(op:Silk[_], stageInfo:StageInfo) {
      val p = stageInfoPath(op)
      info(s"Update: $p $stageInfo")
      cache.update(p, SilkSerializer.serializeObj(stageInfo))
    }

    def get(opid: UUID, index: Int) : Future[Slice] = {
      val p = slicePath(opid, index)

      cache.getOrAwait(p).map{b =>
        if(b == null) {
          // when Slice is not available (reported by poke)
          SilkException.error(s"Failed to retrieve slice ${opid.prefix}/$index")
        }
        else
          SilkSerializer.deserializeObj[Slice](b)
      }
    }

    def poke(opid:UUID, index: Int) {
      cache.update(slicePath(opid, index), null)
    }

    def poke(opid: UUID, partition: Int, index: Int) {
      cache.update(partitionSlicePath(opid, partition, index), null)
    }

    def put(opid: UUID, index: Int, slice: Slice, data:Seq[_]) {
      putRaw(opid, index, slice, SilkSerializer.serializeObj(data))
    }

    def putRaw(opid: UUID, index: Int, slice: Slice, data:Array[Byte]) {
      val path = s"${opid.prefix}/${index}"
      debug(s"put slice $path")
      dataServer.registerByteData(path, data)
      cache.update(slicePath(opid, index), SilkSerializer.serializeObj(slice))
    }



    def contains(op: Silk[_], index: Int) : Boolean = {
      cache.contains(slicePath(op, index))
    }

    def retrieve(opid:UUID, slice: Slice) = {
      val dataID = s"${opid.prefix}/${slice.path}"
      if(slice.nodeName == currentNodeName) {
        SilkClient.client.flatMap { c =>
          debug(s"retrieve $dataID from local DataServer: http://localhost:${c.dataServer.port}/data/${dataID}")
          val result : Option[Seq[_]] = c.dataServer.getData(dataID) map {
            case RawData(s, _) => s.asInstanceOf[Seq[_]]
            case ByteData(b, _) => SilkSerializer.deserializeObj[Seq[_]](b)
            case MmapData(file, _) => {
              val mmapped = LArray.mmap(file, 0, file.length, MMapMode.READ_ONLY)
              SilkSerializer.deserializeObj[Seq[_]](mmapped.toInputStream)
            }
          }
          trace(f"Retrieved ${result.get.size}%,d entries")
          result
        } getOrElse { SilkException.error(s"no slice data is found: [${opid.prefix}] ${slice}") }
      }
      else {
        nodeManager.getNode(slice.nodeName).map { n =>
          val url = new URL(s"http://${n.address}:${n.dataServerPort}/data/${dataID}")
          debug(s"retrieve $dataID from $url (${slice.nodeName})")
          val result = IOUtil.readFully(url.openStream) { data =>
            debug(f"Downloaded ${data.length}%,d bytes")
            val desr = SilkSerializer.deserializeObj[Seq[_]](data)
            debug(f"Deserialized ${desr.length}%,d entries")
            desr
          }
          result
        } getOrElse { SilkException.error(s"invalid node name: ${slice.nodeName}") }
      }
    }

    def putSlice(opid: UUID, partition: Int, index: Int, slice: Slice, data: Seq[_]) {
      val p = partitionSlicePath(opid, partition, index)<commit_msg>Overlap download / deserialization<commit_after>package xerial.silk.cluster.framework

import xerial.silk.framework._
import xerial.core.log.Logger
import xerial.silk.core.SilkSerializer
import xerial.silk.SilkException
import xerial.core.io.IOUtil
import java.net.URL
import xerial.larray.{LArray, MMapMode}
import xerial.silk.cluster.DataServer.{RawData, ByteData, MmapData}
import java.util.UUID
import xerial.silk.cluster.{DataServerComponent, SilkClient}
import xerial.core.util.Timer
import java.io.BufferedInputStream

/**
 * @author <NAME>
 */
trait DistributedSliceStorage extends SliceStorageComponent with IDUtil {
  self: SilkFramework
    with DistributedCache
    with DataServerComponent
    with NodeManagerComponent
    with LocalClientAPI =>
    //with LocalClientComponent =>

  type LocalClient = SilkClient
  val sliceStorage = new SliceStorage

  class SliceStorage extends SliceStorageAPI with Timer with Logger {

    private def slicePath(op:Silk[_], index:Int) = {
      // TODO append session path: s"${session.sessionIDPrefix}/slice/${op.idPrefix}/${index}"
      s"slice/${op.idPrefix}/${index}"
    }
    private def slicePath(opid:UUID, index:Int) = {
      // TODO append session path: s"${session.sessionIDPrefix}/slice/${op.idPrefix}/${index}"
      s"slice/${opid.prefix}/${index}"
    }

    private def partitionSlicePath(opid:UUID, partition:Int, index:Int) = {
      s"slice/${opid.prefix}/${partition}:${index}"
    }

    private def stageInfoPath(op:Silk[_]) = {
      // TODO append session path: s"${session.sessionIDPrefix}/slice/${op.idPrefix}/${index}"
      s"slice/${op.idPrefix}/info"
    }

    def getStageInfo(op:Silk[_]) : Option[StageInfo] = {
      val p = stageInfoPath(op)
      cache.get(p).map(b => SilkSerializer.deserializeObj[StageInfo](b))
    }

    def setStageInfo(op:Silk[_], stageInfo:StageInfo) {
      val p = stageInfoPath(op)
      info(s"Update: $p $stageInfo")
      cache.update(p, SilkSerializer.serializeObj(stageInfo))
    }

    def get(opid: UUID, index: Int) : Future[Slice] = {
      val p = slicePath(opid, index)

      cache.getOrAwait(p).map{b =>
        if(b == null) {
          // when Slice is not available (reported by poke)
          SilkException.error(s"Failed to retrieve slice ${opid.prefix}/$index")
        }
        else
          SilkSerializer.deserializeObj[Slice](b)
      }
    }

    def poke(opid:UUID, index: Int) {
      cache.update(slicePath(opid, index), null)
    }

    def poke(opid: UUID, partition: Int, index: Int) {
      cache.update(partitionSlicePath(opid, partition, index), null)
    }

    def put(opid: UUID, index: Int, slice: Slice, data:Seq[_]) {
      putRaw(opid, index, slice, SilkSerializer.serializeObj(data))
    }

    def putRaw(opid: UUID, index: Int, slice: Slice, data:Array[Byte]) {
      val path = s"${opid.prefix}/${index}"
      debug(s"put slice $path")
      dataServer.registerByteData(path, data)
      cache.update(slicePath(opid, index), SilkSerializer.serializeObj(slice))
    }



    def contains(op: Silk[_], index: Int) : Boolean = {
      cache.contains(slicePath(op, index))
    }

    def retrieve(opid:UUID, slice: Slice) = {
      val dataID = s"${opid.prefix}/${slice.path}"
      if(slice.nodeName == currentNodeName) {
        SilkClient.client.flatMap { c =>
          debug(s"retrieve $dataID from local DataServer: http://localhost:${c.dataServer.port}/data/${dataID}")
          val result : Option[Seq[_]] = c.dataServer.getData(dataID) map {
            case RawData(s, _) => s.asInstanceOf[Seq[_]]
            case ByteData(b, _) => SilkSerializer.deserializeObj[Seq[_]](b)
            case MmapData(file, _) => {
              val mmapped = LArray.mmap(file, 0, file.length, MMapMode.READ_ONLY)
              SilkSerializer.deserializeObj[Seq[_]](mmapped.toInputStream)
            }
          }
          trace(f"Retrieved ${result.get.size}%,d entries")
          result
        } getOrElse { SilkException.error(s"no slice data is found: [${opid.prefix}] ${slice}") }
      }
      else {
        nodeManager.getNode(slice.nodeName).map { n =>
          val url = new URL(s"http://${n.address}:${n.dataServerPort}/data/${dataID}")
          debug(s"retrieve $dataID from $url (${slice.nodeName})")
          val result = IOUtil.withResource(new BufferedInputStream(url.openStream)) { in =>
            val desr = SilkSerializer.deserializeObj[Seq[_]](in)
            debug(f"Deserialized ${desr.length}%,d entries")
            desr
          }
          result
        } getOrElse { SilkException.error(s"invalid node name: ${slice.nodeName}") }
      }
    }

    def putSlice(opid: UUID, partition: Int, index: Int, slice: Slice, data: Seq[_]) {
      val p = partitionSlicePath(opid, partition, index)
{'fraction_non_alphanumeric': 0.13995177691581928, 'fraction_numerical': 0.0002096655833944858, 'mean_word_length': 3.083904109589041, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2375106', 'n_tokens_mistral': 3075, 'n_tokens_neox': 3079, 'n_words': 599}