KAFKA-9338; Fetch session should cache request leader epoch

core/src/main/scala/kafka/coordinator/group/GroupMetadata.scala

@@ -27,8 +27,6 @@ import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.message.JoinGroupResponseData.JoinGroupResponseMember
 import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.protocol.types.SchemaException
-import org.apache.kafka.common.record.RecordBatch.NO_PRODUCER_ID
-import org.apache.kafka.common.requests.JoinGroupRequest
 import org.apache.kafka.common.utils.Time
 
 import scala.collection.{Seq, immutable, mutable}

core/src/main/scala/kafka/server/FetchSession.scala

@@ -76,44 +76,42 @@ class CachedPartition(val topic: String,
                       var maxBytes: Int,
                       var fetchOffset: Long,
                       var highWatermark: Long,
+                      var leaderEpoch: Optional[Integer],
                       var fetcherLogStartOffset: Long,
                       var localLogStartOffset: Long)
     extends ImplicitLinkedHashCollection.Element {
 
   var cachedNext: Int = ImplicitLinkedHashCollection.INVALID_INDEX
   var cachedPrev: Int = ImplicitLinkedHashCollection.INVALID_INDEX
 
-  override def next = cachedNext
-  override def setNext(next: Int) = this.cachedNext = next
-  override def prev = cachedPrev
-  override def setPrev(prev: Int) = this.cachedPrev = prev
+  override def next: Int = cachedNext
+  override def setNext(next: Int): Unit = this.cachedNext = next
+  override def prev: Int = cachedPrev
+  override def setPrev(prev: Int): Unit = this.cachedPrev = prev
 
   def this(topic: String, partition: Int) =
-    this(topic, partition, -1, -1, -1, -1, -1)
+    this(topic, partition, -1, -1, -1, Optional.empty(), -1, -1)
 
   def this(part: TopicPartition) =
     this(part.topic, part.partition)
 
   def this(part: TopicPartition, reqData: FetchRequest.PartitionData) =
-    this(part.topic, part.partition,
-      reqData.maxBytes, reqData.fetchOffset, -1,
-      reqData.logStartOffset, -1)
+    this(part.topic, part.partition, reqData.maxBytes, reqData.fetchOffset, -1,
+      reqData.currentLeaderEpoch, reqData.logStartOffset, -1)
 
   def this(part: TopicPartition, reqData: FetchRequest.PartitionData,
            respData: FetchResponse.PartitionData[Records]) =
-    this(part.topic, part.partition,
-      reqData.maxBytes, reqData.fetchOffset, respData.highWatermark,
-      reqData.logStartOffset, respData.logStartOffset)
+    this(part.topic, part.partition, reqData.maxBytes, reqData.fetchOffset, respData.highWatermark,
+      reqData.currentLeaderEpoch, reqData.logStartOffset, respData.logStartOffset)
 
-  def topicPartition = new TopicPartition(topic, partition)
-
-  def reqData = new FetchRequest.PartitionData(fetchOffset, fetcherLogStartOffset, maxBytes, Optional.empty())
+  def reqData = new FetchRequest.PartitionData(fetchOffset, fetcherLogStartOffset, maxBytes, leaderEpoch)
 
   def updateRequestParams(reqData: FetchRequest.PartitionData): Unit = {
     // Update our cached request parameters.
     maxBytes = reqData.maxBytes
     fetchOffset = reqData.fetchOffset
     fetcherLogStartOffset = reqData.logStartOffset
+    leaderEpoch = reqData.currentLeaderEpoch
   }
 
   /**
@@ -158,19 +156,21 @@ class CachedPartition(val topic: String,
     mustRespond
   }
 
-  override def hashCode = (31 * partition) + topic.hashCode
+  override def hashCode: Int = (31 * partition) + topic.hashCode
 
   def canEqual(that: Any) = that.isInstanceOf[CachedPartition]
 
   override def equals(that: Any): Boolean =
     that match {
-      case that: CachedPartition => that.canEqual(this) &&
-        this.topic.equals(that.topic) &&
-        this.partition.equals(that.partition)
+      case that: CachedPartition =>
+        this.eq(that) ||
+          (that.canEqual(this) &&
+            this.partition.equals(that.partition) &&
+            this.topic.equals(that.topic))
       case _ => false
     }
 
-  override def toString = synchronized {
+  override def toString: String = synchronized {
     "CachedPartition(topic=" + topic +
       ", partition=" + partition +
       ", maxBytes=" + maxBytes +
@@ -197,12 +197,12 @@ class CachedPartition(val topic: String,
   *                     FetchSessionCache#touch.
   * @param epoch        The fetch session sequence number.
   */
-case class FetchSession(val id: Int,
-                        val privileged: Boolean,
-                        val partitionMap: FetchSession.CACHE_MAP,
-                        val creationMs: Long,
-                        var lastUsedMs: Long,
-                        var epoch: Int) {
+class FetchSession(val id: Int,
+                   val privileged: Boolean,
+                   val partitionMap: FetchSession.CACHE_MAP,
+                   val creationMs: Long,
+                   var lastUsedMs: Long,
+                   var epoch: Int) {
   // This is used by the FetchSessionCache to store the last known size of this session.
   // If this is -1, the Session is not in the cache.
   var cachedSize = -1
@@ -405,9 +405,9 @@ class IncrementalFetchContext(private val time: Time,
   override def foreachPartition(fun: (TopicPartition, FetchRequest.PartitionData) => Unit): Unit = {
     // Take the session lock and iterate over all the cached partitions.
     session.synchronized {
-      session.partitionMap.iterator.asScala.foreach(part => {
+      session.partitionMap.iterator.asScala.foreach { part =>
         fun(new TopicPartition(part.topic, part.partition), part.reqData)
-      })
+      }
     }
   }
 
@@ -501,15 +501,12 @@ class IncrementalFetchContext(private val time: Time,
   }
 }
 
-case class LastUsedKey(val lastUsedMs: Long,
-                       val id: Int) extends Comparable[LastUsedKey] {
+case class LastUsedKey(lastUsedMs: Long, id: Int) extends Comparable[LastUsedKey] {
   override def compareTo(other: LastUsedKey): Int =
     (lastUsedMs, id) compare (other.lastUsedMs, other.id)
 }
 
-case class EvictableKey(val privileged: Boolean,
-                        val size: Int,
-                        val id: Int) extends Comparable[EvictableKey] {
+case class EvictableKey(privileged: Boolean, size: Int, id: Int) extends Comparable[EvictableKey] {
   override def compareTo(other: EvictableKey): Int =
     (privileged, size, id) compare (other.privileged, other.size, other.id)
 }

core/src/test/scala/unit/kafka/server/FetchRequestTest.scala

@@ -248,6 +248,55 @@ class FetchRequestTest extends BaseRequestTest {
     assertResponseErrorForEpoch(Errors.FENCED_LEADER_EPOCH, followerId, Optional.of(secondLeaderEpoch - 1))
   }
 
+  @Test
+  def testEpochValidationWithinFetchSession(): Unit = {
+    val topic = "topic"
+    val topicPartition = new TopicPartition(topic, 0)
+    val partitionToLeader = TestUtils.createTopic(zkClient, topic, numPartitions = 1, replicationFactor = 3, servers)
+    val firstLeaderId = partitionToLeader(topicPartition.partition)
+
+    // We need a leader change in order to check epoch fencing since the first epoch is 0 and
+    // -1 is treated as having no epoch at all
+    killBroker(firstLeaderId)
+
+    val secondLeaderId = TestUtils.awaitLeaderChange(servers, topicPartition, firstLeaderId)
+    val secondLeaderEpoch = TestUtils.findLeaderEpoch(secondLeaderId, topicPartition, servers)
+    verifyFetchSessionErrors(topicPartition, secondLeaderEpoch, secondLeaderId)
+
+    val followerId = TestUtils.findFollowerId(topicPartition, servers)
+    verifyFetchSessionErrors(topicPartition, secondLeaderEpoch, followerId)
+  }
+
+  private def verifyFetchSessionErrors(topicPartition: TopicPartition,
+                                       leaderEpoch: Int,
+                                       destinationBrokerId: Int): Unit = {
+    val partitionMap = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    partitionMap.put(topicPartition, new FetchRequest.PartitionData(0L, 0L, 1024,
+      Optional.of(leaderEpoch)))
+    val fetchRequest = FetchRequest.Builder.forConsumer(0, 1, partitionMap)
+      .metadata(JFetchMetadata.INITIAL)
+      .build()
+    val fetchResponse = sendFetchRequest(destinationBrokerId, fetchRequest)
+    val sessionId = fetchResponse.sessionId
+
+    def assertResponseErrorForEpoch(expectedError: Errors,
+                                    sessionFetchEpoch: Int,
+                                    leaderEpoch: Optional[Integer]): Unit = {
+      val partitionMap = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+      partitionMap.put(topicPartition, new FetchRequest.PartitionData(0L, 0L, 1024, leaderEpoch))
+      val fetchRequest = FetchRequest.Builder.forConsumer(0, 1, partitionMap)
+        .metadata(new JFetchMetadata(sessionId, sessionFetchEpoch))
+        .build()
+      val fetchResponse = sendFetchRequest(destinationBrokerId, fetchRequest)
+      val partitionData = fetchResponse.responseData.get(topicPartition)
+      assertEquals(expectedError, partitionData.error)
+    }
+
+    // We only check errors because we do not expect the partition in the response otherwise
+    assertResponseErrorForEpoch(Errors.FENCED_LEADER_EPOCH, 1, Optional.of(leaderEpoch - 1))
+    assertResponseErrorForEpoch(Errors.UNKNOWN_LEADER_EPOCH, 2, Optional.of(leaderEpoch + 1))
+  }
+
   /**
    * Tests that down-conversions don't leak memory. Large down conversions are triggered
    * in the server. The client closes its connection after reading partial data when the

core/src/test/scala/unit/kafka/server/FetchSessionTest.scala

@@ -54,7 +54,7 @@ class FetchSessionTest {
 
   private def dummyCreate(size: Int)() = {
     val cacheMap = new FetchSession.CACHE_MAP(size)
-    for (i <- 0 to (size - 1)) {
+    for (i <- 0 until size) {
       cacheMap.add(new CachedPartition("test", i))
     }
     cacheMap
@@ -128,6 +128,66 @@ class FetchSessionTest {
 
   val EMPTY_PART_LIST = Collections.unmodifiableList(new util.ArrayList[TopicPartition]())
 
+
+  @Test
+  def testCachedLeaderEpoch(): Unit = {
+    val time = new MockTime()
+    val cache = new FetchSessionCache(10, 1000)
+    val fetchManager = new FetchManager(time, cache)
+
+    val tp0 = new TopicPartition("foo", 0)
+    val tp1 = new TopicPartition("foo", 1)
+    val tp2 = new TopicPartition("bar", 1)
+
+    def cachedLeaderEpochs(context: FetchContext): Map[TopicPartition, Optional[Integer]] = {
+      val mapBuilder = Map.newBuilder[TopicPartition, Optional[Integer]]
+      context.foreachPartition((tp, data) => mapBuilder += tp -> data.currentLeaderEpoch)
+      mapBuilder.result()
+    }
+
+    val request1 = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    request1.put(tp0, new FetchRequest.PartitionData(0, 0, 100, Optional.empty()))
+    request1.put(tp1, new FetchRequest.PartitionData(10, 0, 100, Optional.of(1)))
+    request1.put(tp2, new FetchRequest.PartitionData(10, 0, 100, Optional.of(2)))
+
+    val context1 = fetchManager.newContext(JFetchMetadata.INITIAL, request1, EMPTY_PART_LIST, false)
+    val epochs1 = cachedLeaderEpochs(context1)
+    assertEquals(Optional.empty(), epochs1(tp0))
+    assertEquals(Optional.of(1), epochs1(tp1))
+    assertEquals(Optional.of(2), epochs1(tp2))
+
+    val response = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[Records]]
+    response.put(tp0, new FetchResponse.PartitionData(Errors.NONE, 100, 100,
+      100, null, null))
+    response.put(tp1, new FetchResponse.PartitionData(
+      Errors.NONE, 10, 10, 10, null, null))
+    response.put(tp2, new FetchResponse.PartitionData(
+      Errors.NONE, 5, 5, 5, null, null))
+
+    val sessionId = context1.updateAndGenerateResponseData(response).sessionId()
+
+    // With no changes, the cached epochs should remain the same
+    val request2 = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    val context2 = fetchManager.newContext(new JFetchMetadata(sessionId, 1), request2, EMPTY_PART_LIST, false)
+    val epochs2 = cachedLeaderEpochs(context2)
+    assertEquals(Optional.empty(), epochs1(tp0))
+    assertEquals(Optional.of(1), epochs2(tp1))
+    assertEquals(Optional.of(2), epochs2(tp2))
+    context2.updateAndGenerateResponseData(response).sessionId()
+
+    // Now verify we can change the leader epoch and the context is updated
+    val request3 = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    request3.put(tp0, new FetchRequest.PartitionData(0, 0, 100, Optional.of(6)))
+    request3.put(tp1, new FetchRequest.PartitionData(10, 0, 100, Optional.empty()))
+    request3.put(tp2, new FetchRequest.PartitionData(10, 0, 100, Optional.of(3)))
+
+    val context3 = fetchManager.newContext(new JFetchMetadata(sessionId, 2), request3, EMPTY_PART_LIST, false)
+    val epochs3 = cachedLeaderEpochs(context3)
+    assertEquals(Optional.of(6), epochs3(tp0))
+    assertEquals(Optional.empty(), epochs3(tp1))
+    assertEquals(Optional.of(3), epochs3(tp2))
+  }
+
   @Test
   def testFetchRequests(): Unit = {
     val time = new MockTime()