Fix some issues with Ignitor

src/Components/Ignitor/src/CancellableOperation.cs

@@ -5,37 +5,50 @@
 using System.Threading;
 using System.Threading.Tasks;
 
+#nullable enable
 namespace Ignitor
 {
     internal class CancellableOperation<TResult>
     {
-        public CancellableOperation(TimeSpan? timeout)
+        public CancellableOperation(TimeSpan? timeout, CancellationToken cancellationToken)
         {
             Timeout = timeout;
 
             Completion = new TaskCompletionSource<TResult>(TaskContinuationOptions.RunContinuationsAsynchronously);
             Completion.Task.ContinueWith(
                 (task, state) =>
                 {
-                    var operation = (CancellableOperation<TResult>)state;
+                    var operation = (CancellableOperation<TResult>)state!;
                     operation.Dispose();
                 },
                 this,
                 TaskContinuationOptions.ExecuteSynchronously); // We need to execute synchronously to clean-up before anything else continues
 
+
+            Cancellation = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+
             if (Timeout != null && Timeout != System.Threading.Timeout.InfiniteTimeSpan && Timeout != TimeSpan.MaxValue)
             {
-                Cancellation = new CancellationTokenSource(Timeout.Value);
-                CancellationRegistration = Cancellation.Token.Register(
-                    (self) =>
-                    {
-                        var operation = (CancellableOperation<TResult>)self;
-                        operation.Completion.TrySetCanceled(operation.Cancellation.Token);
-                        operation.Cancellation.Dispose();
-                        operation.CancellationRegistration.Dispose();
-                    },
-                    this);
+                Cancellation.CancelAfter(Timeout.Value);
             }
+
+            CancellationRegistration = Cancellation.Token.Register(
+                (self) =>
+                {
+                    var operation = (CancellableOperation<TResult>)self!;
+
+                    if (cancellationToken.IsCancellationRequested)
+                    {
+                        // The operation was externally canceled before it timed out.
+                        Dispose();
+                        return;
+                    }
+
+                    operation.Completion.TrySetException(new TimeoutException($"The operation timed out after {Timeout}."));
+                    operation.Cancellation?.Dispose();
+                    operation.CancellationRegistration.Dispose();
+                },
+                this);
         }
 
         public TimeSpan? Timeout { get; }
@@ -62,3 +75,4 @@ private void Dispose()
         }
     }
 }
+#nullable restore

src/Components/Ignitor/src/ComponentState.cs

@@ -3,6 +3,7 @@
 
 namespace Ignitor
 {
+#nullable enable
     public class ComponentState
     {
         public ComponentState(int componentId)
@@ -11,6 +12,7 @@ public ComponentState(int componentId)
         }
 
         public int ComponentId { get; }
-        public IComponent Component { get; }
+        public IComponent? Component { get; }
     }
+#nullable restore
 }

src/Components/Ignitor/src/ContainerNode.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Generic;
 
+#nullable enable
 namespace Ignitor
 {
     public abstract class ContainerNode : Node
@@ -82,3 +83,4 @@ public void RemoveLogicalChild(int childIndex)
         }
     }
 }
+#nullable restore

src/Components/Ignitor/src/ElementHive.cs

@@ -3,7 +3,9 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
 
+#nullable enable
 namespace Ignitor
 {
     public class ElementHive
@@ -35,7 +37,7 @@ public void Update(RenderBatch batch)
             }
         }
 
