Subscription cleanup

workers/unity/Packages/io.improbable.gdk.core/Subscriptions/RequiredSubscriptionsInjector.cs

@@ -1,7 +1,6 @@
 using System;
 using System.Collections.Generic;
 using Improbable.Gdk.Core;
-using UnityEngine;
 
 namespace Improbable.Gdk.Subscriptions
 {
@@ -33,7 +32,7 @@ public RequiredSubscriptionsInjector(object target, EntityId entityId, Subscript
                 return;
             }
 
-            subscriptions = new SubscriptionAggregate(subscriptionSystem, entityId, info.RequiredTypes);
+            subscriptions = subscriptionSystem.Subscribe(entityId, info.RequiredTypes);
             subscriptions.SetAvailabilityHandler(Handler.Pool.Rent(this));
         }
 
@@ -105,13 +104,13 @@ public void OnUnavailable()
                 injector.HandleSubscriptionsNoLongerSatisfied();
             }
 
-            public class Pool
+            public static class Pool
             {
                 private static readonly Stack<Handler> HandlerPool = new Stack<Handler>();
 
                 public static Handler Rent(RequiredSubscriptionsInjector injector)
                 {
-                    Handler handler = HandlerPool.Count == 0
+                    var handler = HandlerPool.Count == 0
                         ? new Handler()
                         : HandlerPool.Pop();
 

workers/unity/Packages/io.improbable.gdk.core/Subscriptions/SubscriptionAggregate.cs

@@ -1,6 +1,6 @@
 using System;
 using System.Collections.Generic;
-using Improbable.Gdk.Core;
+using UnityEngine.Assertions;
 
 namespace Improbable.Gdk.Subscriptions
 {
@@ -10,28 +10,27 @@ public class SubscriptionAggregate
         private readonly Dictionary<Type, int> typesToSubscriptionIndexes = new Dictionary<Type, int>();
 
         private ISubscriptionAvailabilityHandler availabilityHandler;
-
         private int subscriptionsSatisfied;
 
-        public SubscriptionAggregate(SubscriptionSystem subscriptionSystem, EntityId entityId,
-            params Type[] typesToSubscribeTo)
+        private bool IsSatisfied => subscriptionsSatisfied == subscriptions.Length;
+
+        public SubscriptionAggregate(IReadOnlyList<Type> types, ISubscription[] subscriptions)
         {
-            subscriptions = new ISubscription[typesToSubscribeTo.Length];
+            Assert.AreEqual(types.Count, subscriptions.Length);
 
-            for (int i = 0; i < typesToSubscribeTo.Length; ++i)
+            this.subscriptions = subscriptions;
+            for (var i = 0; i < types.Count; i++)
             {
-                var subscription = subscriptionSystem.Subscribe(entityId, typesToSubscribeTo[i]);
-                typesToSubscriptionIndexes.Add(typesToSubscribeTo[i], i);
-                subscriptions[i] = subscription;
-
-                subscription.SetAvailabilityHandler(Handler.Pool.Rent(this));
+                var type = types[i];
+                typesToSubscriptionIndexes.Add(type, i);
+                subscriptions[i].SetAvailabilityHandler(Handler.Pool.Rent(this));
             }
         }
 
         public void SetAvailabilityHandler(ISubscriptionAvailabilityHandler handler)
         {
             availabilityHandler = handler;
-            if (subscriptionsSatisfied == subscriptions.Length)
+            if (IsSatisfied)
             {
                 availabilityHandler?.OnAvailable();
             }
@@ -44,24 +43,21 @@ public ISubscriptionAvailabilityHandler GetAvailabilityHandler()
 
         public T GetValue<T>()
         {
-            if (subscriptionsSatisfied != subscriptions.Length)
-            {
-                throw new InvalidOperationException("Subscriptions not all satisfied");
-            }
-
-            if (!typesToSubscriptionIndexes.TryGetValue(typeof(T), out var subscriptionIndex))
-            {
-                throw new InvalidOperationException("No subscription with that type");
-            }
+            var subscriptionIndex = GetTypeIndex(typeof(T));
+            var subscription = (Subscription<T>) subscriptions[subscriptionIndex];
 
-            var sub = (Subscription<T>) subscriptions[subscriptionIndex];
-
-            return sub.Value;
+            return subscription.Value;
         }
 
         public object GetErasedValue(Type type)
         {
-            if (subscriptionsSatisfied != subscriptions.Length)
+            var subscriptionIndex = GetTypeIndex(type);
+            return subscriptions[subscriptionIndex].GetErasedValue();
+        }
+
+        private int GetTypeIndex(Type type)
+        {
+            if (!IsSatisfied)
             {
                 throw new InvalidOperationException("Subscriptions not all satisfied");
             }
@@ -71,7 +67,7 @@ public object GetErasedValue(Type type)
                 throw new InvalidOperationException("No subscription with that type");
             }
 
-            return subscriptions[subscriptionIndex].GetErasedValue();
+            return subscriptionIndex;
         }
 
         public void Cancel()
@@ -85,7 +81,7 @@ public void Cancel()
 
         private void HandleSubscriptionUnavailable()
         {
-            if (subscriptionsSatisfied == subscriptions.Length)
+            if (IsSatisfied)
             {
                 foreach (var sub in subscriptions)
                 {
@@ -101,7 +97,7 @@ private void HandleSubscriptionUnavailable()
         private void HandleSubscriptionAvailable()
         {
             ++subscriptionsSatisfied;
-            if (subscriptionsSatisfied == subscriptions.Length)
+            if (IsSatisfied)
             {
                 availabilityHandler?.OnAvailable();
             }
@@ -121,7 +117,7 @@ public void OnUnavailable()
                 aggregate?.HandleSubscriptionUnavailable();
             }
 
-            public class Pool
+            public static class Pool
             {
                 private static readonly Stack<Handler> HandlerPool = new Stack<Handler>();
 

workers/unity/Packages/io.improbable.gdk.core/Subscriptions/Systems/SubscriptionSystem.cs

@@ -1,6 +1,5 @@
 using System;
 using System.Collections.Generic;
-using System.Reflection;
 using Improbable.Gdk.Core;
 using Unity.Entities;
 
@@ -12,14 +11,38 @@ public class SubscriptionSystem : ComponentSystem
         private readonly Dictionary<Type, SubscriptionManagerBase> typeToSubscriptionManager =
             new Dictionary<Type, SubscriptionManagerBase>();
 
+        protected override void OnCreate()
+        {
+            base.OnCreate();
+            Enabled = false;
+
+            AutoRegisterManagers();
+        }
+
+        protected override void OnUpdate()
+        {
+        }
+
+        private void AutoRegisterManagers()
+        {
+            var types = ReflectionUtility.GetNonAbstractTypes(typeof(SubscriptionManagerBase),
+                typeof(AutoRegisterSubscriptionManagerAttribute));
+
+            foreach (var type in types)
+            {
+                var instance = (SubscriptionManagerBase) Activator.CreateInstance(type, World);
+                RegisterSubscriptionManager(instance.SubscriptionType, instance);
+            }
+        }
+
         public void RegisterSubscriptionManager(Type type, SubscriptionManagerBase manager)
         {
             if (typeToSubscriptionManager.ContainsKey(type))
             {
                 throw new InvalidOperationException("Already a manager registered");
             }
 
-            typeToSubscriptionManager[type] = manager;
+            typeToSubscriptionManager.Add(type, manager);
         }
 
         public Subscription<T> Subscribe<T>(EntityId entity)
@@ -42,28 +65,15 @@ public ISubscription Subscribe(EntityId entity, Type type)
             return manager.SubscribeTypeErased(entity);
         }
 
-        protected override void OnCreate()
-        {
-            base.OnCreate();
-            Enabled = false;
-
-            AutoRegisterManagers();
-        }
-
-        protected override void OnUpdate()
-        {
-        }
-
-        private void AutoRegisterManagers()
+        public SubscriptionAggregate Subscribe(EntityId entity, params Type[] types)
         {
-            var types = ReflectionUtility.GetNonAbstractTypes(typeof(SubscriptionManagerBase),
-                typeof(AutoRegisterSubscriptionManagerAttribute));
-
-            foreach (var type in types)
+            var subscriptions = new ISubscription[types.Length];
+            for (var i = 0; i < types.Length; i++)
             {
-                var instance = (SubscriptionManagerBase) Activator.CreateInstance(type, World);
-                RegisterSubscriptionManager(instance.SubscriptionType, instance);
+                subscriptions[i] = Subscribe(entity, types[i]);
             }
+
+            return new SubscriptionAggregate(types, subscriptions);
         }
     }
 }

workers/unity/Packages/io.improbable.gdk.core/Tests/Editmode/Subscriptions/AggregateSubscriptionTests.cs

@@ -14,7 +14,7 @@ public class AggregateSubscriptionTests : SubscriptionTestBase
         public void AggregateSubscription_should_be_available_if_all_constraints_satisfied()
         {
             var handler = new AvailabilityHandler();
-            var aggSub = new SubscriptionAggregate(SubscriptionSystem, entityId, typeof(EntityId), typeof(World));
+            var aggSub = SubscriptionSystem.Subscribe(entityId, typeof(EntityId), typeof(World));
             aggSub.SetAvailabilityHandler(handler);
 
             Assert.IsTrue(handler.IsAvailable);
@@ -26,7 +26,7 @@ public void AggregateSubscription_should_be_available_if_all_constraints_satisfi
         public void AggregateSubscription_should_not_be_available_if_not_all_constraints_satisfied()
         {
             var handler = new AvailabilityHandler();
-            var aggSub = new SubscriptionAggregate(SubscriptionSystem, entityId, typeof(Entity), typeof(World));
+            var aggSub = SubscriptionSystem.Subscribe(entityId, typeof(Entity), typeof(World));
             aggSub.SetAvailabilityHandler(handler);
 
             Assert.IsFalse(handler.IsAvailable);