<commit_before>/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package org.apache.samza.test.integration

import org.apache.samza.task.StreamTask
import org.apache.samza.task.TaskContext
import org.apache.samza.task.InitableTask
import org.apache.samza.config.Config
import scala.collection.JavaConversions._
import org.apache.samza.task.TaskCoordinator
import org.apache.samza.task.MessageCollector
import org.apache.samza.system.IncomingMessageEnvelope
import org.apache.samza.checkpoint.Checkpoint
import org.junit.BeforeClass
import org.junit.AfterClass
import kafka.zk.EmbeddedZookeeper
import kafka.utils.TestUtils
import org.apache.samza.system.SystemStream
import kafka.utils.TestZKUtils
import kafka.server.KafkaConfig
import org.I0Itec.zkclient.ZkClient
import kafka.producer.ProducerConfig
import kafka.server.KafkaServer
import kafka.utils.Utils
import org.apache.samza.storage.kv.KeyValueStore
import org.apache.samza.util._
import org.junit.Test
import kafka.admin.AdminUtils
import kafka.common.ErrorMapping
import org.junit.Assert._
import kafka.utils.ZKStringSerializer
import scala.collection.mutable.ArrayBuffer
import org.apache.samza.job.local.LocalJobFactory
import org.apache.samza.job.ApplicationStatus
import java.util.concurrent.CountDownLatch
import org.apache.samza.job.local.ThreadJob
import org.apache.samza.util.TopicMetadataStore
import org.apache.samza.util.ClientUtilTopicMetadataStore
import org.apache.samza.config.MapConfig
import org.apache.samza.system.kafka.TopicMetadataCache
import org.apache.samza.container.SamzaContainer
import scala.collection.mutable.HashMap
import scala.collection.mutable.SynchronizedMap
import org.apache.samza.Partition
import java.util.concurrent.TimeUnit
import kafka.producer.Producer
import kafka.producer.KeyedMessage
import java.util.concurrent.Semaphore
import java.util.concurrent.CyclicBarrier
import kafka.consumer.Consumer
import kafka.consumer.ConsumerConfig
import java.util.Properties
import java.util.concurrent.Executors
import kafka.message.MessageAndOffset
import kafka.message.MessageAndMetadata
import org.apache.samza.job.StreamJob
import org.apache.samza.task.TaskCoordinator.RequestScope

object TestStatefulTask {
  val INPUT_TOPIC = "input"
  val STATE_TOPIC = "mystore"
  val TOTAL_PARTITIONS = 1
  val REPLICATION_FACTOR = 3

  val zkConnect: String = TestZKUtils.zookeeperConnect
  var zkClient: ZkClient = null
  val zkConnectionTimeout = 6000
  val zkSessionTimeout = 6000

  val brokerId1 = 0
  val brokerId2 = 1
  val brokerId3 = 2
  val ports = TestUtils.choosePorts(3)
  val (port1, port2, port3) = (ports(0), ports(1), ports(2))

  val props1 = TestUtils.createBrokerConfig(brokerId1, port1)
  val props2 = TestUtils.createBrokerConfig(brokerId2, port2)
  val props3 = TestUtils.createBrokerConfig(brokerId3, port3)

  val config = new java.util.Properties()
  val brokers = "localhost:%d,localhost:%d,localhost:%d" format (port1, port2, port3)
  config.put("metadata.broker.list", brokers)
  config.put("producer.type", "sync")
  config.put("request.required.acks", "-1")
  config.put("serializer.class", "kafka.serializer.StringEncoder");
  val producerConfig = new ProducerConfig(config)
  var producer: Producer[String, String] = null
  val cp1 = new Checkpoint(Map(new SystemStream("kafka", "topic") -> "123"))
  val cp2 = new Checkpoint(Map(new SystemStream("kafka", "topic") -> "12345"))
  var zookeeper: EmbeddedZookeeper = null
  var server1: KafkaServer = null
  var server2: KafkaServer = null
  var server3: KafkaServer = null
  var metadataStore: TopicMetadataStore = null

  @BeforeClass
  def beforeSetupServers {
    zookeeper = new EmbeddedZookeeper(zkConnect)
    server1 = TestUtils.createServer(new KafkaConfig(props1))
    server2 = TestUtils.createServer(new KafkaConfig(props2))
    server3 = TestUtils.createServer(new KafkaConfig(props3))
    zkClient = new ZkClient(zkConnect + "/", 6000, 6000, ZKStringSerializer)
    producer = new Producer(producerConfig)
    metadataStore = new ClientUtilTopicMetadataStore(brokers, "some-job-name")

    createTopics
    validateTopics
  }

