Ensure TransformOnObservable preserves ChangeSet order

src/DynamicData.Tests/Cache/TransformOnObservableFixture.cs

@@ -1,7 +1,9 @@
 using System;
 using System.Linq;
 using System.Reactive;
+using System.Reactive.Concurrency;
 using System.Reactive.Linq;
+using System.Reactive.Subjects;
 using System.Threading.Tasks;
 using Bogus;
 using DynamicData.Kernel;
@@ -141,13 +143,27 @@ IObservable<string> CreateChildObs(Animal a, int id) =>
         results.IsCompleted.Should().Be(completeSource && completeChildren);
     }
 
+    [Fact]
+    public void ResultFailsIfChildFails()
+    {
+        // Arrange
+        var expectedError = new Exception("Expected");
+        var throwObservable = Observable.Throw<IChangeSet<Animal, int>>(expectedError);
+
+        // Act
+        using var results = _animalCache.Connect().TransformOnObservable(_ => throwObservable).AsAggregator();
+
+        // Assert
+        results.Error.Should().Be(expectedError);
+    }
+
     [Fact]
     public void ResultFailsIfSourceFails()
     {
         // Arrange
         var expectedError = new Exception("Expected");
         var throwObservable = Observable.Throw<IChangeSet<Animal, int>>(expectedError);
-        using var results = _animalCache.Connect().Concat(throwObservable).TransformOnObservable(animal => Observable.Return(animal)).AsAggregator();
+        using var results = _animalCache.Connect().Concat(throwObservable).TransformOnObservable(Observable.Return).AsAggregator();
 
         // Act
         _animalCache.Dispose();
@@ -156,6 +172,39 @@ public void ResultFailsIfSourceFails()
         results.Error.Should().Be(expectedError);
     }
 
