KAFKA-10244 An new java interface to replace 'kafka.common.MessageReader'

clients/src/main/java/org/apache/kafka/tools/RecordReader.java

@@ -0,0 +1,51 @@
+/*
+ * 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.kafka.tools;
+
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.Configurable;
+
+import java.io.Closeable;
+import java.io.InputStream;
+import java.util.Iterator;
+import java.util.Map;
+
+/**
+ * Typical implementations of this interface convert data from an `InputStream` received via `readRecords` into a
+ * iterator of `ProducerRecord` instance. Noted that the implementations to have a public nullary constructor.
+ *
+ * This is used by the `kafka.tools.ConsoleProducer`.
+ */
+public interface RecordReader extends Closeable, Configurable {
+
+    default void configure(Map<String, ?> configs) {}
+
+    /**
+     * read byte array from input stream and then generate a iterator of producer record
+     * @param inputStream of messages. the implementation does not need to close the input stream.
+     * @return a iterator of producer record. It should implement following rules. 1) the hasNext() method must be idempotent.
+     *         2) the convert error should be thrown by next() method.
+     */
+    Iterator<ProducerRecord<byte[], byte[]>> readRecords(InputStream inputStream);
+
+
+    /**
+     * Closes this reader.
+     * This method is invoked if the iterator from readRecords either has no more records or throws exception.
+     */
+    default void close() {}
+}

core/src/main/scala/kafka/common/MessageReader.scala

@@ -28,6 +28,8 @@ import org.apache.kafka.clients.producer.ProducerRecord
   *
   * This is used by the `ConsoleProducer`.
   */
