Add source operator id to watermark to know which source sent it. (#304)

* Add source operator id to watermark to know which source sent it. * Fix code smell
koralium · Jan 14, 2024 · 9e82acf · 9e82acf
1 parent 9f7d7f3
commit 9e82acf
Show file tree

Hide file tree

Showing 6 changed files with 37 additions and 5 deletions.
diff --git a/src/FlowtideDotNet.Base/Engine/Internal/VertexHandler.cs b/src/FlowtideDotNet.Base/Engine/Internal/VertexHandler.cs
@@ -48,6 +48,8 @@ public VertexHandler(
 
         public ILoggerFactory LoggerFactory { get; }
 
+        public string OperatorId => operatorName;
+
         public Task RegisterTrigger(string name, TimeSpan? scheduledInterval = null)
         {
             return registerTrigger(operatorName, name, scheduledInterval);

diff --git a/src/FlowtideDotNet.Base/IVertexHandler.cs b/src/FlowtideDotNet.Base/IVertexHandler.cs
@@ -18,6 +18,8 @@ namespace FlowtideDotNet.Base
 {
     public interface IVertexHandler
     {
+        string OperatorId { get; }
+
         string StreamName { get; }
 
         IMeter Metrics { get; }

diff --git a/src/FlowtideDotNet.Base/Vertices/Egress/EgressVertex.cs b/src/FlowtideDotNet.Base/Vertices/Egress/EgressVertex.cs
@@ -15,6 +15,7 @@
 using FlowtideDotNet.Base.Vertices.Egress.Internal;
 using FlowtideDotNet.Storage.StateManager;
 using Microsoft.Extensions.Logging;
+using System.Collections.Concurrent;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Diagnostics.Metrics;
@@ -30,7 +31,7 @@ public abstract class EgressVertex<T, TState> : ITargetBlock<IStreamEvent>, IStr
         private IEgressImplementation? _targetBlock;
         private bool _isHealthy = true;
         private CancellationTokenSource? _cancellationTokenSource;
-        private float _lastLatency = 0;
+        private readonly ConcurrentDictionary<string, float> _lastLatency = new ConcurrentDictionary<string, float>();
 
         public string Name { get; private set; }
 
@@ -65,7 +66,16 @@ private void InitializeBlocks()
         private Task HandleWatermark(Watermark watermark)
         {
             var span = DateTimeOffset.UtcNow.Subtract(watermark.StartTime);
-            _lastLatency = (float)span.TotalMilliseconds;
+            var latency = (float)span.TotalMilliseconds;
+            if (watermark.SourceOperatorId != null)
+            {
+                _lastLatency[watermark.SourceOperatorId] = latency;
+            }
+            else
+            {
+                Logger.LogWarning("Recieved watermark without source operator id");
+            }
+
             return OnWatermark(watermark);
         }
 
@@ -166,7 +176,15 @@ public Task Initialize(string name, long restoreTime, long newTime, JsonElement?
             });
             Metrics.CreateObservableGauge("latency", () =>
             {
-                return _lastLatency;
+                List<Measurement<float>> output = new List<Measurement<float>>();
+                foreach(var kv in _lastLatency)
+                {
+                    output.Add(new Measurement<float>(kv.Value, new TagList
+                    {
+                        { "source", kv.Key }
+                    }));
+                }
+                return output;
             });
 
             return InitializeOrRestore(restoreTime, dState, vertexHandler.StateClient);

diff --git a/src/FlowtideDotNet.Base/Vertices/Ingress/IngressOutput.cs b/src/FlowtideDotNet.Base/Vertices/Ingress/IngressOutput.cs
@@ -78,6 +78,8 @@ internal Task SendLockingEvent(ILockingEvent lockingEvent)
 
         public Task SendWatermark(Watermark watermark)
         {
+            Debug.Assert(_ingressState._vertexHandler != null, nameof(_ingressState._vertexHandler));
+            watermark.SourceOperatorId = _ingressState._vertexHandler.OperatorId;
             if (_inLock)
             {
                 return _targetBlock.SendAsync(watermark, CancellationToken);

diff --git a/src/FlowtideDotNet.Base/Vertices/MultipleInput/MultipleInputVertex.cs b/src/FlowtideDotNet.Base/Vertices/MultipleInput/MultipleInputVertex.cs
@@ -212,7 +212,10 @@ private async IAsyncEnumerable<IStreamEvent> HandleWatermark(int targetId, Water
 
             if (_currentWatermark == null) 
             {
-                _currentWatermark = new Watermark(ImmutableDictionary<string, long>.Empty, watermark.StartTime);
+                _currentWatermark = new Watermark(ImmutableDictionary<string, long>.Empty, watermark.StartTime)
+                {
+                    SourceOperatorId = watermark.SourceOperatorId
+                };
             }
             _targetWatermarks[targetId] = watermark;
             var currentDict = _currentWatermark.Watermarks;
@@ -240,7 +243,10 @@ private async IAsyncEnumerable<IStreamEvent> HandleWatermark(int targetId, Water
                     currentDict = currentDict.SetItem(kv.Key, watermarkValue);
                 }
             }
-            var newWatermark = new Watermark(currentDict, watermark.StartTime);
+            var newWatermark = new Watermark(currentDict, watermark.StartTime)
+            {
+                SourceOperatorId = watermark.SourceOperatorId
+            };
 
             // only output watermark if there is a difference in the numbers
             if (!newWatermark.Equals(_currentWatermark))

diff --git a/src/FlowtideDotNet.Base/Watermark.cs b/src/FlowtideDotNet.Base/Watermark.cs
@@ -52,6 +52,8 @@ public Watermark(IImmutableDictionary<string, long> watermarks, DateTimeOffset s
 
         public DateTimeOffset StartTime { get; }
 
+        public string? SourceOperatorId { get; internal set; }
+
         public override bool Equals(object? obj)
         {
             return obj is Watermark watermark &&