  def createTopics {
    AdminUtils.createTopic(
      zkClient,
      INPUT_TOPIC,
      TOTAL_PARTITIONS,
      REPLICATION_FACTOR)

    AdminUtils.createTopic(
      zkClient,
      STATE_TOPIC,
      TOTAL_PARTITIONS,
      REPLICATION_FACTOR)
  }

  def validateTopics {
    val topics = Set(STATE_TOPIC, INPUT_TOPIC)
    var done = false
    var retries = 0

    while (!done && retries < 100) {
      try {
        val topicMetadataMap = TopicMetadataCache.getTopicMetadata(topics, "kafka", metadataStore.getTopicInfo)

        topics.foreach(topic => {
          val topicMetadata = topicMetadataMap(topic)
          val errorCode = topicMetadata.errorCode

          ErrorMapping.maybeThrowException(errorCode)
        })

        done = true
      } catch {
        case e: Exception =>
          System.err.println("Got exception while validating test topics. Waiting and retrying.", e)
          retries += 1
          Thread.sleep(500)
      }
    }

    if (retries >= 100) {
      fail("Unable to successfully create topics. Tried to validate %s times." format retries)
    }
  }

  @AfterClass
  def afterCleanLogDirs {
    server1.shutdown
    server1.awaitShutdown()
    server2.shutdown
    server2.awaitShutdown()
    server3.shutdown
    server3.awaitShutdown()
    Utils.rm(server1.config.logDirs)
    Utils.rm(server2.config.logDirs)
    Utils.rm(server3.config.logDirs)
    zkClient.close
    zookeeper.shutdown
  }
}

/**
 * Test that does the following:
 *
 * 1. Starts ZK, and 3 kafka brokers.
 * 2. Create two topics: input and mystore.
 * 3. Validate that the topics were created successfully and have leaders.
 * 4. Start a single partition of TestTask using LocalJobFactory/ThreadJob.
 * 5. Send four messages to input (1,2,3,2), which contain one dupe (2).
 * 6. Validate that all messages were received by TestTask.
 * 7. Validate that TestTask called store.put() for all four messages, and that the messages ended up in the mystore topic.
 * 8. Kill the job.
 * 9. Start the job again.
 * 10. Validate that the job restored all messages (1,2,3) to the store.
 * 11. Send three more messages to input (4,5,5), and validate that TestTask receives them.
 * 12. Kill the job again.
 */
class TestStatefulTask {
  import TestStatefulTask._

  val jobConfig = Map(
    "job.factory.class" -> "org.apache.samza.job.local.LocalJobFactory",
    "job.name" -> "hello-stateful-world",
    "task.class" -> "org.apache.samza.test.integration.TestTask",
    "task.inputs" -> "kafka.input",
    "serializers.registry.string.class" -> "org.apache.samza.serializers.StringSerdeFactory",
    "stores.mystore.factory" -> "org.apache.samza.storage.kv.KeyValueStorageEngineFactory",
    "stores.mystore.key.serde" -> "string",
    "stores.mystore.msg.serde" -> "string",
    "stores.mystore.changelog" -> "kafka.mystore",

    // Use smallest reset for input streams, so we can fix SAMZA-166.
    "systems.kafka.samza.factory" -> "org.apache.samza.system.kafka.KafkaSystemFactory",
    "systems.kafka.samza.offset.default" -> "oldest",
    "systems.kafka.samza.msg.serde" -> "string",
    "systems.kafka.consumer.zookeeper.connect" -> zkConnect,
    // Use largest offset for reset, so we can test SAMZA-142.
    "systems.kafka.consumer.auto.offset.reset" -> "largest",
    "systems.kafka.producer.metadata.broker.list" -> ("localhost:%s" format port1))

  @Test
  def testShouldStartAndRestore {
    // Have to do this in one test to guarantee ordering.
    testShouldStartTaskForFirstTime
    testShouldRestoreStore
  }

