Fix partitioned session window bug, where restoring from checkpoint triggers exception

Author: arunkm

Sources/Core/Microsoft.StreamProcessing/Operators/SessionWindow/PartitionedSessionWindowPipe.cs

@@ -253,11 +253,18 @@ protected override void UpdatePointers()
             var temp = new List<Tuple<TKey, long, TPartitionKey>>();
             while (this.lastDataTimeDictionary.Iterate(ref iter))
             {
+                var partitionKey = this.getPartitionKey(this.lastDataTimeDictionary.entries[iter].key);
+
                 if (this.stateDictionary.entries[iter].value.Any())
                 {
                     temp.Add(Tuple.Create(
                         this.lastDataTimeDictionary.entries[iter].key,
-                        Math.Min(this.lastDataTimeDictionary.entries[iter].value + this.sessionTimeout, this.windowEndTimeDictionary.entries[iter].value), this.getPartitionKey(this.lastDataTimeDictionary.entries[iter].key)));
+                        Math.Min(this.lastDataTimeDictionary.entries[iter].value + this.sessionTimeout, this.windowEndTimeDictionary.entries[iter].value),
+                        partitionKey));
+                }
+                else if (!this.orderedKeysDictionary.ContainsKey(partitionKey))
+                {
+                    this.orderedKeysDictionary.Add(partitionKey, new LinkedList<TKey>());
                 }
             }
             foreach (var item in temp.OrderBy(o => o.Item2))

Sources/Test/SimpleTesting/Partitioned/PartitionedStreamCheckpointTests.cs

@@ -0,0 +1,91 @@
+// *********************************************************************
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+// Licensed under the MIT License
+// *********************************************************************
+using System;
+using System.Collections.Generic;
+using System.IO;
+using System.Linq;
+using System.Reactive.Linq;
+using System.Reactive.Subjects;
+using Microsoft.StreamProcessing;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+namespace SimpleTesting
+{
+    [TestClass]
+    public class PartitionedStreamCheckpointTests : TestWithConfigSettingsWithoutMemoryLeakDetection
+    {
+        [TestMethod, TestCategory("Gated")]
+        public void CheckpointPartitionedSessionWindow()
+        {
+            Config.DataBatchSize = 1;
+
+            var data = new PartitionedStreamEvent<int, double>[]
+            {
+                PartitionedStreamEvent.CreatePoint(0, 5, 1.0),
+                PartitionedStreamEvent.CreatePunctuation<int, double>(0, 8),
+                PartitionedStreamEvent.CreatePunctuation<int, double>(0, 11),
+                PartitionedStreamEvent.CreatePunctuation<int, double>(0, 14),
+                PartitionedStreamEvent.CreatePunctuation<int, double>(0, 17),
+                PartitionedStreamEvent.CreatePunctuation<int, double>(0, 21),
+                PartitionedStreamEvent.CreatePoint(0, 24, 1.0),
+                PartitionedStreamEvent.CreatePunctuation<int, double>(0, 100),
+            };
+
+            var expected = new PartitionedStreamEvent<int, double>[]
+            {
+                PartitionedStreamEvent.CreateStart(0, 5, 1.0),
+                PartitionedStreamEvent.CreatePunctuation<int, double>(0, 8),
+                PartitionedStreamEvent.CreateEnd(0, 9, 5, 1.0),
+                PartitionedStreamEvent.CreatePunctuation<int, double>(0, 11),
+                PartitionedStreamEvent.CreatePunctuation<int, double>(0, 14),
+                PartitionedStreamEvent.CreatePunctuation<int, double>(0, 17),
+                PartitionedStreamEvent.CreatePunctuation<int, double>(0, 21),
+                PartitionedStreamEvent.CreateStart(0, 24, 1.0),
+                PartitionedStreamEvent.CreateEnd(0, 28, 24, 1.0),
+                PartitionedStreamEvent.CreatePunctuation<int, double>(0, 100),
+            };
+
+            // This index represents after series of punctions and just before new data arrives
+            const int checkpointIndex = 6;
+
+            var subject = new Subject<PartitionedStreamEvent<int, double>>();
+            var output = new List<PartitionedStreamEvent<int, double>>();
+            var process = CreateQueryContainerForPartitionedStream(subject, output);
+
+            for (int i = 0; i < data.Length; i++)
+            {
+                if (i == checkpointIndex)
+                {
+                    using (var ms = new MemoryStream())
+                    {
+                        process.Checkpoint(ms);
+                        ms.Seek(0, SeekOrigin.Begin);
+
+                        subject = new Subject<PartitionedStreamEvent<int, double>>();
+                        process = CreateQueryContainerForPartitionedStream(subject, output, ms);
+                    }
+                }
+
+                subject.OnNext(data[i]);
+            }
+
+            Assert.IsTrue(expected.SequenceEqual(output));
+        }
+
+        private Process CreateQueryContainerForPartitionedStream(
+            Subject<PartitionedStreamEvent<int, double>> subject,
+            List<PartitionedStreamEvent<int, double>> output,
+            Stream stream = null)
+        {
+            var qc = new QueryContainer();
+            var input = qc.RegisterInput(subject);
+            var streamableOutput = input.SessionTimeoutWindow(4, 5).Sum(o => o);
+            var egress = qc.RegisterOutput(streamableOutput).ForEachAsync(o => output.Add(o));
+            var process = qc.Restore(stream);
+
+            return process;
+        }
+    }
+}