Commit offsets to kafka on reset if external state storage used (#35)

epam · Jul 10, 2024 · f685a26 · f685a26
1 parent 9cb6589
commit f685a26
Show file tree

Hide file tree

Showing 17 changed files with 73 additions and 84 deletions.
diff --git a/sample/Epam.Kafka.Sample/Program.cs b/sample/Epam.Kafka.Sample/Program.cs
@@ -34,8 +34,7 @@ private static void Main(string[] args)
             // view metrics in console for demo purposes only
             services.AddOpenTelemetry()
                 .WithMetrics(mb => mb
-                    .AddMeter(PipelineMonitor.StatusMeterName, PipelineMonitor.HealthMeterName, Statistics.MeterName,
-                        Statistics.TopicsMeterName)
+                    .AddMeter(PipelineMonitor.StatusMeterName, PipelineMonitor.HealthMeterName, Statistics.MeterName)
                     .AddConsoleExporter());
 
             KafkaBuilder kafkaBuilder = services.AddKafka()

diff --git a/src/Epam.Kafka.PubSub/Subscription/State/BatchState.cs b/src/Epam.Kafka.PubSub/Subscription/State/BatchState.cs
@@ -19,9 +19,9 @@ public IReadOnlyCollection<ConsumeResult<TKey, TValue>> GetBatch<TKey, TValue>(
 
         topic.ClearIfNotAssigned();
 
-        using (activitySpan.CreateSpan("assign"))
+        using (var span = activitySpan.CreateSpan("assign"))
         {
-            this.AssignConsumer(topic, cancellationToken);
+            this.AssignConsumer(topic, span, cancellationToken);
         }
 
         cancellationToken.ThrowIfCancellationRequested();
@@ -35,6 +35,7 @@ public IReadOnlyCollection<ConsumeResult<TKey, TValue>> GetBatch<TKey, TValue>(
     }
 
     protected abstract void AssignConsumer<TKey, TValue>(SubscriptionTopicWrapper<TKey, TValue> topic,
+        ActivityWrapper activitySpan,
         CancellationToken cancellationToken);
 
     public void CommitResults<TKey, TValue>(SubscriptionTopicWrapper<TKey, TValue> topic,

diff --git a/src/Epam.Kafka.PubSub/Subscription/State/CombinedState.cs b/src/Epam.Kafka.PubSub/Subscription/State/CombinedState.cs
@@ -18,11 +18,12 @@ public CombinedState(TOffsetsStorage offsetsStorage)
     }
 
     protected override void AssignConsumer<TKey, TValue>(SubscriptionTopicWrapper<TKey, TValue> topic,
+        ActivityWrapper activitySpan,
         CancellationToken cancellationToken)
     {
         topic.ExternalState = list => topic.GetAndResetState(this._offsetsStorage, list, cancellationToken);
 
-        base.AssignConsumer(topic, cancellationToken);
+        base.AssignConsumer(topic, activitySpan, cancellationToken);
 
         if (topic.Consumer.Assignment.Count > 0)
         {
@@ -49,6 +50,8 @@ protected override void AssignConsumer<TKey, TValue>(SubscriptionTopicWrapper<TK
             }
 
             topic.OnReset(reset);
+
+            topic.CommitOffsetIfNeeded(activitySpan, reset);
         }
     }
 

diff --git a/src/Epam.Kafka.PubSub/Subscription/State/ExternalState.cs b/src/Epam.Kafka.PubSub/Subscription/State/ExternalState.cs
@@ -19,6 +19,7 @@ public ExternalState(TOffsetsStorage offsetsStorage)
     }
 
     protected override void AssignConsumer<TKey, TValue>(SubscriptionTopicWrapper<TKey, TValue> topic,
+        ActivityWrapper activitySpan,
         CancellationToken cancellationToken)
     {
         if (topic == null)
@@ -63,6 +64,8 @@ protected override void AssignConsumer<TKey, TValue>(SubscriptionTopicWrapper<TK
         }
 
         topic.OnReset(reset);
+
+        topic.CommitOffsetIfNeeded(activitySpan, reset);
     }
 
     protected override IReadOnlyCollection<TopicPartitionOffset> CommitState<TKey, TValue>(

diff --git a/src/Epam.Kafka.PubSub/Subscription/State/ExternalStateExtensions.cs b/src/Epam.Kafka.PubSub/Subscription/State/ExternalStateExtensions.cs
@@ -17,10 +17,14 @@ public static IReadOnlyCollection<TopicPartitionOffset> CommitState<TKey, TValue
         CancellationToken cancellationToken)
     {
         if (topic == null)
+        {
             throw new ArgumentNullException(nameof(topic));
+        }
 
         if (offsets == null)
+        {
             throw new ArgumentNullException(nameof(offsets));
+        }
 
         var reset = new List<TopicPartitionOffset>();
         var committed = new List<TopicPartitionOffset>();
@@ -53,21 +57,53 @@ public static IReadOnlyCollection<TopicPartitionOffset> CommitState<TKey, TValue
 
         topic.OnReset(reset);
 
+        topic.CommitOffsetIfNeeded(activitySpan, newState);
+
+        return committed;
+    }
+
+    public static void CommitOffsetIfNeeded<TKey, TValue>(
+        this SubscriptionTopicWrapper<TKey, TValue> topic, 
+        ActivityWrapper activitySpan,
+        IReadOnlyCollection<TopicPartitionOffset> offsets)
+    {
         if (topic.Options.ExternalStateCommitToKafka)
         {
             try
             {
-                topic.CommitOffsets(activitySpan, committed);
+                List<TopicPartitionOffset> toCommit = new();
+
+                foreach (TopicPartitionOffset item in offsets)
+                {
+                    if (item.Offset.Value >= 0)
+                    {
+                        toCommit.Add(item);
+                    }
+                    else if (item.Offset == Offset.Beginning)
+                    {
+                        var w = topic.Consumer.GetWatermarkOffsets(item.TopicPartition);
+
+                        if (w.Low == Offset.Unset)
+                        {
+                            w = topic.Consumer.QueryWatermarkOffsets(item.TopicPartition, TimeSpan.FromSeconds(5));
+                        }
+
+                        if (w.Low.Value >= 0)
+                        {
+                            toCommit.Add(new TopicPartitionOffset(item.TopicPartition, w.Low));
+                        }
+                    }
+                }
+
+                topic.CommitOffsets(activitySpan, toCommit);
             }
             catch (KafkaException exception)
             {
-                topic.Logger.KafkaCommitFailed(exception, topic.Monitor.Name, committed);
+                topic.Logger.KafkaCommitFailed(exception, topic.Monitor.Name, offsets);
 
                 // ignore exception because external provider is a single point of truth for offsets.
                 // failed commit will trigger offset reset on next batch iteration
             }
         }
-
-        return committed;
     }
 }
