kafka.tools.TestLogCleaning.scala Maven / Gradle / Ivy
The newest version!
/**
* 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 kafka.tools
import joptsimple.OptionParser
import java.util.Properties
import java.util.Random
import java.io._
import kafka.consumer._
import kafka.serializer._
import kafka.utils._
import kafka.log.FileMessageSet
import kafka.log.Log
import org.apache.kafka.clients.producer.{ProducerRecord, KafkaProducer, ProducerConfig}
/**
* This is a torture test that runs against an existing broker. Here is how it works:
*
* It produces a series of specially formatted messages to one or more partitions. Each message it produces
* it logs out to a text file. The messages have a limited set of keys, so there is duplication in the key space.
*
* The broker will clean its log as the test runs.
*
* When the specified number of messages have been produced we create a consumer and consume all the messages in the topic
* and write that out to another text file.
*
* Using a stable unix sort we sort both the producer log of what was sent and the consumer log of what was retrieved by the message key.
* Then we compare the final message in both logs for each key. If this final message is not the same for all keys we
* print an error and exit with exit code 1, otherwise we print the size reduction and exit with exit code 0.
*/
object TestLogCleaning {
def main(args: Array[String]) {
val parser = new OptionParser
val numMessagesOpt = parser.accepts("messages", "The number of messages to send or consume.")
.withRequiredArg
.describedAs("count")
.ofType(classOf[java.lang.Long])
.defaultsTo(Long.MaxValue)
val numDupsOpt = parser.accepts("duplicates", "The number of duplicates for each key.")
.withRequiredArg
.describedAs("count")
.ofType(classOf[java.lang.Integer])
.defaultsTo(5)
val brokerOpt = parser.accepts("broker", "Url to connect to.")
.withRequiredArg
.describedAs("url")
.ofType(classOf[String])
val topicsOpt = parser.accepts("topics", "The number of topics to test.")
.withRequiredArg
.describedAs("count")
.ofType(classOf[java.lang.Integer])
.defaultsTo(1)
val percentDeletesOpt = parser.accepts("percent-deletes", "The percentage of updates that are deletes.")
.withRequiredArg
.describedAs("percent")
.ofType(classOf[java.lang.Integer])
.defaultsTo(0)
val zkConnectOpt = parser.accepts("zk", "Zk url.")
.withRequiredArg
.describedAs("url")
.ofType(classOf[String])
val sleepSecsOpt = parser.accepts("sleep", "Time to sleep between production and consumption.")
.withRequiredArg
.describedAs("ms")
.ofType(classOf[java.lang.Integer])
.defaultsTo(0)
val dumpOpt = parser.accepts("dump", "Dump the message contents of a topic partition that contains test data from this test to standard out.")
.withRequiredArg
.describedAs("directory")
.ofType(classOf[String])
val options = parser.parse(args:_*)
if(args.length == 0)
CommandLineUtils.printUsageAndDie(parser, "An integration test for log cleaning.")
if(options.has(dumpOpt)) {
dumpLog(new File(options.valueOf(dumpOpt)))
System.exit(0)
}
CommandLineUtils.checkRequiredArgs(parser, options, brokerOpt, zkConnectOpt, numMessagesOpt)
// parse options
val messages = options.valueOf(numMessagesOpt).longValue
val percentDeletes = options.valueOf(percentDeletesOpt).intValue
val dups = options.valueOf(numDupsOpt).intValue
val brokerUrl = options.valueOf(brokerOpt)
val topicCount = options.valueOf(topicsOpt).intValue
val zkUrl = options.valueOf(zkConnectOpt)
val sleepSecs = options.valueOf(sleepSecsOpt).intValue
val testId = new Random().nextInt(Int.MaxValue)
val topics = (0 until topicCount).map("log-cleaner-test-" + testId + "-" + _).toArray
println("Producing %d messages...".format(messages))
val producedDataFile = produceMessages(brokerUrl, topics, messages, dups, percentDeletes)
println("Sleeping for %d seconds...".format(sleepSecs))
Thread.sleep(sleepSecs * 1000)
println("Consuming messages...")
val consumedDataFile = consumeMessages(zkUrl, topics)
val producedLines = lineCount(producedDataFile)
val consumedLines = lineCount(consumedDataFile)
val reduction = 1.0 - consumedLines.toDouble/producedLines.toDouble
println("%d rows of data produced, %d rows of data consumed (%.1f%% reduction).".format(producedLines, consumedLines, 100 * reduction))
println("De-duplicating and validating output files...")
validateOutput(producedDataFile, consumedDataFile)
producedDataFile.delete()
consumedDataFile.delete()
}
def dumpLog(dir: File) {
require(dir.exists, "Non-existent directory: " + dir.getAbsolutePath)
for(file <- dir.list.sorted; if file.endsWith(Log.LogFileSuffix)) {
val ms = new FileMessageSet(new File(dir, file))
for(entry <- ms) {
val key = Utils.readString(entry.message.key)
val content =
if(entry.message.isNull)
null
else
Utils.readString(entry.message.payload)
println("offset = %s, key = %s, content = %s".format(entry.offset, key, content))
}
}
}
def lineCount(file: File): Int = io.Source.fromFile(file).getLines.size
def validateOutput(producedDataFile: File, consumedDataFile: File) {
val producedReader = externalSort(producedDataFile)
val consumedReader = externalSort(consumedDataFile)
val produced = valuesIterator(producedReader)
val consumed = valuesIterator(consumedReader)
val producedDedupedFile = new File(producedDataFile.getAbsolutePath + ".deduped")
val producedDeduped = new BufferedWriter(new FileWriter(producedDedupedFile), 1024*1024)
val consumedDedupedFile = new File(consumedDataFile.getAbsolutePath + ".deduped")
val consumedDeduped = new BufferedWriter(new FileWriter(consumedDedupedFile), 1024*1024)
var total = 0
var mismatched = 0
while(produced.hasNext && consumed.hasNext) {
val p = produced.next()
producedDeduped.write(p.toString)
producedDeduped.newLine()
val c = consumed.next()
consumedDeduped.write(c.toString)
consumedDeduped.newLine()
if(p != c)
mismatched += 1
total += 1
}
producedDeduped.close()
consumedDeduped.close()
require(!produced.hasNext, "Additional values produced not found in consumer log.")
require(!consumed.hasNext, "Additional values consumed not found in producer log.")
println("Validated " + total + " values, " + mismatched + " mismatches.")
require(mismatched == 0, "Non-zero number of row mismatches.")
// if all the checks worked out we can delete the deduped files
producedDedupedFile.delete()
consumedDedupedFile.delete()
}
def valuesIterator(reader: BufferedReader) = {
new IteratorTemplate[TestRecord] {
def makeNext(): TestRecord = {
var next = readNext(reader)
while(next != null && next.delete)
next = readNext(reader)
if(next == null)
allDone()
else
next
}
}
}
def readNext(reader: BufferedReader): TestRecord = {
var line = reader.readLine()
if(line == null)
return null
var curr = new TestRecord(line)
while(true) {
line = peekLine(reader)
if(line == null)
return curr
val next = new TestRecord(line)
if(next == null || next.topicAndKey != curr.topicAndKey)
return curr
curr = next
reader.readLine()
}
null
}
def peekLine(reader: BufferedReader) = {
reader.mark(4096)
val line = reader.readLine
reader.reset()
line
}
def externalSort(file: File): BufferedReader = {
val builder = new ProcessBuilder("sort", "--key=1,2", "--stable", "--buffer-size=20%", "--temporary-directory=" + System.getProperty("java.io.tmpdir"), file.getAbsolutePath)
val process = builder.start()
new Thread() {
override def run() {
val exitCode = process.waitFor()
if(exitCode != 0) {
System.err.println("Process exited abnormally.")
while(process.getErrorStream.available > 0) {
System.err.write(process.getErrorStream().read())
}
}
}
}.start()
new BufferedReader(new InputStreamReader(process.getInputStream()), 10*1024*1024)
}
def produceMessages(brokerUrl: String,
topics: Array[String],
messages: Long,
dups: Int,
percentDeletes: Int): File = {
val producerProps = new Properties
producerProps.setProperty(ProducerConfig.BLOCK_ON_BUFFER_FULL_CONFIG, "true")
producerProps.setProperty(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerUrl)
producerProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArraySerializer")
producerProps.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArraySerializer")
val producer = new KafkaProducer[Array[Byte],Array[Byte]](producerProps)
val rand = new Random(1)
val keyCount = (messages / dups).toInt
val producedFile = File.createTempFile("kafka-log-cleaner-produced-", ".txt")
println("Logging produce requests to " + producedFile.getAbsolutePath)
val producedWriter = new BufferedWriter(new FileWriter(producedFile), 1024*1024)
for(i <- 0L until (messages * topics.length)) {
val topic = topics((i % topics.length).toInt)
val key = rand.nextInt(keyCount)
val delete = i % 100 < percentDeletes
val msg =
if(delete)
new ProducerRecord[Array[Byte],Array[Byte]](topic, key.toString.getBytes(), null)
else
new ProducerRecord[Array[Byte],Array[Byte]](topic, key.toString.getBytes(), i.toString.getBytes())
producer.send(msg)
producedWriter.write(TestRecord(topic, key, i, delete).toString)
producedWriter.newLine()
}
producedWriter.close()
producer.close()
producedFile
}
def makeConsumer(zkUrl: String, topics: Array[String]): ZookeeperConsumerConnector = {
val consumerProps = new Properties
consumerProps.setProperty("group.id", "log-cleaner-test-" + new Random().nextInt(Int.MaxValue))
consumerProps.setProperty("zookeeper.connect", zkUrl)
consumerProps.setProperty("consumer.timeout.ms", (20*1000).toString)
consumerProps.setProperty("auto.offset.reset", "smallest")
new ZookeeperConsumerConnector(new ConsumerConfig(consumerProps))
}
def consumeMessages(zkUrl: String, topics: Array[String]): File = {
val connector = makeConsumer(zkUrl, topics)
val streams = connector.createMessageStreams(topics.map(topic => (topic, 1)).toMap, new StringDecoder, new StringDecoder)
val consumedFile = File.createTempFile("kafka-log-cleaner-consumed-", ".txt")
println("Logging consumed messages to " + consumedFile.getAbsolutePath)
val consumedWriter = new BufferedWriter(new FileWriter(consumedFile))
for(topic <- topics) {
val stream = streams(topic).head
try {
for(item <- stream) {
val delete = item.message == null
val value = if(delete) -1L else item.message.toLong
consumedWriter.write(TestRecord(topic, item.key.toInt, value, delete).toString)
consumedWriter.newLine()
}
} catch {
case e: ConsumerTimeoutException =>
}
}
consumedWriter.close()
connector.shutdown()
consumedFile
}
}
case class TestRecord(val topic: String, val key: Int, val value: Long, val delete: Boolean) {
def this(pieces: Array[String]) = this(pieces(0), pieces(1).toInt, pieces(2).toLong, pieces(3) == "d")
def this(line: String) = this(line.split("\t"))
override def toString() = topic + "\t" + key + "\t" + value + "\t" + (if(delete) "d" else "u")
def topicAndKey = topic + key
}