<commit_before>package com.twitter.finagle.stress

import scala.collection.mutable.{ArrayBuffer, HashMap}

import java.util.concurrent.atomic.AtomicInteger

import org.jboss.netty.handler.codec.http._

import com.twitter.ostrich.stats.{Stats => OstrichStats}
import com.twitter.ostrich.stats.{StatsCollection, StatsProvider}
import com.twitter.util.{Duration, CountDownLatch, Return, Throw, Time}
import com.twitter.conversions.time._

import com.twitter.finagle.builder.ClientBuilder
import com.twitter.finagle.http.Http
import com.twitter.finagle.Service
import com.twitter.finagle.stats.{NullStatsReceiver, OstrichStatsReceiver}
import com.twitter.finagle.util.Timer
import com.twitter.finagle.util.Conversions._

object LoadBalancerTest {
  def main(args: Array[String]) {
    // Make the type enforced by the *codec*

    runSuite(
      ClientBuilder()
        .requestTimeout(100.milliseconds)
        .retries(10)
    )

    // TODO: proper resource releasing, etc.
  }

  def runSuite(clientBuilder: ClientBuilder[_, _, _, _, _]) {
    println("testing " + clientBuilder)
    println("\n== baseline (warmup) ==\n")
    new LoadBalancerTest(clientBuilder)({ case _ => }).run()

    println("\n== baseline ==\n")
    new LoadBalancerTest(clientBuilder)({ case _ => }).run()

    println("\n== 1 server goes offline ==\n")
    new LoadBalancerTest(clientBuilder)({
      case (10000, servers) =>
        servers(1).stop()
    }).run()

    println("\n== 1 application becomes nonresponsive ==\n")
    new LoadBalancerTest(clientBuilder)({
      case (10000, servers) =>
        servers(1).becomeApplicationNonresponsive()
    }).run()

    println("\n== 1 connection becomes nonresponsive ==\n")
    new LoadBalancerTest(clientBuilder)({
      case (10000, servers) =>
        servers(1).becomeConnectionNonresponsive()
    }).run()

    println("\n== 1 server has a protocol error ==\n")
    new LoadBalancerTest(clientBuilder)({
      case (10000, servers) =>
        servers(1).becomeBelligerent()
    }).run()
  }
}

