applied changes to fix issue LGouellec#328

core/Kafka/Internal/StreamsRebalanceListener.cs

@@ -55,16 +55,19 @@ public void PartitionsRevoked(IConsumer<byte[], byte[]> consumer, List<TopicPart
             DateTime start = DateTime.Now;
             lock (manager._lock)
             {
-                manager.RebalanceInProgress = true;
-                manager.RevokeTasks(new List<TopicPartition>(partitions.Select(p => p.TopicPartition)));
-                Thread.SetState(ThreadState.PARTITIONS_REVOKED);
-                manager.RebalanceInProgress = false;
+                if (Thread.IsRunning)
+                {
+                    manager.RebalanceInProgress = true;
+                    manager.RevokeTasks(new List<TopicPartition>(partitions.Select(p => p.TopicPartition)));
+                    Thread.SetState(ThreadState.PARTITIONS_REVOKED);
+                    manager.RebalanceInProgress = false;
 
-                StringBuilder sb = new StringBuilder();
-                sb.AppendLine($"Partition revocation took {DateTime.Now - start} ms");
-                sb.AppendLine(
-                    $"\tCurrent suspended active tasks: {string.Join(",", partitions.Select(p => $"{p.Topic}-{p.Partition}"))}");
-                log.LogInformation(sb.ToString());
+                    StringBuilder sb = new StringBuilder();
+                    sb.AppendLine($"Partition revocation took {DateTime.Now - start} ms");
+                    sb.AppendLine(
+                        $"\tCurrent suspended active tasks: {string.Join(",", partitions.Select(p => $"{p.Topic}-{p.Partition}"))}");
+                    log.LogInformation(sb.ToString());
+                }
             }
         }
 

core/KafkaStream.cs

@@ -259,6 +259,7 @@ public override string ToString()
         private readonly StreamMetricsRegistry metricsRegistry;
 
         private readonly CancellationTokenSource _cancelSource = new();
+        private readonly SequentiallyGracefullyShutdownHook shutdownHook;
 
         internal State StreamState { get; private set; }
 
@@ -401,6 +402,13 @@ string Protect(string str)
             queryableStoreProvider = new QueryableStoreProvider(stateStoreProviders, globalStateStoreProvider);
 
             StreamState = State.CREATED;
+
+            shutdownHook = new SequentiallyGracefullyShutdownHook(
+                threads,
+                globalStreamThread,
+                externalStreamThread,
+                _cancelSource
+            );
         }
 
         /// <summary>