diff --git a/src/Epam.Kafka.PubSub/Subscription/State/InternalKafkaState.cs b/src/Epam.Kafka.PubSub/Subscription/State/InternalKafkaState.cs
@@ -11,6 +11,7 @@ namespace Epam.Kafka.PubSub.Subscription.State;
 internal class InternalKafkaState : BatchState
 {
     protected override void AssignConsumer<TKey, TValue>(SubscriptionTopicWrapper<TKey, TValue> topic,
+        ActivityWrapper activitySpan,
         CancellationToken cancellationToken)
     {
         if (topic == null)

diff --git a/src/Epam.Kafka.PubSub/Subscription/Topics/SubscriptionTopicWrapper.cs b/src/Epam.Kafka.PubSub/Subscription/Topics/SubscriptionTopicWrapper.cs
@@ -346,10 +346,13 @@ private void ConfigureConsumerConfig(ConsumerConfig config)
 
     public void CommitOffsets(ActivityWrapper activitySpan, IReadOnlyCollection<TopicPartitionOffset> offsets)
     {
-        using (activitySpan.CreateSpan("commit_kafka"))
+        if (offsets.Count > 0)
         {
+            using var span = activitySpan.CreateSpan("commit_kafka");
             // commit to kafka also
             this.Consumer.Commit(offsets);
+
+            span.SetResult(offsets);
         }
     }
 

diff --git a/src/Epam.Kafka/Epam.Kafka.csproj b/src/Epam.Kafka/Epam.Kafka.csproj
@@ -12,6 +12,7 @@ Simplify Confluent.Kafka usage by leveraging best patterns and practices:
 	<ItemGroup>
 	  <PackageReference Include="Confluent.Kafka" Version="2.4.0" />
 	  <PackageReference Include="Confluent.SchemaRegistry" Version="2.4.0" />
+	  <PackageReference Include="System.Text.Json" Version="8.0.4" />
 	</ItemGroup>
 
 	<ItemGroup Condition="'$(TargetFramework)' == 'net6.0' or '$(TargetFramework)' == 'net462' or '$(TargetFramework)' == 'netstandard2.0'">
@@ -21,10 +22,6 @@ Simplify Confluent.Kafka usage by leveraging best patterns and practices:
 		<PackageReference Include="System.Diagnostics.DiagnosticSource" Version="6.0.0" />
 	</ItemGroup>
 
-	<ItemGroup Condition="'$(TargetFramework)' == 'net462' or '$(TargetFramework)' == 'netstandard2.0'">
-		<PackageReference Include="System.Text.Json" Version="6.0.0" />
-	</ItemGroup>
-
 	<ItemGroup Condition="'$(TargetFramework)' == 'net8.0'">
 		<PackageReference Include="Microsoft.Extensions.Configuration.Abstractions" Version="8.0.0" />
 		<PackageReference Include="Microsoft.Extensions.Logging.Abstractions" Version="8.0.0" />

diff --git a/src/Epam.Kafka/Internals/Metrics/ConsumerTopicsMetrics.cs b/src/Epam.Kafka/Internals/Metrics/ConsumerTopicsMetrics.cs
diff --git a/src/Epam.Kafka/Internals/Metrics/StatisticsMetrics.cs b/src/Epam.Kafka/Internals/Metrics/StatisticsMetrics.cs
@@ -50,20 +50,4 @@ public void OnNext(Statistics value)
     }
 
     protected abstract void Create();
-
-    protected IEnumerable<KeyValuePair<string, object?>> BuildTpTags(string topic, long partition)
-    {
-        return this.Tags.Concat(
-            new Dictionary<string, object?>
-            {
-                { "topic", topic },
-                { "partition", partition }
-            });
-    }
-
-    protected Measurement<long> CreateStatusMetric(string value)
-    {
-        return new Measurement<long>(this.Latest.EpochTimeSeconds,
-            this.Tags.Concat(Enumerable.Repeat(new KeyValuePair<string, object?>("state", value), 1)));
-    }
 }