+@deprecated("This class has been deprecated and will be removed in 4.0." +
+  "Please use org.apache.kafka.tools.RecordReader instead", "3.5.0")
 trait MessageReader {
 
   def init(inputStream: InputStream, props: Properties): Unit = {}

core/src/main/scala/kafka/tools/ConsoleProducer.scala

@@ -24,33 +24,81 @@ import java.util.regex.Pattern
 import joptsimple.{OptionException, OptionParser, OptionSet}
 import kafka.common.MessageReader
 import kafka.utils.Implicits._
-import kafka.utils.{Exit, ToolsUtils}
+import kafka.utils.{Exit, Logging, ToolsUtils}
 import org.apache.kafka.clients.producer.internals.ErrorLoggingCallback
-import org.apache.kafka.clients.producer.{KafkaProducer, ProducerConfig, ProducerRecord}
+import org.apache.kafka.clients.producer.{KafkaProducer, Producer, ProducerConfig, ProducerRecord}
 import org.apache.kafka.common.KafkaException
 import org.apache.kafka.common.record.CompressionType
 import org.apache.kafka.common.utils.Utils
 import org.apache.kafka.server.util.{CommandDefaultOptions, CommandLineUtils}
+import org.apache.kafka.tools.RecordReader
+
+import scala.annotation.nowarn
+
+@nowarn("cat=deprecation")
+object ConsoleProducer extends Logging {
+
+  private[tools] def newReader(className: String, prop: Properties): RecordReader = {
+    val reader = Class.forName(className).getDeclaredConstructor().newInstance()
+    reader match {
+      case r: RecordReader =>
+        r.configure(prop.asInstanceOf[java.util.Map[String, _]])
+        r
+      case r: MessageReader =>
+        logger.warn("MessageReader is deprecated. Please use org.apache.kafka.server.tools.RecordReader instead")
+        new RecordReader {
+          private[this] var initialized = false
+
+          override def readRecords(inputStream: InputStream): java.util.Iterator[ProducerRecord[Array[Byte], Array[Byte]]] = {
+            if (initialized) throw new IllegalStateException("It is invalid to call readRecords again when the reader is based on deprecated MessageReader")
+            if (!initialized) {
+              r.init(inputStream, prop)
+              initialized = true
+            }
+            new java.util.Iterator[ProducerRecord[Array[Byte], Array[Byte]]] {
+              private[this] var current: ProducerRecord[Array[Byte], Array[Byte]] = null
+              // a flag used to avoid accessing readMessage again after it does return null
+              private[this] var done: Boolean = false
+
+              override def hasNext: Boolean = {
+                if (current != null) true
+                else if (done) false
+                else {
+                  current = r.readMessage()
+                  done = current == null
+                  !done
+                }
+              }
+
+              override def next(): ProducerRecord[Array[Byte], Array[Byte]] =
+                try if (hasNext) current
+                else throw new NoSuchElementException("no more records from input stream")
+                finally current = null
+            }
+          }
+          override def close(): Unit = r.close()
+        }
+      case _ => throw new IllegalArgumentException(f"the reader must extend ${classOf[RecordReader].getName}")
+    }
+  }
 
-object ConsoleProducer {
+  private[tools] def loopReader(producer: Producer[Array[Byte], Array[Byte]],
+                               reader: RecordReader,
+                                inputStream: InputStream,
+                               sync: Boolean): Unit = {
+    val iter = reader.readRecords(inputStream)
+    try while (iter.hasNext) send(producer, iter.next(), sync) finally reader.close()
+  }
 
   def main(args: Array[String]): Unit = {
 
     try {
-        val config = new ProducerConfig(args)
-        val reader = Class.forName(config.readerClass).getDeclaredConstructor().newInstance().asInstanceOf[MessageReader]
-        reader.init(System.in, getReaderProps(config))
-
-        val producer = new KafkaProducer[Array[Byte], Array[Byte]](producerProps(config))
-
-    Exit.addShutdownHook("producer-shutdown-hook", producer.close)
-
-        var record: ProducerRecord[Array[Byte], Array[Byte]] = null
-        do {
-          record = reader.readMessage()
-          if (record != null)
-            send(producer, record, config.sync)
-        } while (record != null)
+      val config = new ProducerConfig(args)
+      val input = System.in
+      val producer = new KafkaProducer[Array[Byte], Array[Byte]](producerProps(config))
+      try loopReader(producer, newReader(config.readerClass, getReaderProps(config)), input, config.sync)
+      finally producer.close()
+      Exit.exit(0)
     } catch {
       case e: joptsimple.OptionException =>
         System.err.println(e.getMessage)
@@ -59,10 +107,9 @@ object ConsoleProducer {
         e.printStackTrace
         Exit.exit(1)
     }
-    Exit.exit(0)
   }
 
-  private def send(producer: KafkaProducer[Array[Byte], Array[Byte]],
+  private def send(producer: Producer[Array[Byte], Array[Byte]],
                          record: ProducerRecord[Array[Byte], Array[Byte]], sync: Boolean): Unit = {
     if (sync)
       producer.send(record).get()

core/src/test/scala/unit/kafka/tools/ConsoleProducerTest.scala

@@ -17,14 +17,21 @@
 
 package kafka.tools
 
+import kafka.common.MessageReader
+
 import java.nio.file.Files
 import kafka.tools.ConsoleProducer.LineMessageReader
 import kafka.utils.{Exit, TestUtils}
-import org.apache.kafka.clients.producer.ProducerConfig
+import org.apache.kafka.clients.producer.{Producer, ProducerConfig, ProducerRecord}
+import org.apache.kafka.tools.RecordReader
 import org.junit.jupiter.api.Assertions.{assertEquals, assertThrows, assertTrue}
 import org.junit.jupiter.api.Test
+import org.mockito.Mockito
 
+import java.io.InputStream
 import java.util
+import java.util.Properties
+import scala.annotation.nowarn
 
 class ConsoleProducerTest {
 
@@ -220,4 +227,61 @@ class ConsoleProducerTest {
       producerConfig.getInt(ProducerConfig.BATCH_SIZE_CONFIG))
   }
 
+  @Test
+  def testNewReader(): Unit = {
+    ConsoleProducerTest.configureCount = 0
+    ConsoleProducerTest.closeCount = 0
+    val reader = ConsoleProducer.newReader(classOf[ConsoleProducerTest.DumbMessageReader].getName, new Properties())
+    // the deprecated MessageReader get configured when creating records
+    assertEquals(0, ConsoleProducerTest.configureCount)
+    reader.readRecords(System.in)
+    assertEquals(1, ConsoleProducerTest.configureCount)
+    assertEquals(0, ConsoleProducerTest.closeCount)
+    assertThrows(classOf[IllegalStateException], () => reader.readRecords(System.in))
+    reader.close()
+    assertEquals(1, ConsoleProducerTest.closeCount)
+
+    ConsoleProducerTest.configureCount = 0
+    ConsoleProducerTest.closeCount = 0
+
+    val reader1 = ConsoleProducer.newReader(classOf[ConsoleProducerTest.DumbRecordReader].getName, new Properties())
+    assertEquals(1, ConsoleProducerTest.configureCount)
+    assertEquals(0, ConsoleProducerTest.closeCount)
+    reader1.close()
+    assertEquals(1, ConsoleProducerTest.closeCount)
+  }
+
+  @Test
+  def testLoopReader(): Unit = {
+    ConsoleProducerTest.configureCount = 0
+    ConsoleProducerTest.closeCount = 0
+    val reader = ConsoleProducer.newReader(classOf[ConsoleProducerTest.DumbRecordReader].getName, new Properties())
+
+    ConsoleProducer.loopReader(Mockito.mock(classOf[Producer[Array[Byte], Array[Byte]]]),
+      reader, System.in, false)
+
+    assertEquals(1, ConsoleProducerTest.configureCount)
+    assertEquals(1, ConsoleProducerTest.closeCount)
+  }
+}
+
+@nowarn("cat=deprecation")
+object ConsoleProducerTest {
+  var configureCount = 0
+  var closeCount = 0
+  class DumbMessageReader extends MessageReader {
+    override def init(inputStream: InputStream, props: Properties): Unit = configureCount += 1
+    override def readMessage(): ProducerRecord[Array[Byte], Array[Byte]] = null
+
+    override def close(): Unit = closeCount += 1
+
+  }
+
+  class DumbRecordReader extends RecordReader {
+    override def configure(configs: util.Map[String, _]): Unit = configureCount += 1
+    override def readRecords(inputStream: InputStream): java.util.Iterator[ProducerRecord[Array[Byte], Array[Byte]]] =
+      java.util.Collections.emptyIterator()
+
+    override def close(): Unit = closeCount += 1
+  }
 }