@@ -424,13 +432,9 @@ public void Start(CancellationToken? token = null)
         public async Task StartAsync(CancellationToken? token = null)
         {
             if (token.HasValue)
-            {
-                token.Value.Register(() => {
-                    _cancelSource.Cancel();
-                    Dispose();
-                });
-            }
-            await Task.Factory.StartNew(async () =>
+                token.Value.Register(Dispose);
+
+            await Task.Factory.StartNew(() =>
             {
                 if (SetState(State.REBALANCING))
                 {
@@ -448,12 +452,12 @@ await Task.Factory.StartNew(async () =>
                             SetState(State.PENDING_SHUTDOWN);
                             SetState(State.ERROR);
                         }
-                        return;
+                        return Task.CompletedTask;
                     }
 
                     RunMiddleware(true, true);
 
-                    globalStreamThread?.Start(_cancelSource.Token);
+                    globalStreamThread?.Start();
                     externalStreamThread?.Start(_cancelSource.Token);
 
                     foreach (var t in threads)
@@ -463,14 +467,17 @@ await Task.Factory.StartNew(async () =>
 
                     RunMiddleware(false, true);
                 }
+
+
+                return Task.CompletedTask;
             }, token ?? _cancelSource.Token);
 
 
             try
             {
                 // Allow time for streams thread to run
-                await Task.Delay(TimeSpan.FromMilliseconds(configuration.StartTaskDelayMs),
-                    token ?? _cancelSource.Token);
+                //await Task.Delay(TimeSpan.FromMilliseconds(configuration.StartTaskDelayMs),
+                //    token ?? _cancelSource.Token);
             }
             catch
             {
@@ -484,16 +491,8 @@ await Task.Delay(TimeSpan.FromMilliseconds(configuration.StartTaskDelayMs),
         /// </summary>
         public void Dispose()
         {
-            Task.Factory.StartNew(() =>
-            {
-                if (!_cancelSource.IsCancellationRequested)
-                {
-                    _cancelSource.Cancel();
-                }
-
-                Close();
-                _cancelSource.Dispose();
-            }).Wait(TimeSpan.FromSeconds(30));
+            Close();
+            _cancelSource.Dispose();
         }
 
         /// <summary>
@@ -511,14 +510,8 @@ private void Close()
             else
             {
                 RunMiddleware(true, false);
-
-                foreach (var t in threads)
-                {
-                    t.Dispose();
-                }
 
-                externalStreamThread?.Dispose();
-                globalStreamThread?.Dispose();
+                shutdownHook.Shutdown();
 
                 RunMiddleware(false, false);
                 metricsRegistry.RemoveClientSensors();

core/Mock/ClusterInMemoryTopologyDriver.cs

@@ -212,7 +212,7 @@ void stateChangedHandeler(IThread thread, ThreadStateTransitionValidator old,
 
             InitializeInternalTopicManager();
 
-            globalStreamThread?.Start(token);
+            globalStreamThread?.Start();
             externalStreamThread?.Start(token);
 
             threadTopology.Start(token);

core/Mock/Kafka/MockCluster.cs

@@ -760,7 +760,8 @@ internal ConsumeResult<byte[], byte[]> Consume(MockConsumer mockConsumer, TimeSp
                                         Key = record.Key,
                                         Value = record.Value,
                                         Timestamp = new Timestamp(
-                                            record.Timestamp ?? DateTime.Now, TimestampType.NotAvailable)
+                                            record.Timestamp ?? DateTime.Now, TimestampType.NotAvailable),
+                                        Headers = record.Headers
                                     }
                                 };
                                 ++offset.OffsetConsumed;
@@ -799,8 +800,8 @@ internal DeliveryReport<byte[], byte[]> Produce(string topic, Message<byte[], by
                 partition = RandomGenerator.GetInt32(topics[topic].PartitionNumber);
             else
                 partition = Math.Abs(MurMurHash3.Hash(new MemoryStream(message.Key))) % topics[topic].PartitionNumber;
-            
-            topics[topic].AddMessage(message.Key, message.Value, partition, message.Timestamp.UnixTimestampMs);
+
+            topics[topic].AddMessage(message.Key, message.Value, partition, message.Timestamp.UnixTimestampMs, message.Headers);
 
             r.Message = message;
             r.Partition = partition;
@@ -822,13 +823,13 @@ internal DeliveryReport<byte[], byte[]> Produce(TopicPartition topicPartition, M
             CreateTopic(topicPartition.Topic);
             if (topics[topicPartition.Topic].PartitionNumber > topicPartition.Partition)
             {
-                topics[topicPartition.Topic].AddMessage(message.Key, message.Value, topicPartition.Partition, message.Timestamp.UnixTimestampMs);
+                topics[topicPartition.Topic].AddMessage(message.Key, message.Value, topicPartition.Partition, message.Timestamp.UnixTimestampMs, message.Headers); 
                 r.Status = PersistenceStatus.Persisted;
             }
             else
             {
                 topics[topicPartition.Topic].CreateNewPartitions(topicPartition.Partition);
-                topics[topicPartition.Topic].AddMessage(message.Key, message.Value, topicPartition.Partition, message.Timestamp.UnixTimestampMs);
+                topics[topicPartition.Topic].AddMessage(message.Key, message.Value, topicPartition.Partition, message.Timestamp.UnixTimestampMs, message.Headers);
                 r.Status = PersistenceStatus.Persisted;
             }
 

core/Mock/Kafka/MockPartition.cs

@@ -7,7 +7,7 @@ namespace Streamiz.Kafka.Net.Mock.Kafka
 {
     internal class MockPartition
     {
-        private readonly List<(byte[], byte[], long)> log = new();
+        private readonly List<(byte[], byte[], long, Headers)> log = new();
         private readonly Dictionary<long, long> mappingOffsets = new();
 
         public MockPartition(int indice)
@@ -21,10 +21,10 @@ public MockPartition(int indice)
         public long LowOffset { get; private set; } = Offset.Unset;
         public long HighOffset { get; private set; } = Offset.Unset;
 
-        internal void AddMessageInLog(byte[] key, byte[] value, long timestamp)
+        internal void AddMessageInLog(byte[] key, byte[] value, long timestamp, Headers headers)
         {
             mappingOffsets.Add(Size, log.Count);
-            log.Add((key, value, timestamp));
+            log.Add((key, value, timestamp, headers));
             ++Size;
             UpdateOffset();
         }
@@ -47,7 +47,8 @@ internal TestRecord<byte[], byte[]> GetMessage(long offset)
                 {
                     Key = record.Item1,
                     Value = record.Item2,
-                    Timestamp = record.Item3.FromMilliseconds()
+                    Timestamp = record.Item3.FromMilliseconds(),
+                    Headers = record.Item4
                 };
             }
 

core/Mock/Kafka/MockTopic.cs

@@ -23,9 +23,9 @@ public MockTopic(string topic, int part)
         public int PartitionNumber { get; private set; }
         public IEnumerable<MockPartition> Partitions => partitions.AsReadOnly();
 
-        public void AddMessage(byte[] key, byte[] value, int partition, long timestamp = 0)
+        public void AddMessage(byte[] key, byte[] value, int partition, long timestamp = 0, Headers headers = null)
         {
-            partitions[partition].AddMessageInLog(key, value, timestamp);
+            partitions[partition].AddMessageInLog(key, value, timestamp, headers);
         }
 
         public TestRecord<byte[], byte[]> GetMessage(int partition, long consumerOffset)

core/Mock/Sync/SyncPipeBuilder.cs

@@ -34,11 +34,13 @@ public void PublishRecord(string topic, byte[] key, byte[] value, DateTime times
             public void Flush()
             {
                 long now = DateTime.Now.GetMilliseconds();
+                TaskManager.CurrentTask = task;
                 while (task.CanProcess(now))
                     task.Process();
 
                 task.PunctuateStreamTime();
                 task.PunctuateSystemTime();
+                TaskManager.CurrentTask = null;
             }
 
             public void Close()

core/Processors/GlobalStreamThread.cs

@@ -109,8 +109,7 @@ public void Close()
         private readonly string logPrefix;
         private readonly string threadClientId;
         private readonly IConsumer<byte[], byte[]> globalConsumer;
-        private CancellationToken token;
-        private readonly object stateLock = new object();
+        private readonly object stateLock = new();
         private readonly IStreamConfig configuration;
         private StateConsumer stateConsumer;
         private readonly IGlobalStateMaintainer globalStateMaintainer;
@@ -140,7 +139,7 @@ private void Run()
             SetState(GlobalThreadState.RUNNING);
             try
             {
-                while (!token.IsCancellationRequested && State.IsRunning())
+                while (State.IsRunning())
                 {
                     stateConsumer.PollAndUpdate();
 
@@ -166,11 +165,11 @@ private void Run()
                     // https://docs.microsoft.com/en-us/visualstudio/code-quality/ca1065
                 }
 
-                Dispose(false);
+                //Dispose(false);
             }
         }
 
-        public void Start(CancellationToken token)
+        public void Start()
         {
             log.LogInformation("{LogPrefix}Starting", logPrefix);
 
@@ -185,8 +184,6 @@ public void Start(CancellationToken token)
                 throw;
             }
 
-            this.token = token;
-
             thread.Start();
         }
 

core/Processors/Internal/SequentiallyGracefullyShutdownHook.cs

@@ -0,0 +1,44 @@
+using System.Threading;
+using Microsoft.Extensions.Logging;
+using Streamiz.Kafka.Net.Crosscutting;
+
+namespace Streamiz.Kafka.Net.Processors.Internal
+{
+    internal class SequentiallyGracefullyShutdownHook
+    {
+        private readonly IThread[] _streamThreads;
+        private readonly GlobalStreamThread _globalStreamThread;
+        private readonly IThread _externalStreamThread;
+        private readonly CancellationTokenSource _tokenSource;
+        private readonly ILogger log = Logger.GetLogger(typeof(SequentiallyGracefullyShutdownHook));
+
+        public SequentiallyGracefullyShutdownHook(
+            IThread[] streamThreads,
+            GlobalStreamThread globalStreamThread,
+            IThread externalStreamThread,
+            CancellationTokenSource tokenSource
+            )
+        {
+            _streamThreads = streamThreads;
+            _globalStreamThread = globalStreamThread;
+            _externalStreamThread = externalStreamThread;
+            _tokenSource = tokenSource;
+        }
+
+        public void Shutdown()
+        {
+            log.LogInformation($"Request shutdown gracefully");
+            _tokenSource.Cancel();
+
+            foreach (var t in _streamThreads)
+            {
+                t.Dispose();
+            }
+
+            _externalStreamThread?.Dispose();
+            _globalStreamThread?.Dispose();
+
+            log.LogInformation($"Shutdown gracefully successful");
+        }
+    }
+}

core/StreamConfig.cs

@@ -2719,6 +2719,7 @@ public MetricsRecordingLevel MetricsRecording
         /// Time wait before completing the start task of <see cref="KafkaStream"/>. (default: 5000)
         /// </summary>
         [StreamConfigProperty("" + startTaskDelayMsCst)]
+        [Obsolete]
         public long StartTaskDelayMs
         {
             get => configProperties[startTaskDelayMsCst];