StableCompositeChangeToken.cs

#nullable enable
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Text;
using System.Threading;
using Microsoft.Extensions.Primitives;

namespace Microsoft.Extensions.Primitives
{
	/// <summary>
	/// An <see cref="IChangeToken"/> which represents one or more <see cref="IChangeToken"/> instances.
	/// Important! Use this class instead of <see cref="CompositeChangeToken"/>.
	/// This class behaves nearly same as <see cref="CancellationTokenSource.CreateLinkedTokenSource(CancellationToken[])"/>,
	/// while <see cref="CompositeChangeToken"/> may cause a deadlock if multiple source tokens changes at the same time
	/// because of mutual registration disposes (see https://github.com/dotnet/runtime/issues/43438).
	/// </summary>
	public class StableCompositeChangeToken : IChangeToken
	{
		private static readonly Action<object> _onChangeDelegate = OnChange;

		private static CancellationTokenSource CreateCancelledCts()
		{
			var result = new CancellationTokenSource();
			result.Cancel();
			return result;
		}

		private static readonly CancellationTokenSource _cancellationTokenSourceSentinel = CreateCancelledCts();

		private readonly object _callbackLock = new object();
		private CancellationTokenSource? _cancellationTokenSource;
		private List<IDisposable>? _disposables;


		/// <summary>
		/// Creates a new instance of <see cref="StableCompositeChangeToken"/>.
		/// </summary>
		/// <param name="changeTokens">The list of <see cref="IChangeToken"/> to compose.</param>
		public StableCompositeChangeToken(IReadOnlyList<IChangeToken> changeTokens)
		{
			ChangeTokens = changeTokens ?? throw new ArgumentNullException(nameof(changeTokens));

			// see https://github.com/dotnet/runtime/issues/40677
			ActiveChangeCallbacks = true;
			for (int i = 0; i < ChangeTokens.Count; i++)
			{
				if (!ChangeTokens[i].ActiveChangeCallbacks)
				{
					ActiveChangeCallbacks = false;
					break;
				}
			}
		}

		/// <summary>
		/// Returns the list of <see cref="IChangeToken"/> which compose the current <see cref="StableCompositeChangeToken"/>.
		/// </summary>
		public IReadOnlyList<IChangeToken> ChangeTokens { get; }

		/// <inheritdoc />
		public IDisposable RegisterChangeCallback(Action<object> callback, object state)
		{
			EnsureCallbacksInitialized();
			return _cancellationTokenSource!.Token.Register(callback, state);
		}

		/// <inheritdoc />
		public bool HasChanged
		{
			get
			{
				var cts = _cancellationTokenSource;
				if (cts != null && cts.IsCancellationRequested)
				{
					return true;
				}

				for (int i = 0; i < ChangeTokens.Count; i++)
				{
					if (ChangeTokens[i].HasChanged)
					{
						OnChange(this);
						return true;
					}
				}

				return false;
			}
		}

		/// <inheritdoc />
		public bool ActiveChangeCallbacks { get; }

		private void EnsureCallbacksInitialized()
		{
			if (_cancellationTokenSource != null)
			{
				return;
			}

			List<IDisposable>? disposables = null;

			try
			{
				// using lock to not to subscribe to ChangeTokens twice
				lock (_callbackLock)
				{
					if (_cancellationTokenSource != null)
					{
						return;
					}

					disposables = new List<IDisposable>(ChangeTokens.Count);
					for (int i = 0; i < ChangeTokens.Count; i++)
					{
						if (ChangeTokens[i].ActiveChangeCallbacks)
						{
							IDisposable disposable = ChangeTokens[i].RegisterChangeCallback(_onChangeDelegate, this);
							disposables.Add(disposable);
						}
					}

					Debug.Assert(_disposables == null);
					_disposables = disposables;
					var prevCts = Interlocked.CompareExchange(ref _cancellationTokenSource, new CancellationTokenSource(), null);

					if (prevCts == null)
					{
						disposables = null;
					}
					else
					{
						// OnChange has been already raised
						Debug.Assert(prevCts == _cancellationTokenSourceSentinel);
					}
				}

			}
			finally
			{
				if (disposables != null)
				{
					DisposeAll(disposables);
				}
			}
		}

		private void OnChange()
		{
			var prevCts = Interlocked.Exchange(ref _cancellationTokenSource, _cancellationTokenSourceSentinel);

			if (prevCts == _cancellationTokenSourceSentinel || prevCts == null)
			{
				return;
			}

			try
			{
				prevCts.Cancel();
			}
			finally
			{
				var disposables = _disposables;
				Debug.Assert(disposables != null);
				DisposeAll(disposables!);
			}
		}

		private void DisposeAll(List<IDisposable> disposables)
		{
			for (int i = 0; i < disposables.Count; i++)
			{
				disposables[i].Dispose();
			}
		}

		private static void OnChange(object state)
		{
			((StableCompositeChangeToken)state).OnChange();
		}
	}
}