+    [Theory]
+    [InlineData(true)]
+    [InlineData(false)]
+    public void OrderOfChangesIsPreserved(bool removeFirst)
+    {
+        // Arrange
+        using var results = _animalCache.Connect().TransformOnObservable(Observable.Return).AsAggregator();
+        (var firstReason, var nextReason, var expectedChanges) = removeFirst 
+            ? (ChangeReason.Remove, ChangeReason.Add, InitialCount * 2)
+            : (ChangeReason.Add, ChangeReason.Remove, InitialCount * 3);
+
+        // Act
+        _animalCache.Edit(updater =>
+        {
+            if (removeFirst)
+            {
+                updater.Clear();
+                updater.AddOrUpdate(_animalFaker.Generate(InitialCount));
+            }
+            else
+            {
+                updater.AddOrUpdate(_animalFaker.Generate(InitialCount));
+                updater.Clear();
+            }
+        });
+
+        // Assert
+        results.Messages.Count.Should().Be(2);
+        results.Messages[1].Count.Should().Be(expectedChanges);
+        results.Messages[1].Take(InitialCount).All(change => change.Reason == firstReason).Should().BeTrue();
+        results.Messages[1].Skip(InitialCount).All(change => change.Reason == nextReason).Should().BeTrue();
+    }
+
     public void Dispose()
     {
         _animalCache.Dispose();

src/DynamicData/Cache/Internal/TransformOnObservable.cs

@@ -2,63 +2,146 @@
 // Roland Pheasant licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for full license information.
 
+using System.Diagnostics;
 using System.Reactive.Disposables;
 using System.Reactive.Linq;
 using DynamicData.Internal;
 
 namespace DynamicData.Cache.Internal;
 
-internal sealed class TransformOnObservable<TSource, TKey, TDestination>(IObservable<IChangeSet<TSource, TKey>> source, Func<TSource, TKey, IObservable<TDestination>> transform)
+internal sealed class TransformOnObservable<TSource, TKey, TDestination>(IObservable<IChangeSet<TSource, TKey>> source, Func<TSource, TKey, IObservable<TDestination>> transform, bool transformOnRefresh = false)
     where TSource : notnull
     where TKey : notnull
     where TDestination : notnull
 {
-    public IObservable<IChangeSet<TDestination, TKey>> Run() => Observable.Create<IChangeSet<TDestination, TKey>>(observer =>
+    public IObservable<IChangeSet<TDestination, TKey>> Run() =>
+        Observable.Create<IChangeSet<TDestination, TKey>>(observer => new Subscription(source, transform, observer, transformOnRefresh));
+
+    // Maintains state for a single subscription
+    private sealed class Subscription : IDisposable
     {
-        var cache = new ChangeAwareCache<TDestination, TKey>();
-        var locker = InternalEx.NewLock();
-        var parentUpdate = false;
+#if NET9_0_OR_GREATER
+        private readonly Lock _synchronize = new();
+#else
+        private readonly object _synchronize = new();
+#endif
+        private readonly ChangeAwareCache<TDestination, TKey> _cache = new();
+        private readonly KeyedDisposable<TKey> _transformSubscriptions = new();
+        private readonly Func<TSource, TKey, IObservable<TDestination>> _transform;
+        private readonly IDisposable _sourceSubscription;
+        private readonly IObserver<IChangeSet<TDestination, TKey>> _observer;
+        private readonly bool _transformOnRefresh;
+        private int _subscriptionCounter = 1;
+        private int _updateCounter;
+
+        public Subscription(IObservable<IChangeSet<TSource, TKey>> source, Func<TSource, TKey, IObservable<TDestination>> transform, IObserver<IChangeSet<TDestination, TKey>> observer, bool transformOnRefresh)
+        {
+            _observer = observer;
+            _transform = transform;
+            _transformOnRefresh = transformOnRefresh;
+            _sourceSubscription = source
+                .Do(_ => IncrementUpdates())
+                .Synchronize(_synchronize)
+                .SubscribeSafe(ProcessSourceChangeSet, observer.OnError, CheckCompleted);
+        }
+
+        public void Dispose()
+        {
+            lock (_synchronize)
+            {
+                _sourceSubscription.Dispose();
+                _transformSubscriptions.Dispose();
+            }
+        }
 
-        // Helper to emit any pending changes when appropriate
-        void EmitChanges(bool fromParent)
+        private void ProcessSourceChangeSet(IChangeSet<TSource, TKey> changes)
         {
-            if (fromParent || !parentUpdate)
+            // Process all the changes at once to preserve the changeset order
+            foreach (var change in changes.ToConcreteType())
             {
-                var changes = cache!.CaptureChanges();
+                switch (change.Reason)
+                {
+                    // Shutdown existing sub (if any) and create a new one that
+                    // Will update the cache and emit the changes
+                    case ChangeReason.Add or ChangeReason.Update:
+                        CreateTransformSubscription(change.Current, change.Key);
+                        break;
+
+                    // Shutdown the existing subscription and remove from the cache
+                    case ChangeReason.Remove:
+                        _transformSubscriptions.Remove(change.Key);
+                        _cache.Remove(change.Key);
+                        break;
+
+                    case ChangeReason.Refresh:
+                        if (_transformOnRefresh)
+                        {
+                            CreateTransformSubscription(change.Current, change.Key);
+                        }
+                        else
+                        {
+                            // Let the downstream decide what this means
+                            _cache.Refresh(change.Key);
+                        }
+                        break;
+                }
+            }
+
+            // Emit any pending changes
+            EmitChanges();
+        }
+
+        private void IncrementUpdates() => Interlocked.Increment(ref _updateCounter);
+
+        private void EmitChanges()
+        {
+            if (Interlocked.Decrement(ref _updateCounter) == 0)
+            {
+                var changes = _cache.CaptureChanges();
                 if (changes.Count > 0)
                 {
-                    observer.OnNext(changes);
+                    _observer.OnNext(changes);
                 }
+            }
 
-                parentUpdate = false;
+            Debug.Assert(_updateCounter >= 0, "Should never be negative");
+        }
+
+        private void CheckCompleted()
+        {
+            if (Interlocked.Decrement(ref _subscriptionCounter) == 0)
+            {
+                _observer.OnCompleted();
             }
+
+            Debug.Assert(_subscriptionCounter >= 0, "Should never be negative");
         }
 
         // Create the sub-observable that takes the result of the transformation,
         // filters out unchanged values, and then updates the cache
-        IObservable<TDestination> CreateSubObservable(TSource obj, TKey key) =>
-            transform(obj, key)
+        private void CreateTransformSubscription(TSource obj, TKey key)
+        {
+            // Add a new subscription.  Do first so cleanup of existing subs doesn't trigger OnCompleted.
+            Interlocked.Increment(ref _subscriptionCounter);
+
+            // Create a container for the Disposable and add to the KeyedDisposable
+            var disposableContainer = _transformSubscriptions.Add(key, new SingleAssignmentDisposable());
+
+            // Create the transformation observable for the source item, filter unchanged, and update the cache
+            // Will Dispose immediately if OnCompleted fires upon subscription because OnCompleted disposes the container
+            // Remove the TransformSubscription if it completes because its not needed anymore
+            disposableContainer.Disposable = _transform(obj, key)
                 .DistinctUntilChanged()
-                .Synchronize(locker!)
-                .Do(val => cache!.AddOrUpdate(val, key));
-
-        // Flag a parent update is happening once inside the lock
-        var shared = source
-            .Synchronize(locker!)
-            .Do(_ => parentUpdate = true)
-            .Publish();
-
-        // MergeMany automatically handles Add/Update/Remove and OnCompleted/OnError correctly
-        var subMerged = shared
-            .MergeMany(CreateSubObservable)
-            .SubscribeSafe(_ => EmitChanges(fromParent: false), observer.OnError, observer.OnCompleted);
-
-        // Subscribe to the shared Observable to handle Remove events.  MergeMany will unsubscribe from the sub-observable,
-        // but the corresponding key value needs to be removed from the Cache so the remove is observed downstream.
-        var subRemove = shared
-            .OnItemRemoved((_, key) => cache!.Remove(key), invokeOnUnsubscribe: false)
-            .SubscribeSafe(_ => EmitChanges(fromParent: true), observer.OnError);
-
-        return new CompositeDisposable(shared.Connect(), subMerged, subRemove);
-    });
+                .Do(_ => IncrementUpdates())
+                .Synchronize(_synchronize)
+                .Finally(CheckCompleted)
+                .SubscribeSafe(val => TransformOnNext(val, key), _observer.OnError, () => _transformSubscriptions.Remove(key));
+        }
+
+        private void TransformOnNext(TDestination latestValue, TKey key)
+        {
+            _cache.AddOrUpdate(latestValue, key);
+            EmitChanges();
+        }
+    }
 }