  def testShouldStartTaskForFirstTime {
    val (job, task) = startJob

    // Validate that restored is empty.
    assertEquals(0, task.initFinished.getCount)
    assertEquals(0, task.restored.size)
    assertEquals(0, task.received.size)<commit_msg>SAMZA-224. Fix race condition in TestStatefulTask<commit_after>/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package org.apache.samza.test.integration

import org.apache.samza.task.StreamTask
import org.apache.samza.task.TaskContext
import org.apache.samza.task.InitableTask
import org.apache.samza.config.Config
import scala.collection.JavaConversions._
import org.apache.samza.task.TaskCoordinator
import org.apache.samza.task.MessageCollector
import org.apache.samza.system.IncomingMessageEnvelope
import org.apache.samza.checkpoint.Checkpoint
import org.junit.BeforeClass
import org.junit.AfterClass
import kafka.zk.EmbeddedZookeeper
import kafka.utils.TestUtils
import org.apache.samza.system.SystemStream
import kafka.utils.TestZKUtils
import kafka.server.KafkaConfig
import org.I0Itec.zkclient.ZkClient
import kafka.producer.ProducerConfig
import kafka.server.KafkaServer
import kafka.utils.Utils
import org.apache.samza.storage.kv.KeyValueStore
import org.apache.samza.util._
import org.junit.Test
import kafka.admin.AdminUtils
import kafka.common.ErrorMapping
import org.junit.Assert._
import kafka.utils.ZKStringSerializer
import scala.collection.mutable.ArrayBuffer
import org.apache.samza.job.local.LocalJobFactory
import org.apache.samza.job.ApplicationStatus
import java.util.concurrent.CountDownLatch
import org.apache.samza.job.local.ThreadJob
import org.apache.samza.util.TopicMetadataStore
import org.apache.samza.util.ClientUtilTopicMetadataStore
import org.apache.samza.config.MapConfig
import org.apache.samza.system.kafka.TopicMetadataCache
import org.apache.samza.container.SamzaContainer
import scala.collection.mutable.HashMap
import scala.collection.mutable.SynchronizedMap
import org.apache.samza.Partition
import java.util.concurrent.TimeUnit
import kafka.producer.Producer
import kafka.producer.KeyedMessage
import java.util.concurrent.Semaphore
import java.util.concurrent.CyclicBarrier
import kafka.consumer.Consumer
import kafka.consumer.ConsumerConfig
import java.util.Properties
import java.util.concurrent.Executors
import kafka.message.MessageAndOffset
import kafka.message.MessageAndMetadata
import org.apache.samza.job.StreamJob
import org.apache.samza.task.TaskCoordinator.RequestScope

object TestStatefulTask {
  val INPUT_TOPIC = "input"
  val STATE_TOPIC = "mystore"
  val TOTAL_PARTITIONS = 1
  val REPLICATION_FACTOR = 3

  val zkConnect: String = TestZKUtils.zookeeperConnect
  var zkClient: ZkClient = null
  val zkConnectionTimeout = 6000
  val zkSessionTimeout = 6000

  val brokerId1 = 0
  val brokerId2 = 1
  val brokerId3 = 2
  val ports = TestUtils.choosePorts(3)
  val (port1, port2, port3) = (ports(0), ports(1), ports(2))

  val props1 = TestUtils.createBrokerConfig(brokerId1, port1)
  val props2 = TestUtils.createBrokerConfig(brokerId2, port2)
  val props3 = TestUtils.createBrokerConfig(brokerId3, port3)

  val config = new java.util.Properties()
  val brokers = "localhost:%d,localhost:%d,localhost:%d" format (port1, port2, port3)
  config.put("metadata.broker.list", brokers)
  config.put("producer.type", "sync")
  config.put("request.required.acks", "-1")
  config.put("serializer.class", "kafka.serializer.StringEncoder");
  val producerConfig = new ProducerConfig(config)
  var producer: Producer[String, String] = null
  val cp1 = new Checkpoint(Map(new SystemStream("kafka", "topic") -> "123"))
  val cp2 = new Checkpoint(Map(new SystemStream("kafka", "topic") -> "12345"))
  var zookeeper: EmbeddedZookeeper = null
  var server1: KafkaServer = null
  var server2: KafkaServer = null
  var server3: KafkaServer = null
  var metadataStore: TopicMetadataStore = null

  @BeforeClass
  def beforeSetupServers {
    zookeeper = new EmbeddedZookeeper(zkConnect)
    server1 = TestUtils.createServer(new KafkaConfig(props1))
    server2 = TestUtils.createServer(new KafkaConfig(props2))
    server3 = TestUtils.createServer(new KafkaConfig(props3))
    zkClient = new ZkClient(zkConnect + "/", 6000, 6000, ZKStringSerializer)
    producer = new Producer(producerConfig)
    metadataStore = new ClientUtilTopicMetadataStore(brokers, "some-job-name")

    createTopics
    validateTopics
  }

  def createTopics {
    AdminUtils.createTopic(
      zkClient,
      INPUT_TOPIC,
      TOTAL_PARTITIONS,
      REPLICATION_FACTOR)

    AdminUtils.createTopic(
      zkClient,
      STATE_TOPIC,
      TOTAL_PARTITIONS,
      REPLICATION_FACTOR)
  }

  def validateTopics {
    val topics = Set(STATE_TOPIC, INPUT_TOPIC)
    var done = false
    var retries = 0

    while (!done && retries < 100) {
      try {
        val topicMetadataMap = TopicMetadataCache.getTopicMetadata(topics, "kafka", metadataStore.getTopicInfo)

        topics.foreach(topic => {
          val topicMetadata = topicMetadataMap(topic)
          val errorCode = topicMetadata.errorCode

          ErrorMapping.maybeThrowException(errorCode)
        })

        done = true
      } catch {
        case e: Exception =>
          System.err.println("Got exception while validating test topics. Waiting and retrying.", e)
          retries += 1
          Thread.sleep(500)
      }
    }

    if (retries >= 100) {
      fail("Unable to successfully create topics. Tried to validate %s times." format retries)
    }
  }

  @AfterClass
  def afterCleanLogDirs {
    server1.shutdown
    server1.awaitShutdown()
    server2.shutdown
    server2.awaitShutdown()
    server3.shutdown
    server3.awaitShutdown()
    Utils.rm(server1.config.logDirs)
    Utils.rm(server2.config.logDirs)
    Utils.rm(server3.config.logDirs)
    zkClient.close
    zookeeper.shutdown
  }
}

/**
 * Test that does the following:
 *
 * 1. Starts ZK, and 3 kafka brokers.
 * 2. Create two topics: input and mystore.
 * 3. Validate that the topics were created successfully and have leaders.
 * 4. Start a single partition of TestTask using LocalJobFactory/ThreadJob.
 * 5. Send four messages to input (1,2,3,2), which contain one dupe (2).
 * 6. Validate that all messages were received by TestTask.
 * 7. Validate that TestTask called store.put() for all four messages, and that the messages ended up in the mystore topic.
 * 8. Kill the job.
 * 9. Start the job again.
 * 10. Validate that the job restored all messages (1,2,3) to the store.
 * 11. Send three more messages to input (4,5,5), and validate that TestTask receives them.
 * 12. Kill the job again.
 */
class TestStatefulTask {
  import TestStatefulTask._

  val jobConfig = Map(
    "job.factory.class" -> "org.apache.samza.job.local.LocalJobFactory",
    "job.name" -> "hello-stateful-world",
    "task.class" -> "org.apache.samza.test.integration.TestTask",
    "task.inputs" -> "kafka.input",
    "serializers.registry.string.class" -> "org.apache.samza.serializers.StringSerdeFactory",
    "stores.mystore.factory" -> "org.apache.samza.storage.kv.KeyValueStorageEngineFactory",
    "stores.mystore.key.serde" -> "string",
    "stores.mystore.msg.serde" -> "string",
    "stores.mystore.changelog" -> "kafka.mystore",

    "systems.kafka.samza.factory" -> "org.apache.samza.system.kafka.KafkaSystemFactory",
    // Always start consuming at offset 0. This avoids a race condition between
    // the producer and the consumer in this test (SAMZA-166, SAMZA-224).
    "systems.kafka.samza.offset.default" -> "oldest", // applies to a nonempty topic
    "systems.kafka.consumer.auto.offset.reset" -> "smallest", // applies to an empty topic
    "systems.kafka.samza.msg.serde" -> "string",
    "systems.kafka.consumer.zookeeper.connect" -> zkConnect,
    "systems.kafka.producer.metadata.broker.list" -> ("localhost:%s" format port1))

  @Test
  def testShouldStartAndRestore {
    // Have to do this in one test to guarantee ordering.
    testShouldStartTaskForFirstTime
    testShouldRestoreStore
  }

  def testShouldStartTaskForFirstTime {
    val (job, task) = startJob

    // Validate that restored is empty.
    assertEquals(0, task.initFinished.getCount)
    assertEquals(0, task.restored.size)
    assertEquals(0, task.received.size)
{'fraction_non_alphanumeric': 0.09666131252868289, 'fraction_numerical': 0.014570904084442405, 'mean_word_length': 4.016690647482014, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 43, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, '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': '986318', 'n_tokens_mistral': 5558, 'n_tokens_neox': 5340, 'n_words': 1363}