diff --git a/src/Epam.Kafka/Internals/Observable/ObservableConsumer.cs b/src/Epam.Kafka/Internals/Observable/ObservableConsumer.cs
@@ -13,7 +13,6 @@ internal class ObservableConsumer<TKey, TValue> : ObservableClient, IConsumer<TK
 {
     private readonly IConsumer<TKey, TValue> _inner;
     private readonly Meter? _meter;
-    private readonly Meter? _topicsMeter;
 
     public ObservableConsumer(ConsumerBuilder<TKey, TValue> builder, bool metrics)
     {
@@ -36,9 +35,7 @@ public ObservableConsumer(ConsumerBuilder<TKey, TValue> builder, bool metrics)
             if (metrics)
             {
                 this._meter = new Meter(Statistics.MeterName);
-                this._topicsMeter = new Meter(Statistics.TopicsMeterName);
                 this.StatObservers.Add(new ConsumerMetrics(this._meter));
-                this.StatObservers.Add(new ConsumerTopicsMetrics(this._topicsMeter));
             }
         }
         catch (InvalidOperationException)
@@ -62,7 +59,6 @@ public override void Dispose()
             this.ClearObservers();
 
             this._meter?.Dispose();
-            this._topicsMeter?.Dispose();
         }
     }
 

diff --git a/src/Epam.Kafka/Stats/Statistics.cs b/src/Epam.Kafka/Stats/Statistics.cs
@@ -18,11 +18,6 @@ public class Statistics
     /// </summary>
     public const string MeterName = "Epam.Kafka.Statistics";
 
-    /// <summary>
-    /// Name of <see cref="Meter"/> used to expose topics statistics if <see cref="KafkaConfigExtensions.DotnetStatisticMetricsKey"/> enable it.
-    /// </summary>
-    public const string TopicsMeterName = "Epam.Kafka.Statistics.Topics";
-
     /// <summary>
     /// Create new instance of <see cref="Statistics"/> object from json representation.
     /// </summary>