-        public bool TryFindElementById(string id, out ElementNode element)
+        public bool TryFindElementById(string id, [NotNullWhen(true)] out ElementNode? element)
         {
             foreach (var kvp in Components)
             {
@@ -49,7 +51,7 @@ public bool TryFindElementById(string id, out ElementNode element)
             element = null;
             return false;
 
-            bool TryGetElementFromChildren(Node node, out ElementNode foundNode)
+            bool TryGetElementFromChildren(Node node, out ElementNode? foundNode)
             {
                 if (node is ElementNode elementNode &&
                     elementNode.Attributes.TryGetValue("id", out var elementId) &&
@@ -81,6 +83,7 @@ private void UpdateComponent(RenderBatch batch, int componentId, ArrayBuilderSeg
             {
                 component = new ComponentNode(componentId);
                 Components.Add(componentId, component);
+
             }
 
             ApplyEdits(batch, component, 0, edits);
@@ -199,7 +202,7 @@ private void ApplyEdits(RenderBatch batch, ContainerNode parent, int childIndex,
 
                     case RenderTreeEditType.StepOut:
                         {
-                            parent = parent.Parent;
+                            parent = parent.Parent ?? throw new InvalidOperationException($"Cannot step out of {parent}");
                             currentDepth--;
                             childIndexAtCurrentDepth = currentDepth == 0 ? childIndex : 0; // The childIndex is only ever nonzero at zero depth
                             break;
@@ -469,3 +472,4 @@ public PermutationListEntry(int from, int to)
         }
     }
 }
+#nullable restore

src/Components/Ignitor/src/ElementNode.cs

@@ -7,6 +7,7 @@
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.SignalR.Client;
 
+#nullable enable
 namespace Ignitor
 {
     public class ElementNode : ContainerNode
@@ -87,7 +88,7 @@ internal Task SelectAsync(HubConnection connection, string value)
             return DispatchEventCore(connection, Serialize(webEventDescriptor), Serialize(args));
         }
 
-        public Task ClickAsync(HubConnection connection)
+        internal Task ClickAsync(HubConnection connection)
         {
             if (!Events.TryGetValue("click", out var clickEventDescriptor))
             {
@@ -129,3 +130,4 @@ public ElementEventDescriptor(string eventName, ulong eventId)
         }
     }
 }
+#nullable restore

src/Components/Ignitor/src/Error.cs

@@ -1,10 +1,12 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
+#nullable enable
 namespace Ignitor
 {
     public class Error
     {
-        public string Stack { get; set; }
+        public string? Stack { get; set; }
     }
 }
+#nullable restore

src/Components/Ignitor/src/IComponent.cs

@@ -1,6 +1,5 @@
-using System;
-using System.Collections.Generic;
-using System.Text;
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 namespace Ignitor
 {

src/Components/Ignitor/src/Node.cs

@@ -1,10 +1,14 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
+#nullable enable
+
 namespace Ignitor
 {
     public abstract class Node
     {
-        public virtual ContainerNode Parent { get; set; }
+        public virtual ContainerNode? Parent { get; set; }
     }
 }
+
+#nullable restore

src/Components/Ignitor/src/NodeSerializer.cs

@@ -4,6 +4,8 @@
 using System;
 using System.IO;
 
+#nullable enable
+
 namespace Ignitor
 {
     internal static class NodeSerializer
@@ -96,7 +98,7 @@ private void SerializeElement(ElementNode elementNode)
                     if (attribute.Value != null)
                     {
                         Write("=\"");
-                        Write(attribute.Value.ToString());
+                        Write(attribute.Value.ToString()!);
                         Write("\"");
                     }
                 }
@@ -113,7 +115,7 @@ private void SerializeElement(ElementNode elementNode)
                         if (properties.Value != null)
                         {
                             Write("=\"");
-                            Write(properties.Value.ToString());
+                            Write(properties.Value.ToString()!);
                             Write("\"");
                         }
                     }
@@ -194,3 +196,4 @@ private void WriteIndent()
         }
     }
 }
+#nullable restore

src/Components/Ignitor/src/Operations.cs

@@ -3,6 +3,7 @@
 
 using System.Collections.Concurrent;
 
+#nullable enable
 namespace Ignitor
 {
     public sealed class Operations
@@ -18,3 +19,4 @@ public sealed class Operations
         public ConcurrentQueue<CapturedJSInteropCall> JSInteropCalls { get; } = new ConcurrentQueue<CapturedJSInteropCall>();
     }
 }
+#nullable restore

src/Components/Ignitor/src/RenderBatchReader.cs

@@ -4,6 +4,8 @@
 using System;
 using System.Text;
 
+#nullable enable
+
 namespace Ignitor
 {
     public static class RenderBatchReader
@@ -206,7 +208,7 @@ private static ArrayRange<ulong> ReadDisposedEventHandlerIds(ReadOnlySpan<byte>
             return new ArrayRange<ulong>(Array.Empty<ulong>(), 0);
         }
 
-        private static string ReadString(ReadOnlySpan<byte> data, string[] strings)
+        private static string? ReadString(ReadOnlySpan<byte> data, string[] strings)
         {
             var index = BitConverter.ToInt32(data.Slice(0, 4));
             return index >= 0 ? strings[index] : null;
@@ -279,3 +281,4 @@ public ReadOnlySpan<byte> GetStringTableIndexes(ReadOnlySpan<byte> data)
         }
     }
 }
+#nullable restore