Revert #418 (#431)

* Implement OnPartitionsLost event handler from #418

* Bump AkkaVersion to 1.5.37 (#422)

(cherry picked from commit d2aa70f)

* Implement dual targeting for netstandard2.0 and net6.0 (#424)

(cherry picked from commit 6c895c5)

* Cleanup linq by introducing Partition (#425)

* Make sure that we prevent downstream to propagate null completion exception (#427)

(cherry picked from commit 562212d)

* Update RELEASE_NOTES.md for 1.5.37 release (#423)

* Update RELEASE_NOTES.md for 1.5.37 release

* Update RELEASE_NOTES.md

(cherry picked from commit d6b0a42)

* Bump Confluent.Kafka to 2.8.0 (#428)

(cherry picked from commit 02c002e)

* Make sure that restricted consumer is populated properly (#429)

(cherry picked from commit 8e41f71)

* Add CancellationToken support to MockCustomer (#430)

(cherry picked from commit 2d11235)

* Fix missing OnPartitionLost handler

* Simplify exception handling

---------

Co-authored-by: Aaron Stannard <aaron@petabridge.com>

Author: Arkatufus

src/Akka.Streams.Kafka.Tests/Integration/PlainPartitionedSourceIntegrationTests.cs

@@ -216,7 +216,7 @@ public async Task PlainPartitionedSource_should_not_leave_gaps_when_subsource_is
                 .TakeWhile(m => m < totalMessages, inclusive: true)
                 .RunWith(Sink.Last<int>(), Materializer);
 
-            var consumedMessages = await consumedMessagesTask.ShouldCompleteWithin(60.Seconds());
+            var consumedMessages = await consumedMessagesTask.ShouldCompleteWithin(10.Seconds());
             consumedMessages.Should().Be(totalMessages);
         }
 

src/Akka.Streams.Kafka/Extensions/CollectionExtensions.cs

@@ -43,7 +43,7 @@ public static IImmutableSet<T> ToImmutableSet<T>(this IEnumerable<T> collection)
         /// <summary>
         /// Split a list into two list depending on the boolean return value of the predicate.
         /// </summary>
-        public static (IImmutableList<T> True, IImmutableList<T> False) Partition<T>(this IImmutableList<T> list, Predicate<T> predicate)
+        public static (List<T> True, List<T> False) Partition<T>(this List<T> list, Predicate<T> predicate)
         {
             var left = new List<T>();
             var right = new List<T>();
@@ -55,7 +55,7 @@ public static (IImmutableList<T> True, IImmutableList<T> False) Partition<T>(thi
                     right.Add(t);
             }
 
-            return (left.ToImmutableList(), right.ToImmutableList());
+            return (left, right);
         }        
     }
 }

src/Akka.Streams.Kafka/Stages/Consumers/Actors/KafkaConsumerActor.cs

@@ -82,7 +82,7 @@ internal class KafkaConsumerActor<K, V> : ActorBase, ILogReceive
         /// While `true`, committing is delayed.
         /// Changed by `onPartitionsRevoked` and `onPartitionsAssigned` callbacks
         /// </summary>
-        private AtomicBoolean _rebalanceInProgress = new();
+        private bool _rebalanceInProgress = false;
         /// <summary>
         /// Keeps commit offsets during rebalances for later commit.
         /// </summary>
@@ -116,11 +116,42 @@ public KafkaConsumerActor(IActorRef owner, ConsumerSettings<K, V> settings, Deci
         }
 
         #region Rebalance listener
+        
+        internal sealed class PartitionAssigned
+        {
+            public PartitionAssigned(IImmutableSet<TopicPartition> partitions)
+            {
+                Partitions = partitions;
+            }
+
+            public IImmutableSet<TopicPartition> Partitions { get; }
+        }
+        
+        internal sealed class PartitionRevoked
+        {
+            public PartitionRevoked(IImmutableSet<TopicPartitionOffset> partitions)
+            {
+                Partitions = partitions;
+            }
+
+            public IImmutableSet<TopicPartitionOffset> Partitions { get; }
+        }
+        
+        internal sealed class PartitionLost
+        {
+            public PartitionLost(IImmutableSet<TopicPartitionOffset> partitions)
+            {
+                Partitions = partitions;
+            }
+
+            public IImmutableSet<TopicPartitionOffset> Partitions { get; }
+        }
+    
         // This is RebalanceListener.OnPartitionAssigned on JVM
         private void PartitionsAssignedHandler(IImmutableSet<TopicPartition> partitions)
         {
             var assignment = _consumer.Assignment;
-            var partitionsToPause = partitions.Where(p => assignment.Contains(p)).ToImmutableList();
+            var partitionsToPause = partitions.Where(p => assignment.Contains(p)).ToList();
             PausePartitions(partitionsToPause);
 
             _commitRefreshing.AssignedPositions(partitions, _consumer, _settings.PositionTimeout);
@@ -130,7 +161,7 @@ private void PartitionsAssignedHandler(IImmutableSet<TopicPartition> partitions)
             watch.Stop();
             CheckDuration(watch, "onAssign");
 
-            _rebalanceInProgress.GetAndSet(false);
+            _rebalanceInProgress = false;
         }
 
         // This is RebalanceListener.OnPartitionRevoked on JVM
@@ -142,10 +173,10 @@ private void PartitionsRevokedHandler(IImmutableSet<TopicPartitionOffset> partit
             CheckDuration(watch, "onRevoke");
 
             _commitRefreshing.Revoke(partitions.Select(tp => tp.TopicPartition).ToImmutableHashSet());
-            _rebalanceInProgress.GetAndSet(true);
+            _rebalanceInProgress = true;
         }
 
-        // This is RebalanceListener.OnPartitionLost on JVM
+        // This is RebalanceListener.OnPartitionRevoked on JVM
         private void PartitionsLostHandler(IImmutableSet<TopicPartitionOffset> partitions)
         {
             var watch = Stopwatch.StartNew();
@@ -154,12 +185,12 @@ private void PartitionsLostHandler(IImmutableSet<TopicPartitionOffset> partition
             CheckDuration(watch, "onLost");
 
             _commitRefreshing.Revoke(partitions.Select(tp => tp.TopicPartition).ToImmutableHashSet());
-            _rebalanceInProgress.GetAndSet(true);
+            _rebalanceInProgress = true;
         }
-        
+
         private void RebalancePostStop()
         {
-            var currentTopicPartitions = _consumer.Assignment.ToImmutableList();
+            var currentTopicPartitions = _consumer.Assignment;
             PausePartitions(currentTopicPartitions);
 
             var watch = Stopwatch.StartNew();
@@ -274,6 +305,23 @@ protected override bool Receive(object message)
                     Sender.Tell(HandleMetadataRequest(req));
                     return true;
 
+                // Rebalance callbacks
+                case PartitionAssigned evt:
+                    PartitionsAssignedHandler(evt.Partitions);
+                    return true;
+                
+                case PartitionRevoked evt:
+                    PartitionsRevokedHandler(evt.Partitions);
+                    return true;
+                
+                case PartitionLost evt:
+                    PartitionsLostHandler(evt.Partitions);
+                    return true;
+                
+                case Status.Failure fail:
+                    ProcessExceptions(fail.Cause);
+                    return true;
+                
                 default:
                     return false;
             }
@@ -305,11 +353,12 @@ private void ApplySettings(ConsumerSettings<K, V> updatedSettings)
                 if (_log.IsDebugEnabled)
                     _log.Debug($"Creating Kafka consumer with settings: {JsonConvert.SerializeObject(_settings)}");
 
+                var localSelf = Self;
                 _consumer = _settings.CreateKafkaConsumer(
-                    consumeErrorHandler: (c, e) => ProcessExceptions(new KafkaException(e)),
-                    partitionAssignedHandler: (c, tp) => PartitionsAssignedHandler(tp.ToImmutableHashSet()),
-                    partitionRevokedHandler: (c, tp) => PartitionsRevokedHandler(tp.ToImmutableHashSet()),
-                    partitionLostHandler: (c, tp) => PartitionsLostHandler(tp.ToImmutableHashSet()),
+                    consumeErrorHandler: (c, e) => localSelf.Tell(new Status.Failure(new KafkaException(e))),
+                    partitionAssignedHandler: (c, tp) => localSelf.Tell(new PartitionAssigned(tp.ToImmutableHashSet())),
+                    partitionRevokedHandler: (c, tp) => localSelf.Tell(new PartitionRevoked(tp.ToImmutableHashSet())),
+                    partitionLostHandler: (c, tp) => localSelf.Tell(new PartitionLost(tp.ToImmutableHashSet())),
                     statisticHandler: (c, json) => _statisticsHandler.OnStatistics(c, json));
 
                 var restrictedConsumerTimeoutMs = Math.Round(_settings.PartitionHandlerWarning.TotalMilliseconds * 0.95);
@@ -455,68 +504,62 @@ private void ReceivePoll(Internal.Poll<K, V> poll)
 
         private void Poll()
         {
-            var currentAssignment = _consumer.Assignment.ToImmutableList();
-            var initialRebalanceInProcess = _rebalanceInProgress.Value;
-
-            if (_requests.IsEmpty())
+            try
             {
-                if(_log.IsDebugEnabled)
-                    _log.Debug("Requests are empty - attempting to consume.");
-                PausePartitions(currentAssignment);
-                try
+                var currentAssignment = _consumer.Assignment;
+                var initialRebalanceInProcess = _rebalanceInProgress;
+
+                if (_requests.IsEmpty())
                 {
+                    if (_log.IsDebugEnabled)
+                        _log.Debug("Requests are empty - attempting to consume.");
+                    PausePartitions(currentAssignment);
                     var consumed = _consumer.Consume(0);
                     if (consumed != null)
                         throw new IllegalActorStateException("Consumed message should be null");
                 }
-                catch (Exception e)
-                {
-                    ProcessExceptions(e);
-                }
-            }
-            else
-            {
-                // Seek has to be done here because they can somehow fail.
-                // Would need to see if we can move this somewhere else
-                // because a seek can take up to 200ms to complete
-                foreach (var tpo in _seekedOffset.Select(kvp => kvp.Value))
+                else
                 {
-                    try
-                    {
-                        if(_log.IsDebugEnabled)
-                            _log.Debug("Seeking offset {0} in partition {1} for topic {2}", tpo.Offset, tpo.Partition, tpo.Topic);
-                        _consumer.Seek(tpo);
-                    }
-                    catch (Exception ex)
+                    // Seek has to be done here because they can somehow fail.
+                    // Would need to see if we can move this somewhere else
+                    // because a seek can take up to 200ms to complete
+                    foreach (var tpo in _seekedOffset.Select(kvp => kvp.Value))
                     {
-                        _log.Error(ex, $"{tpo.TopicPartition} Failed to seek to {tpo.Offset}: {ex}");
-                        throw;
+                        try
+                        {
+                            if (_log.IsDebugEnabled)
+                                _log.Debug("Seeking offset {0} in partition {1} for topic {2}", tpo.Offset,
+                                    tpo.Partition, tpo.Topic);
+                            _consumer.Seek(tpo);
+                        }
+                        catch (Exception ex)
+                        {
+                            _log.Error(ex, $"{tpo.TopicPartition} Failed to seek to {tpo.Offset}: {ex}");
+                            throw;
+                        }
                     }
-                }
-                
-                // resume partitions to fetch
-                IImmutableSet<TopicPartition> partitionsToFetch = _requests.Values.SelectMany(v => v.Topics).ToImmutableHashSet();
-                var (resumeThese, pauseThese) = currentAssignment.Partition(partitionsToFetch.Contains);
-                PausePartitions(pauseThese);
-                ResumePartitions(resumeThese);
 
-                using (var cts = new CancellationTokenSource(_settings.PollTimeout))
-                {
-                    var (polled, exception) = PollKafka(cts.Token);
-                    try
+                    // resume partitions to fetch
+                    IImmutableSet<TopicPartition> partitionsToFetch =
+                        _requests.Values.SelectMany(v => v.Topics).ToImmutableHashSet();
+                    var (resumeThese, pauseThese) = currentAssignment.Partition(partitionsToFetch.Contains);
+                    PausePartitions(pauseThese);
+                    ResumePartitions(resumeThese);
+
+                    using (var cts = new CancellationTokenSource(_settings.PollTimeout))
                     {
+                        var (polled, exception) = PollKafka(cts.Token);
                         ProcessResult(partitionsToFetch, polled);
-                    }
-                    catch (Exception e)
-                    {
-                        ProcessExceptions(e);
+                        ProcessExceptions(exception);
                     }
 
-                    ProcessExceptions(exception);
+                    CheckRebalanceState(initialRebalanceInProcess);
                 }
             }
-            
-            CheckRebalanceState(initialRebalanceInProcess);
+            catch (Exception e)
+            {
+                ProcessExceptions(e);
+            }
 
             if (_stopInProgress)
             {
@@ -661,7 +704,7 @@ private void Commit(IImmutableSet<TopicPartitionOffset> commitMap, Action<object
         /// </summary>
         private void CheckRebalanceState(bool initialRebalanceInProgress)
         {
-            if (initialRebalanceInProgress && !_rebalanceInProgress.Value && _rebalanceCommitSenders.Any())
+            if (initialRebalanceInProgress && !_rebalanceInProgress && _rebalanceCommitSenders.Any())
             {
                 _log.Debug($"Comitting stash {string.Join(", ", _rebalanceCommitStash)} replying to {string.Join(", ", _rebalanceCommitSenders)}");
                 var replyTo = _rebalanceCommitSenders;
@@ -671,7 +714,7 @@ private void CheckRebalanceState(bool initialRebalanceInProgress)
             }
         }
 
-        private void PausePartitions(IImmutableList<TopicPartition> partitions)
+        private void PausePartitions(List<TopicPartition> partitions)
         {
             if (partitions.Count == 0)
                 return;
@@ -682,17 +725,14 @@ private void PausePartitions(IImmutableList<TopicPartition> partitions)
             _resumedPartitions = _resumedPartitions.Except(partitions);
         }
 
-        private void ResumePartitions(IImmutableList<TopicPartition> partitions)
+        private void ResumePartitions(List<TopicPartition> partitions)
         {
             if (partitions.Count == 0)
                 return;
 
             var partitionsToResume = partitions.Except(_resumedPartitions).ToList();
-            if(partitionsToResume.Count == 0 && _log.IsDebugEnabled)
-            {
-                _log.Debug("Requested partitions already resumed. Resume request: [{0}], already resumed: [{1}]", string.Join(",", partitions), string.Join(",", _resumedPartitions));
+            if(partitionsToResume.Count == 0)
                 return;
-            }
 
             if(_log.IsDebugEnabled)
                 _log.Debug("Resuming partitions [{0}]", string.Join(",", partitionsToResume));