class LoadBalancerTest(
  clientBuilder: ClientBuilder[_, _, _, _, _],
  serverLatency: Duration = 0.seconds,
  numRequests: Int = 100000,
  concurrency: Int = 20)(behavior: PartialFunction[(Int, Seq[EmbeddedServer]), Unit])
{
  private[this] val requestNumber = new AtomicInteger(0)
  private[this] val requestCount  = new AtomicInteger(numRequests)
  private[this] val latch         = new CountDownLatch(concurrency)
  private[this] val stats         = new StatsCollection
  private[this] val gaugeValues   = new ArrayBuffer[(Int, Map[String, Float])]

  private[this] def dispatch(
      client: Service[HttpRequest, HttpResponse],
      servers: Seq[EmbeddedServer],
      f: PartialFunction[(Int, Seq[EmbeddedServer]), Unit]) {
    val num = requestNumber.incrementAndGet()
    if (f.isDefinedAt((num, servers)))
      f((num, servers))

    val beginTime = Time.now

    client(new DefaultHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.GET, "/")) respond { result =>
      result match {
        case Return(_) =>
          val duration = beginTime.untilNow
          stats.addMetric("request_msec", duration.inMilliseconds.toInt)
          stats.incr("success")
        case Throw(exc) =>
          stats.incr("fail")
          stats.incr("fail_%s".format(exc.getClass.getName.split('.').last))
      }

      if (requestCount.decrementAndGet() > 0)
        dispatch(client, servers, f)
      else
        latch.countDown()
    }
  }

  def run() {<commit_msg>[split] LoadBalancerTest: add total count<commit_after>package com.twitter.finagle.stress

import scala.collection.mutable.{ArrayBuffer, HashMap}

import java.util.concurrent.atomic.AtomicInteger

import org.jboss.netty.handler.codec.http._

import com.twitter.ostrich.stats.{Stats => OstrichStats}
import com.twitter.ostrich.stats.{StatsCollection, StatsProvider}
import com.twitter.util.{Duration, CountDownLatch, Return, Throw, Time}
import com.twitter.conversions.time._

import com.twitter.finagle.builder.ClientBuilder
import com.twitter.finagle.http.Http
import com.twitter.finagle.Service
import com.twitter.finagle.stats.{NullStatsReceiver, OstrichStatsReceiver}
import com.twitter.finagle.util.Timer
import com.twitter.finagle.util.Conversions._

object LoadBalancerTest {
  val totalRequests = new AtomicInteger(0)

  def main(args: Array[String]) {
    // Make the type enforced by the *codec*

    runSuite(
      ClientBuilder()
        .requestTimeout(100.milliseconds)
        .retries(10)
    )

    // TODO: proper resource releasing, etc.
  }

  def runSuite(clientBuilder: ClientBuilder[_, _, _, _, _]) {
    println("testing " + clientBuilder)
    println("\n== baseline (warmup) ==\n")
    new LoadBalancerTest(clientBuilder)({ case _ => }).run()

    println("\n== baseline ==\n")
    new LoadBalancerTest(clientBuilder)({ case _ => }).run()

    println("\n== 1 server goes offline ==\n")
    new LoadBalancerTest(clientBuilder)({
      case (10000, servers) =>
        servers(1).stop()
    }).run()

    println("\n== 1 application becomes nonresponsive ==\n")
    new LoadBalancerTest(clientBuilder)({
      case (10000, servers) =>
        servers(1).becomeApplicationNonresponsive()
    }).run()

    println("\n== 1 connection becomes nonresponsive ==\n")
    new LoadBalancerTest(clientBuilder)({
      case (10000, servers) =>
        servers(1).becomeConnectionNonresponsive()
    }).run()

    println("\n== 1 server has a protocol error ==\n")
    new LoadBalancerTest(clientBuilder)({
      case (10000, servers) =>
        servers(1).becomeBelligerent()
    }).run()
  }
}

class LoadBalancerTest(
  clientBuilder: ClientBuilder[_, _, _, _, _],
  serverLatency: Duration = 0.seconds,
  numRequests: Int = 100000,
  concurrency: Int = 20)(behavior: PartialFunction[(Int, Seq[EmbeddedServer]), Unit])
{
  private[this] val requestNumber = new AtomicInteger(0)
  private[this] val requestCount  = new AtomicInteger(numRequests)
  private[this] val latch         = new CountDownLatch(concurrency)
  private[this] val stats         = new StatsCollection
  private[this] val gaugeValues   = new ArrayBuffer[(Int, Map[String, Float])]

  private[this] def dispatch(
      client: Service[HttpRequest, HttpResponse],
      servers: Seq[EmbeddedServer],
      f: PartialFunction[(Int, Seq[EmbeddedServer]), Unit]) {
    val num = requestNumber.incrementAndGet()
    LoadBalancerTest.totalRequests.incrementAndGet()
    if (f.isDefinedAt((num, servers)))
      f((num, servers))

    val beginTime = Time.now

    client(new DefaultHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.GET, "/")) respond { result =>
      result match {
        case Return(_) =>
          val duration = beginTime.untilNow
          stats.addMetric("request_msec", duration.inMilliseconds.toInt)
          stats.incr("success")
        case Throw(exc) =>
          stats.incr("fail")
          stats.incr("fail_%s".format(exc.getClass.getName.split('.').last))
      }

      if (requestCount.decrementAndGet() > 0)
        dispatch(client, servers, f)
      else
        latch.countDown()
    }
  }

  def run() {
{'fraction_non_alphanumeric': 0.13194542129694753, 'fraction_numerical': 0.013082008721339147, 'mean_word_length': 3.5752895752895753, 'pattern_counts': {'":': 0, '<': 3, '<?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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6544172', 'n_tokens_mistral': 2250, 'n_tokens_neox': 2120, 'n_words': 461}