src/DynamicData/Internal/KeyedDisposable.cs

@@ -0,0 +1,80 @@
+// Copyright (c) 2011-2023 Roland Pheasant. All rights reserved.
+// Roland Pheasant licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for full license information.
+
+namespace DynamicData.Internal;
+
+/// <summary>
+/// Manages Disposables by Key:
+/// 1) Adding a disposable with the same key will dispose/replace the previous one.
+/// 2) Adding when the container is Disposed will Dispose it immediately.
+/// </summary>
+/// <typeparam name="TKey">Type to use for the Key.</typeparam>
+internal sealed class KeyedDisposable<TKey> : IDisposable
+    where TKey : notnull
+{
+    private readonly Dictionary<TKey, IDisposable> _disposables = [];
+    private bool _disposedValue;
+
+    public int Count => _disposables.Count;
+
+    public IEnumerable<TKey> Keys => _disposables.Keys;
+
+    public bool ContainsKey(TKey key) => _disposables.ContainsKey(key);
+
+    public bool IsDisposed => _disposedValue;
+
+    public TDisposable Add<TDisposable>(TKey key, TDisposable disposable)
+        where TDisposable : IDisposable
+    {
+        disposable.ThrowArgumentNullExceptionIfNull(nameof(disposable));
+
+        if (!_disposedValue)
+        {
+            Remove(key);
+            _disposables.Add(key, disposable);
+        }
+        else
+        {
+            disposable.Dispose();
+        }
+
+        return disposable;
+    }
+
+    public void Remove(TKey key)
+    {
+#if NET6_0_OR_GREATER
+        if (_disposables.Remove(key, out var disposable))
+        {
+            disposable.Dispose();
+        }
+#else
+        if (_disposables.TryGetValue(key, out var disposable))
+        {
+            disposable.Dispose();
+            _disposables.Remove(key);
+        }
+#endif
+    }
+
+    public void Dispose()
+    {
+        // Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
+        Dispose(disposing: true);
+        GC.SuppressFinalize(this);
+    }
+
+    private void Dispose(bool disposing)
+    {
+        if (!_disposedValue)
+        {
+            _disposedValue = true;
+            if (disposing)
+            {
+                _disposables.Values.ForEach(d => d.Dispose());
+                _disposables.Clear();
+            }
+        }
+    }
+}