<commit_before>package com.wavesplatform.network

import com.wavesplatform.metrics.LatencyHistogram
import com.wavesplatform.state2.ByteStr
import io.netty.channel.{ChannelHandlerContext, ChannelInboundHandlerAdapter}
import io.netty.util.concurrent.ScheduledFuture
import kamon.Kamon
import kamon.metric.instrument
import scorex.block.Block
import scorex.utils.ScorexLogging

import scala.collection.mutable
import scala.concurrent.duration.FiniteDuration

class ExtensionBlocksLoader(
    blockSyncTimeout: FiniteDuration,
    peerDatabase: PeerDatabase) extends ChannelInboundHandlerAdapter with ScorexLogging {
  private var pendingSignatures = Map.empty[ByteStr, Int]
  private var targetExtensionIds = Option.empty[ExtensionIds]
  private val blockBuffer = mutable.TreeMap.empty[Int, Block]
  private var currentTimeout = Option.empty[ScheduledFuture[_]]
  private val extensionsFetchingTimeStats = new LatencyHistogram(Kamon.metrics.histogram("extensions-fetching-time", instrument.Time.Milliseconds))

  private def cancelTimeout(): Unit = {
    currentTimeout.foreach(_.cancel(false))
    currentTimeout = None
  }

  private def resetTimeout(ctx: ChannelHandlerContext, xid: ExtensionIds): Unit = {
    cancelTimeout()
    currentTimeout = Some(ctx.executor().schedule(blockSyncTimeout) {
      if (targetExtensionIds.contains(xid)) {
        peerDatabase.blacklistAndClose(ctx.channel(), "Timeout loading blocks")
      }
    })
  }

  override def channelInactive(ctx: ChannelHandlerContext) = {
    cancelTimeout()
    super.channelInactive(ctx)
  }

  override def channelRead(ctx: ChannelHandlerContext, msg: AnyRef) = msg match {
    case xid@ExtensionIds(_, newIds) if pendingSignatures.isEmpty =>
        if (newIds.nonEmpty) {
          targetExtensionIds = Some(xid)
          pendingSignatures = newIds.zipWithIndex.toMap
          resetTimeout(ctx, xid)
          extensionsFetchingTimeStats.start()
          newIds.foreach(s => ctx.write(GetBlock(s)))
          ctx.flush()
        } else {
          log.debug(s"${id(ctx)} No new blocks to load")
          ctx.fireChannelRead(ExtensionBlocks(Seq.empty))
        }
    case b: Block if pendingSignatures.contains(b.uniqueId) =>
      blockBuffer += pendingSignatures(b.uniqueId) -> b
      pendingSignatures -= b.uniqueId
      if (pendingSignatures.isEmpty) {
        cancelTimeout()
        extensionsFetchingTimeStats.record()
        log.trace(s"${id(ctx)} Loaded all blocks, doing a pre-check")

        val newBlocks = blockBuffer.values.toSeq<commit_msg>NODE-229: timeout per each block of extension<commit_after>package com.wavesplatform.network

import com.wavesplatform.metrics.LatencyHistogram
import com.wavesplatform.state2.ByteStr
import io.netty.channel.{ChannelHandlerContext, ChannelInboundHandlerAdapter}
import io.netty.util.concurrent.ScheduledFuture
import kamon.Kamon
import kamon.metric.instrument
import scorex.block.Block
import scorex.utils.ScorexLogging

import scala.collection.mutable
import scala.concurrent.duration.FiniteDuration

class ExtensionBlocksLoader(
    blockSyncTimeout: FiniteDuration,
    peerDatabase: PeerDatabase) extends ChannelInboundHandlerAdapter with ScorexLogging {
  private var pendingSignatures = Map.empty[ByteStr, Int]
  private var targetExtensionIds = Option.empty[ExtensionIds]
  private val blockBuffer = mutable.TreeMap.empty[Int, Block]
  private var currentTimeout = Option.empty[ScheduledFuture[_]]
  private val extensionsFetchingTimeStats = new LatencyHistogram(Kamon.metrics.histogram("extensions-fetching-time", instrument.Time.Milliseconds))

  private def cancelBlaclkistTask(): Unit = {
    currentTimeout.foreach(_.cancel(false))
    currentTimeout = None
  }

  private def rescheduleBlacklistOnTimeput(ctx: ChannelHandlerContext): Unit = {
    cancelBlaclkistTask()
    currentTimeout = Some(ctx.executor().schedule(blockSyncTimeout) {
        peerDatabase.blacklistAndClose(ctx.channel(), "Timeout loading blocks")
    })
  }

  override def channelInactive(ctx: ChannelHandlerContext) = {
    cancelBlaclkistTask()
    super.channelInactive(ctx)
  }

  override def channelRead(ctx: ChannelHandlerContext, msg: AnyRef) = msg match {
    case xid@ExtensionIds(_, newIds) if pendingSignatures.isEmpty =>
        if (newIds.nonEmpty) {
          targetExtensionIds = Some(xid)
          pendingSignatures = newIds.zipWithIndex.toMap
          rescheduleBlacklistOnTimeput(ctx)
          extensionsFetchingTimeStats.start()
          newIds.foreach(s => ctx.write(GetBlock(s)))
          ctx.flush()
        } else {
          log.debug(s"${id(ctx)} No new blocks to load")
          ctx.fireChannelRead(ExtensionBlocks(Seq.empty))
        }
    case b: Block if pendingSignatures.contains(b.uniqueId) =>
      blockBuffer += pendingSignatures(b.uniqueId) -> b
      pendingSignatures -= b.uniqueId
      rescheduleBlacklistOnTimeput(ctx)
      if (pendingSignatures.isEmpty) {
        cancelBlaclkistTask()
        extensionsFetchingTimeStats.record()
        log.trace(s"${id(ctx)} Loaded all blocks, doing a pre-check")

        val newBlocks = blockBuffer.values.toSeq
{'fraction_non_alphanumeric': 0.092151503319016, 'fraction_numerical': 0.0009761811792268645, 'mean_word_length': 4.169525731584258, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 11, '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': '6440805', 'n_tokens_mistral': 1443, 'n_tokens_neox': 1430, 'n_words': 287}