diff --git a/tests/Epam.Kafka.PubSub.IntegrationTests/Subscription/CombinedState/ResetTests.cs b/tests/Epam.Kafka.PubSub.IntegrationTests/Subscription/CombinedState/ResetTests.cs
@@ -82,7 +82,7 @@ public async Task OneRunningOneStarting()
 
         // iteration 3
         observer.AssertStart();
-        observer.AssertAssign();
+        observer.AssertAssign(true);
         observer.AssertRead(5);
         observer.AssertProcess();
         observer.AssertCommitExternal();
@@ -139,7 +139,7 @@ public async Task AllPausedStartOne()
 
         // iteration 3
         observer.AssertStart();
-        observer.AssertAssign();
+        observer.AssertAssign(true);
         observer.AssertRead(5);
         observer.AssertProcess();
         observer.AssertCommitExternal();
@@ -207,7 +207,7 @@ public async Task AllPausedStopOne(bool onCommit)
         observer.AssertStop(SubscriptionBatchResult.Processed);
 
         // iteration 3
-        observer.AssertSubEmpty();
+        observer.AssertSubEmpty(!onCommit);
     }
 
     [Theory]
@@ -275,7 +275,7 @@ public async Task ResetToBeginning(bool onCommit)
 
         // iteration 3
         observer.AssertStart();
-        observer.AssertAssign();
+        observer.AssertAssign(!onCommit);
         observer.AssertRead(5);
         observer.AssertProcess();
         observer.AssertCommitExternal();

diff --git a/tests/Epam.Kafka.PubSub.IntegrationTests/Subscription/ExternalState/ResetTests.cs b/tests/Epam.Kafka.PubSub.IntegrationTests/Subscription/ExternalState/ResetTests.cs
@@ -80,7 +80,7 @@ public async Task OneRunningOneStarting()
 
         // iteration 2
         observer.AssertStart();
-        observer.AssertAssign();
+        observer.AssertAssign(true);
         observer.AssertRead(5);
         observer.AssertProcess();
         observer.AssertCommitExternal();
@@ -134,7 +134,7 @@ public async Task AllPausedStartOne()
 
         // iteration 2
         observer.AssertStart();
-        observer.AssertAssign();
+        observer.AssertAssign(true);
         observer.AssertRead(5);
         observer.AssertProcess();
         observer.AssertCommitExternal();
@@ -203,6 +203,6 @@ public async Task AllPausedStopOne(bool onCommit)
         observer.AssertStop(SubscriptionBatchResult.Processed);
 
         // iteration 2
-        observer.AssertSubEmpty();
+        observer.AssertSubEmpty(!onCommit);
     }
 }
diff --git a/tests/Epam.Kafka.PubSub.Tests/Helpers/AssertExtensions.cs b/tests/Epam.Kafka.PubSub.Tests/Helpers/AssertExtensions.cs
@@ -17,9 +17,15 @@ public static void AssertSubNotAssigned(this TestObserver observer)
         observer.AssertStop(SubscriptionBatchResult.NotAssigned);
     }
 
-    public static void AssertAssign(this TestObserver observer)
+    public static void AssertAssign(this TestObserver observer, bool offsetsCommit = false)
     {
         observer.AssertNextActivity("assign.Start");
+
+        if (offsetsCommit)
+        {
+            observer.AssertCommitKafka();
+        }
+
         observer.AssertNextActivity("assign.Stop");
     }
 
@@ -54,10 +60,10 @@ public static void AssertProcess(this TestObserver observer)
         observer.AssertNextActivity("process.Stop");
     }
 
-    public static void AssertSubEmpty(this TestObserver observer)
+    public static void AssertSubEmpty(this TestObserver observer, bool offsetsCommit = false)
     {
         observer.AssertStart();
-        observer.AssertAssign();
+        observer.AssertAssign(offsetsCommit);
         observer.AssertRead(0);
         observer.AssertStop(SubscriptionBatchResult.Empty);
     }

diff --git a/tests/Epam.Kafka.Tests/Epam.Kafka.approved.txt b/tests/Epam.Kafka.Tests/Epam.Kafka.approved.txt
@@ -141,7 +141,6 @@ namespace Epam.Kafka.Stats
     public class Statistics
     {
         public const string MeterName = "Epam.Kafka.Statistics";
-        public const string TopicsMeterName = "Epam.Kafka.Statistics.Topics";
         public Statistics() { }
         [System.Text.Json.Serialization.JsonPropertyName("age")]
         public long AgeMicroseconds { get; set; }