Fix receive getting stuck after seek in MultiTopicsConsumer (#336)

* Fix MultiTopicsConsumer seek stuck

* Add test "Seek wont stuck the receive"

* Improve codes

* Address comments

Author: RobertIndie

src/Pulsar.Client/Internal/MultiTopicsConsumerImpl.fs

@@ -375,15 +375,24 @@ type internal MultiTopicsConsumerImpl<'T> (consumerConfig: ConsumerConfiguration
         | Ok msg -> incomingMessagesSize <- incomingMessagesSize + msg.Data.LongLength
         | _ -> ()
         incomingMessages.Enqueue(m)
+        
+    let tryResumePoller() =
+        if isPollingAllowed() && (waitingPoller <> defaultWaitingPoller) then
+            Log.Logger.LogDebug("{0} resume poller", prefix)
+            waitingPoller.SetResult()
+            waitingPoller <- defaultWaitingPoller
+    
+    let clearIncomingMessages() =
+        incomingMessages.Clear()
+        incomingMessagesSize <- 0L
+        tryResumePoller()
 
     let dequeueMessage() =
         let m = incomingMessages.Dequeue()
         match m with
         | Ok msg -> incomingMessagesSize <- incomingMessagesSize - msg.Data.LongLength
         | _ -> ()
-        if isPollingAllowed() && (waitingPoller <> defaultWaitingPoller) then
-            waitingPoller.SetResult()
-            waitingPoller <- defaultWaitingPoller
+        tryResumePoller()
         m
 
     let hasEnoughMessagesForBatchReceive() =
@@ -678,6 +687,7 @@ type internal MultiTopicsConsumerImpl<'T> (consumerConfig: ConsumerConfiguration
                         replyWithBatch ch
                 // check if should reply to poller immediately
                 if isPollingAllowed() |> not then
+                    Log.Logger.LogDebug("{0} paused poller, incomingMessages={1}, sharedQueueResumeThreshold={2}", prefix, incomingMessages.Count, sharedQueueResumeThreshold)
                     waitingPoller <- pollerChannel
                 else
                     pollerChannel.SetResult()
@@ -755,9 +765,8 @@ type internal MultiTopicsConsumerImpl<'T> (consumerConfig: ConsumerConfiguration
                             |> Seq.map(fun (KeyValue(_, (consumer, _))) -> consumer.RedeliverUnacknowledgedMessagesAsync())
                             |> Task.WhenAll
                         unAckedMessageTracker.Clear()
-                        incomingMessages.Clear()
+                        clearIncomingMessages()
                         currentStream.RestartCompletedTasks()
-                        incomingMessagesSize <- 0L
                         channel |> Option.map _.SetResult() |> ignore
                     with ex ->
                         Log.Logger.LogError(ex, "{0} RedeliverUnacknowledgedMessages failed", prefix)
@@ -838,16 +847,16 @@ type internal MultiTopicsConsumerImpl<'T> (consumerConfig: ConsumerConfiguration
                     Log.Logger.LogDebug("{0} Seek {1}", prefix, seekData)
                     backgroundTask {
                         try
+                            unAckedMessageTracker.Clear()
+                            clearIncomingMessages()
                             let! _ =
                                 consumers
                                 |> Seq.map (fun (KeyValue(_, (consumer, _))) ->
                                     match seekData with
                                     | SeekType.Timestamp ts -> consumer.SeekAsync(ts)
                                     | SeekType.MessageId msgId -> consumer.SeekAsync(msgId))
                                 |> Task.WhenAll
-                            unAckedMessageTracker.Clear()
-                            incomingMessages.Clear()
-                            incomingMessagesSize <- 0L
+                            currentStream.RestartCompletedTasks()
                             channel.SetResult()
                         with Flatten ex ->
                             channel.SetException ex
@@ -856,13 +865,13 @@ type internal MultiTopicsConsumerImpl<'T> (consumerConfig: ConsumerConfiguration
             | SeekWithResolver (resolver, channel) ->
                 backgroundTask {
                     try
+                        unAckedMessageTracker.Clear()
+                        clearIncomingMessages()
                         let! _ =
                             consumers
                             |> Seq.map (fun (KeyValue(_, (consumer, _))) -> consumer.SeekAsync(resolver))
                             |> Task.WhenAll
-                        unAckedMessageTracker.Clear()
-                        incomingMessages.Clear()
-                        incomingMessagesSize <- 0L
+                        currentStream.RestartCompletedTasks()
                         channel.SetResult()
                     with Flatten ex ->
                         channel.SetException ex

tests/IntegrationTests/Seek.fs

@@ -3,6 +3,7 @@ module Pulsar.Client.IntegrationTests.Seek
 open System
 open System.Threading
 open System.Diagnostics
+open System.Collections.Generic
 
 open Expecto
 open Expecto.Flip
@@ -240,5 +241,94 @@ let tests =
         testTask "Seek randomly works without batching " {
             do! testRandomSeek true 
         }
+        
+        
+        testTask "Seek won't get stuck at the receive in MultiTopicsConsumer" {
+            Log.Debug("Started Seek won't get stuck at the receive in MultiTopicsConsumer")
+            let client = getClient()
+            let topicName = "persistent://public/default/multi-topic-seek"
+            let producerName = "seekStuckProducer"
+            let consumerName = "seekStuckConsumer"
+            let numberOfMessages = 30
+            let subscriptionName = "test-seek-stuck-" + Guid.NewGuid().ToString("N")
+            
+            let seekWithRetry (consumer: IConsumer<byte[]>) (targetTimestamp: TimeStamp) (maxRetries: int) =
+                task {
+                    let mutable retryCount = 0
+                    let mutable success = false
+                    while retryCount < maxRetries && not success do
+                        try
+                            do! consumer.SeekAsync(targetTimestamp)
+                            success <- true
+                        with Flatten ex ->
+                            match ex with
+                            | :? NotConnectedException as notConnectedEx ->
+                                retryCount <- retryCount + 1
+                                if retryCount >= maxRetries then
+                                    Log.Error("SeekAsync failed after {0} retries: {1}", maxRetries, notConnectedEx.Message)
+                                    raise notConnectedEx
+                                else
+                                    Log.Debug("SeekAsync failed (attempt {0}/{1}): {2}. Retrying in 1 second...", retryCount, maxRetries, notConnectedEx.Message)
+                                    do! Task.Delay(1000)
+                            | _ ->
+                                raise ex
+                }
+            
+            let! consumer =
+                client.NewConsumer()
+                    .Topic(topicName)
+                    .ConsumerName(consumerName)
+                    .SubscriptionName(subscriptionName)
+                    .ReceiverQueueSize(10)
+                    .SubscribeAsync()
+            let! (producer : IProducer<byte[]>) =
+                client.NewProducer()
+                    .Topic(topicName)
+                    .ProducerName(producerName)
+                    .EnableBatching(false)
+                    .CreateAsync()
+            
+            let expectedMessages = HashSet<string>()
+            for i in 1..numberOfMessages do
+                let messageContent = sprintf "Message-%i-%s" i (Guid.NewGuid().ToString("N"))
+                expectedMessages.Add(messageContent) |> ignore
+                let messageBytes = Encoding.UTF8.GetBytes(messageContent)
+                let! (_ : MessageId) = producer.SendAsync(messageBytes)
+                ()
+            do! Task.Delay(1000)
+            
+            let targetTimestamp = %(DateTimeOffset.UtcNow.ToUnixTimeMilliseconds() - 1000L * 3600L * 24L)
+            Log.Debug("Seeking to timestamp: {0}", targetTimestamp)
+            do! seekWithRetry consumer targetTimestamp 10
+            
+            let receivedMessages = HashSet<string>()
+            let cts = new CancellationTokenSource(TimeSpan.FromSeconds(30.0))
+            
+            try
+                for _ in 1..numberOfMessages do
+                    let! (message : Message<byte[]>) = consumer.ReceiveAsync(cts.Token)
+                    let received = Encoding.UTF8.GetString(message.Data)
+                    Log.Debug("{0} received {1}", consumerName, received)
+                    receivedMessages.Add(received) |> ignore
+                    do! consumer.AcknowledgeAsync(message.MessageId)
+                
+                Expect.equal "" numberOfMessages receivedMessages.Count
+                for expectedMsg in expectedMessages do
+                    Expect.isTrue "" (receivedMessages.Contains(expectedMsg))
+                
+                cts.Dispose()
+            with
+            | :? OperationCanceledException
+            | :? TaskCanceledException ->
+                cts.Dispose()
+                let errorMsg = $"Test timeout: Only received {receivedMessages.Count} out of {numberOfMessages} messages within 30 seconds"
+                Log.Error(errorMsg)
+                failwith errorMsg
+            | ex ->
+                cts.Dispose()
+                raise ex
+            
+            Log.Debug("Finished Seek won't get stuck at the receive in MultiTopicsConsumer")
+        }
 
     ]

tests/compose/standalone/scripts/init-standalone.sh

@@ -22,6 +22,7 @@ bin/pulsar-admin topics create-partitioned-topic persistent://public/default/par
 bin/pulsar-admin topics create-partitioned-topic persistent://public/default/partitioned6 --partitions 2
 bin/pulsar-admin topics create-partitioned-topic persistent://public/default/partitioned-dl-test --partitions 2
 bin/pulsar-admin topics create-partitioned-topic persistent://public/deduplication/partitioned --partitions 3
+bin/pulsar-admin topics create-partitioned-topic persistent://public/default/multi-topic-seek --partitions 3
 
 echo "Initializing transaction coordinator metadata..."
 bin/pulsar initialize-transaction-coordinator-metadata -cs standalone:2181 -c standalone --initial-num-transaction-coordinators 2