dispose wait set after invoking IContext.OnShutdown

This prevents violating the documentation.

src/ros2cs/ros2cs_core/Context.cs

@@ -60,6 +60,14 @@ public sealed class Context : IContext
         /// </remarks>
         private HashSet<GuardCondition> GuardConditions = new HashSet<GuardCondition>();
 
+        /// <summary>
+        /// Collection of wait sets active in this context;
+        /// </summary>
+        /// <remarks>
+        /// Also used for synchronisation when creating / removing guard conditions.
+        /// </remarks>
+        private HashSet<WaitSet> WaitSets = new HashSet<WaitSet>();
+
         /// <summary>
         /// Get the current RMW implementation.
         /// </summary>
@@ -171,6 +179,40 @@ internal bool RemoveGuardCondition(GuardCondition guardCondition)
             }
         }
 
+        /// <summary>
+        /// Create a wait set.
+        /// </summary>
+        /// <remarks>
+        /// This method is thread safe.
+        /// </remarks>
+        /// <returns> A new wait set instance. </returns>
+        internal WaitSet CreateWaitSet()
+        {
+            lock (this.WaitSets)
+            {
+                WaitSet waitSet = new WaitSet(this);
+                this.WaitSets.Add(waitSet);
+                return waitSet;
+            }
+        }
+
+        /// <summary>
+        /// Remove a wait set.
+        /// </summary>
+        /// <remarks>
+        /// This method is intended to be used by <see cref="WaitSet.Dispose"/> and does not dispose the wait set.
+        /// Furthermore, it is thread safe.
+        /// </remarks>
+        /// <param name="waitSet"> Wait set to remove. </param>
+        /// <returns> If the wait set existed in this context and has been removed. </returns>
+        internal bool RemoveWaitSet(WaitSet waitSet)
+        {
+            lock (this.WaitSets)
+            {
+                return this.WaitSets.Remove(waitSet);
+            }
+        }
+
         /// <remarks>
         /// This method is not thread safe.
         /// Do not call while the context or any entities
@@ -197,7 +239,7 @@ private void Dispose(bool disposing)
             {
                 Utils.CheckReturnEnum(ret);
             }
-            // only continue if ROSNodes and GuardConditions has not been finalized
+            // only continue if the collections of the active primitives have not been finalized
             if (disposing)
             {
                 this.OnShutdown?.Invoke();
@@ -211,7 +253,12 @@ private void Dispose(bool disposing)
                     guardCondition.DisposeFromContext();
                 }
                 this.GuardConditions.Clear();
-                // only safe when all nodes are gone, not calling Dispose() will leak the Handle
+                foreach (var waitSet in this.WaitSets)
+                {
+                    waitSet.DisposeFromContext();
+                }
+                this.WaitSets.Clear();
+                // only safe when all primitives are gone, not calling Dispose() will leak the Handle
                 Utils.CheckReturnEnum(NativeRcl.rcl_context_fini(this.Handle));
                 this.FreeHandles();
             }

src/ros2cs/ros2cs_core/WaitSet.cs

@@ -16,6 +16,7 @@
 using System;
 using System.Collections;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Linq;
 
 namespace ROS2
@@ -49,7 +50,12 @@ internal sealed class WaitSet : IReadOnlyCollection<IWaitable>, IExtendedDisposa
         /// <summary>
         /// Context associated with this wait set.
         /// </summary>
-        public IContext Context { get; private set; }
+        public IContext Context
+        {
+            get => this.ROSContext;
+        }
+
+        private Context ROSContext;
 
         /// <inheritdoc />
         public bool IsDisposed
@@ -99,7 +105,7 @@ public int Count
         /// <param name="context">Associated context</param>
         internal WaitSet(Context context)
         {
-            this.Context = context;
+            this.ROSContext = context;
             this.Handle = NativeRclInterface.rclcs_get_zero_initialized_wait_set();
             int ret = NativeRcl.rcl_wait_set_init(
                 this.Handle,
@@ -117,7 +123,6 @@ internal WaitSet(Context context)
                 this.FreeHandles();
                 Utils.CheckReturnEnum(ret);
             }
-            context.OnShutdown += this.Dispose;
         }
 
         /// <inheritdoc />
@@ -314,16 +319,37 @@ private void Dispose(bool disposing)
                 return;
             }
 
+            if (disposing)
+            {
+                bool success = this.ROSContext.RemoveWaitSet(this);
+                Debug.Assert(success, "failed to remove wait set");
+                this.ClearCollections();
+            }
+
             Utils.CheckReturnEnum(NativeRcl.rcl_wait_set_fini(this.Handle));
             this.FreeHandles();
+        }
 
-            if (disposing)
+        /// <summary> Dispose without modifying the context. </summary>
+        internal void DisposeFromContext()
+        {
+            if (this.Handle == IntPtr.Zero)
             {
-                this.Context.OnShutdown -= this.Dispose;
-                this.Subscriptions.Clear();
-                this.Clients.Clear();
-                this.Services.Clear();
+                return;
             }
+
+            this.ClearCollections();
+
+            Utils.CheckReturnEnum(NativeRcl.rcl_wait_set_fini(this.Handle));
+            this.FreeHandles();
+        }
+
+        private void ClearCollections()
+        {
+            this.Subscriptions.Clear();
+            this.Clients.Clear();
+            this.Services.Clear();
+            this.GuardConditions.Clear();
         }
 
         private void FreeHandles()

src/ros2cs/ros2cs_core/executors/ManualExecutor.cs

@@ -114,7 +114,7 @@ public bool IsReadOnly
         /// <param name="context"> Context to associate with. </param>
         /// <exception cref="ObjectDisposedException"> If <paramref name="context"/> is disposed. </exception>
         public ManualExecutor(Context context) : this(
-            new WaitSet(context),
+            context.CreateWaitSet(),
             context.CreateGuardCondition(() => { })
         )
         { }

src/ros2cs/ros2cs_tests/src/WaitSetTest.cs

@@ -36,7 +36,7 @@ public class WaitSetTest
         public void SetUp()
         {
             this.Context = new Context();
-            this.WaitSet = new WaitSet(this.Context);
+            this.WaitSet = this.Context.CreateWaitSet();
         }
 
         [TearDown]
@@ -65,6 +65,14 @@ public void DoubleDisposeWaitSet()
             Assert.That(this.WaitSet.IsDisposed, Is.True);
         }
 
+        [Test]
+        public void OnShutdownDisposal()
+        {
+            this.Context.OnShutdown += () => Assert.That(this.WaitSet.IsDisposed, Is.False);
+
+            this.Context.Dispose();
+        }
+
         [Test]
         public void TestSubscriptionCollection()
         {