Add an optional timeout to the periodic scheduler (#539)

This can cause extractor shutdown to hang forever if there's an issue with shutting down a task, so we should allow setting a timeout.

Should we set a default? It would technically be a breaking change, but hanging forever seems unlikely to be desirable behavior.

Author: einarmo

Cognite.Common/PeriodicScheduler.cs

@@ -254,7 +254,7 @@ public async Task WaitForTermination(string name)
         /// Same as calling ExitAndWaitForTermination on all tasks.
         /// </summary>
         /// <returns>Task which completes once all tasks are done</returns>
-        public async Task ExitAllAndWait()
+        public async Task ExitAllAndWait(int timeoutMs = 0)
         {
             var tasks = new List<PeriodicTask>(_tasks.Count);
             lock (_taskListMutex)
@@ -268,7 +268,19 @@ public async Task ExitAllAndWait()
             }
             try
             {
-                await Task.WhenAll(tasks.Select(tsk => tsk.Task)).ConfigureAwait(false);
+                var timeout = timeoutMs > 0 ? TimeSpan.FromMilliseconds(timeoutMs) : Timeout.InfiniteTimeSpan;
+                var tasksToWaitFor = tasks.Select(tsk => tsk.Task).ToList();
+                // Do not pass the cancellation token, since we are likely cancelled, and this method bypasses that.
+                var waitTask = CommonUtils.WaitAsync(_newTaskEvent, timeout, CancellationToken.None);
+                var res = await Task.WhenAny(Task.WhenAll(tasksToWaitFor), waitTask).ConfigureAwait(false);
+                if (res == waitTask)
+                {
+                    var notFinishedTasks = tasks.Where(t => !t.Task.IsCompleted).Select(t => t.Name).ToList();
+
+                    throw new TimeoutException(
+                        $"Timed out waiting for task scheduler to finish. The following tasks failed to finish within {timeout}: {string.Join(", ", notFinishedTasks)}"
+                    );
+                }
             }
             catch (TaskCanceledException) { }
             catch (AggregateException aex)

ExtractorUtils.Test/unit/SchedulerTest.cs

@@ -205,5 +205,14 @@ public async Task TestBlockingResourceIncCapacity()
 
             Assert.Equal(1, await task2);
         }
+
+        [Fact]
+        public async Task TestSchedulerTimeout()
+        {
+            using var source = new CancellationTokenSource();
+            using var scheduler = new PeriodicScheduler(source.Token);
+            scheduler.ScheduleTask("forever", (_t) => Task.Delay(Timeout.Infinite, source.Token));
+            await Assert.ThrowsAsync<TimeoutException>(() => scheduler.ExitAllAndWait(200));
+        }
     }
 }

version

@@ -1 +1 @@
-1.32.0
+1.33.0