RebalanceExtSpec failure: `SubSource` caught in die / recreate loop for partition

[Aaronontheweb]

Version Information
Version of Akka.NET? dev branch (v1.5.39-beta2)
Which Akka.NET Modules?

Describe the bug

Looks like a CommittablePartitionedSource got stuck in a loop on #464 - see a lot of this in the logs:

[DEBUG][03/16/2025 23:45:03.061Z][Thread 0029][RebalanceExtTests (akka://RebalanceExtTests)] [consumer-1] Consuming partitioned source clientId: consumer-1, for tp: topic-1-1b046062-6731-496d-8524-58cef5321605 [[1]]
[INFO][03/16/2025 23:45:03.062Z][Thread 0029][SourceShape`1([] [out])(akka://RebalanceExtTests/user/StreamSupervisor-49)] 51 Starting SubSource for partition topic-1-1b046062-6731-496d-8524-58cef5321605 [[1]]
[DEBUG][03/16/2025 23:45:03.062Z][Thread 0029][RebalanceExtTests (akka://RebalanceExtTests)] [consumer-1] Consuming partitioned source clientId: consumer-1, for tp: topic-1-1b046062-6731-496d-8524-58cef5321605 [[1]]
[INFO][03/16/2025 23:45:03.062Z][Thread 0029][SourceShape`1([] [out])(akka://RebalanceExtTests/user/StreamSupervisor-49)] 51 Starting SubSource for partition topic-1-1b046062-6731-496d-8524-58cef5321605 [[1]]
[DEBUG][03/16/2025 23:45:03.062Z][Thread 0029][RebalanceExtTests (akka://RebalanceExtTests)] [consumer-1] Consuming partitioned source clientId: consumer-1, for tp: topic-1-1b046062-6731-496d-8524-58cef5321605 [[1]]
[INFO][03/16/2025 23:45:03.063Z][Thread 0029][SourceShape`1([] [out])(akka://RebalanceExtTests/user/StreamSupervisor-49)] 51 Starting SubSource for partition topic-1-1b046062-6731-496d-8524-58cef5321605 [[1]]
[DEBUG][03/16/2025 23:45:03.063Z][Thread 0029][RebalanceExtTests (akka://RebalanceExtTests)] [consumer-1] Consuming partitioned source clientId: consumer-1, for tp: topic-1-1b046062-6731-496d-8524-58cef5321605 [[1]]

To Reproduce

Links to working reproductions on Github / Gitlab are very much appreciated

Expected behavior

Stream should complete normally.

Actual behavior

Stream was stuck in a terminal loop.

Logs

See attached:

Standard_Console_Output (6).log

Environment

Azure DevOps

Additional context

It's very possible that the old BatchCommiterFlow logic might be responsible for this in and the fixes in #458 resolve it.

[Aaronontheweb]

Looks like the failure happened here:

   at System.Threading.Tasks.Task.ThrowIfExceptional(Boolean includeTaskCanceledExceptions)
   at System.Threading.Tasks.Task`1.GetResultCore(Boolean waitCompletionNotification)
   at Akka.TestKit.TestKitBase.WithinAsync[T](TimeSpan min, TimeSpan max, Func`1 function, String hint, Nullable`1 epsilonValue, CancellationToken cancellationToken)
   at Akka.TestKit.TestKitBase.WithinAsync(TimeSpan max, Func`1 actionAsync, Nullable`1 epsilonValue, CancellationToken cancellationToken)
   at Akka.Streams.Kafka.Tests.Integration.RebalanceExtTests.FetchedRecordsMustNotBeLostUponAbort() in /home/vsts/work/1/s/src/Akka.Streams.Kafka.Tests/Integration/RebalanceExtTests.cs:line 216
--- End of stack trace from previous location ---
----- Inner Stack Trace -----
   at Akka.TestKit.Xunit2.XunitAssertions.Fail(String format, Object[] args)
   at Akka.TestKit.TestKitBase.InternalExpectMsgEnvelopeAsync[T](Nullable`1 timeout, Action`2 assert, String hint, CancellationToken cancellationToken, Boolean shouldLog)
   at Akka.TestKit.TestKitBase.InternalExpectMsgEnvelopeAsync[T](Nullable`1 timeout, Action`1 msgAssert, Action`1 senderAssert, String hint, CancellationToken cancellationToken)
   at Akka.TestKit.TestKitBase.InternalExpectMsgAsync[T](Nullable`1 timeout, Action`1 msgAssert, Action`1 senderAssert, String hint, CancellationToken cancellationToken)
   at Akka.Streams.Kafka.Tests.Integration.RebalanceExtTests.<FetchedRecordsMustNotBeLostUponAbort>b__11_0() in /home/vsts/work/1/s/src/Akka.Streams.Kafka.Tests/Integration/RebalanceExtTests.cs:line 313
   at Akka.TestKit.TestKitBase.<>c__DisplayClass217_0.<<WithinAsync>b__0>d.MoveNext()

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

Lines 313 to 315 in d60c4f7

	await probe2RebalanceActor.ExpectMsgAsync<TopicPartitionsRevoked>();
	var assigned2again = await probe2RebalanceActor.ExpectMsgAsync<TopicPartitionsAssigned>();
	Assert.Equal(partitionCount, assigned2again.Partitions.Count);

[Aaronontheweb]

Logs that are being written out over and over again are

Akka.Streams.Kafka/src/Akka.Streams.Kafka/Stages/Consumers/Abstract/SubSourceLogic.cs

Line 665 in 7f906ed

Log.Info("{0} Starting SubSource for partition {1}", _actorNumber, _topicPartition);

and

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

Line 79 in 7f906ed

Log.Debug("{0} Consuming partitioned source clientId: {1}, for tp: {2}", logPrefix, clientId, tp);

So from the looks of this, it appears that the CommittablePartitionedSource keeps emitting partition topic-1-1b046062-6731-496d-8524-58cef5321605 [[1]] constantly.

[Aaronontheweb]

Pretty sure this is what's driving the re-create loop:

Akka.Streams.Kafka/src/Akka.Streams.Kafka/Stages/Consumers/Abstract/SubSourceLogic.cs

Lines 400 to 429 in 7f906ed

	private void HandleSubsourceCancelled((TopicPartition, ISubSourceCancellationStrategy) obj)
	{
	var (topicPartition, cancellationStrategy) = obj;
	_subSources = _subSources.Remove(topicPartition);
	_partitionsInStartup = _partitionsInStartup.Remove(topicPartition);
	switch (cancellationStrategy)
	{
	case SeekToOffsetAndReEmit seek:
	var offset = seek.Offset;
	// re-add this partition to pending partitions so it can be re-emitted
	_pendingPartitions = _pendingPartitions.Add(topicPartition);
	if (Log.IsDebugEnabled)
	Log.Debug("#{0} Seeking {1} to {2} after partition SubSource cancelled", _actorNumber,
	topicPartition, offset);
	var topicPartitionOffset = new TopicPartitionOffset(topicPartition, offset);
	SeekAndEmitSubSources(
	formerlyUnknown: ImmutableHashSet<TopicPartition>.Empty,
	offsets: ImmutableList.Create(topicPartitionOffset).ToImmutableHashSet());
	break;
	case ReEmit:
	// re-add this partition to pending partitions so it can be re-emitted
	_pendingPartitions = _pendingPartitions.Add(topicPartition);
	EmitSubSourcesForPendingPartitions();
	break;
	case DoNothing _:
	break;
	}
	}

Edit: specifically the ReEmit handler

[Aaronontheweb]

Earlier in the log we can see the partition being revoked:

[DEBUG][03/16/2025 23:45:01.242Z][Thread 0029][SourceShape`1([] [out])(akka://RebalanceExtTests/user/StreamSupervisor-49)] Filtering out messages from revoked partitions [topic-1-1b046062-6731-496d-8524-58cef5321605 [[1]] @3]

This is right before the spin starts. I wonder if the race condition here actually comes from #452

[Aaronontheweb]

Looks like a downstream stage is completing and this is what triggers the cancellation of the sub-stream

Akka.Streams.Kafka/src/Akka.Streams.Kafka/Stages/Consumers/Abstract/SubSourceLogic.cs

Lines 652 to 657 in 7f906ed

	SetHandler(shape.Outlet, onPull: Pump, onDownstreamFinish: ex =>
	{
	subSourceCancelledCallback((
	topicPartition, OnDownstreamFinishSubSourceCancellationStrategy));
	base.InternalOnDownstreamFinish(ex);
	});

Probably because we invoked its kill-switch right before the area where the Exception occurred:

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

Lines 293 to 294 in 0691502

	// abort Consumer-1
	sharedKillSwitch1.Abort(new Exception($"abort {ConsumerClientId1} messageId=2"));

So each time the stream gets materialized it's going to complete immediately.

[Aaronontheweb]

Looks like this an issue with the RebelanceExtTest - there's a race condition between the stage being cancelled and the partition being revoked.

The underlying KillSwitch the substream is being materialized with has already completed, and each time the stream is "cancelled" it gets recreated because the parent (partitioned stream) has to do its cleanup of revoked substreams on a timer - default duration of 1 second:

Akka.Streams.Kafka/src/Akka.Streams.Kafka/Stages/Consumers/Abstract/SubSourceLogic.cs

Lines 393 to 398 in 7f906ed

	private void HandlePartitionsRevoked(IImmutableSet<TopicPartitionOffset> revoked)
	{
	_partitionsToRevoke = _partitionsToRevoke.Union(revoked.Select(r => r.TopicPartition));
	ScheduleOnce(CloseRevokedPartitions.Instance, _settings.WaitClosePartition);
	}

The cancelling / re-creation of this stream is happening thousands of times per second after the original revocation. Not really sure this is fixable TBH - if you're re-creating a stream and against an already-completed stage that's an end-user error in and of itself.

[Aaronontheweb]

One easy fix here would be simply not allowing the re-creation of partitions that are in the _partitionsToRevoke list, since they're already on their way out.