[API] Allow to multihtread register and close. (#26)

mandel-macaque · web-flow · commit ba7e8deb010b · 2024-08-19T15:23:13.000-04:00
diff --git a/Marille.Tests/CancellationTests.cs b/Marille.Tests/CancellationTests.cs
@@ -3,11 +3,13 @@ namespace Marille.Tests;
 public class CancellationTests {
 
 	Hub _hub;
+	SemaphoreSlim _semaphoreSlim;
 	TopicConfiguration configuration;
 
 	public CancellationTests ()
 	{
-		_hub = new ();
+		_semaphoreSlim = new (1);
+		_hub = new (_semaphoreSlim);
 		configuration = new();
 	}
 	[Fact]
@@ -100,4 +102,54 @@ public async Task CloseAllWorkersNoEvents ()
 		Assert.Equal (0, worker1.ConsumedCount);
 		Assert.Equal (0, worker2.ConsumedCount);
 	}
+
+	[Fact]
+	public async Task MultithreadedClose ()
+	{
+		var threadCount = 100;
+		var results = new List<Task<bool>> (100);
+
+		Random random = new Random ();
+		
+		// create the topic and then try to close if from several threads ensuring that only one of them
+		// closes the channel.
+		configuration.Mode = ChannelDeliveryMode.AtLeastOnce;
+		var topic = nameof (MultithreadedClose);
+		await _hub.CreateAsync<WorkQueuesEvent> (topic, configuration);
+		
+		// block the closing until we have created all the needed threads
+		await _semaphoreSlim.WaitAsync ();
+		
+		for (var index = 0; index < threadCount; index++) {
+			var tcs = new TaskCompletionSource<bool> ();
+			results.Add (tcs.Task);
+			// try to register from diff threads and ensure there are no unexpected issues
+			// this means that we DO NOT have two true values
+			// DO NOT AWAIT THE TASKS OR ELSE YOU WILL DEADLOCK
+			Task.Run (async () => {
+				// random sleep to ensure that the other thread is also trying to create
+				var sleep = random.Next (1000);
+				await Task.Delay (TimeSpan.FromMilliseconds (sleep));
+				var closed = await _hub.CloseAsync <WorkQueuesEvent> (topic);
+				tcs.TrySetResult (closed);
+			});
+		}
+		
+		_semaphoreSlim.Release ();
+		var closed = await Task.WhenAll (results);
+		bool? positive = null;
+		var finalResult = true;
+		// ensure that we have a true and a false, that means that an && should be false
+		for (var index = 0; index < threadCount; index++) {
+			finalResult &= closed[index];
+			if (closed[index] && positive is null) {
+				positive = true;
+				continue;
+			}
+			if (closed[index] && positive is true) {
+				Assert.Fail ("More than one close happened.");
+			}
+		}
+		Assert.False (finalResult);
+	}
 }
diff --git a/Marille.Tests/RegistrationTests.cs b/Marille.Tests/RegistrationTests.cs
@@ -3,10 +3,12 @@ namespace Marille.Tests;
 public class RegistrationTests {
 
 	Hub _hub;
+	SemaphoreSlim _semaphoreSlim;
 
 	public RegistrationTests ()
 	{
-		_hub = new ();
+		_semaphoreSlim = new(1);
+		_hub = new (_semaphoreSlim);
 	} 
 
 	[Fact]
@@ -70,4 +72,50 @@ public async Task MultipleOneToOneRegistrationWithLambda ()
 		Assert.True (await _hub.RegisterAsync (topic, worker1));
 		Assert.False(await _hub.RegisterAsync (topic, action));
 	}
+
+	[Fact]
+	public async Task MutithreadCreate ()
+	{
+		var threadCount = 100;
+		var results = new List<Task<bool>> (100);
+
+		Random random = new Random ();
+		var topic = nameof (MutithreadCreate );
+
+		TopicConfiguration configuration = new() { Mode = ChannelDeliveryMode.AtMostOnceAsync };
+
+		await _semaphoreSlim.WaitAsync ();
+		for (var index = 0; index < threadCount; index++) {
+			var tcs = new TaskCompletionSource<bool> ();
+			results.Add (tcs.Task);
+			// try to register from diff threads and ensure there are no unexpected issues
+			// this means that we DO NOT have two true values
+			// DO NOT AWAIT THE TASKS OR ELSE YOU WILL DEADLOCK
+			Task.Run (async () => {
+				// random sleep to ensure that the other thread is also trying to create
+				var sleep = random.Next (1000);
+				await Task.Delay (TimeSpan.FromMilliseconds (sleep));
+				var created = await _hub.CreateAsync<WorkQueuesEvent> (topic, configuration);
+				tcs.TrySetResult (created);
+			});
+		}
+		
+		// release the semaphore so that we can move on
+		_semaphoreSlim.Release ();
+		var added = await Task.WhenAll (results);
+		bool? positive = null;
+		var finalResult = true;
+		// ensure that we have a true and a false, that means that an && should be false
+		for (var index = 0; index < threadCount; index++) {
+			finalResult &= added[index];
+			if (added[index] && positive is null) {
+				positive = true;
+				continue;
+			}
+			if (added[index] && positive is true) {
+				Assert.Fail ("More than one addition happened.");
+			}
+		}
+		Assert.False (finalResult);
+	}
 }
diff --git a/Marille/Hub.cs b/Marille/Hub.cs
@@ -7,11 +7,17 @@ namespace Marille;
 /// 
 /// </summary>
 public class Hub : IHub {
+	readonly SemaphoreSlim semaphoreSlim = new(1);
 	readonly Dictionary<string, Topic> topics = new();
 	readonly Dictionary<(string Topic, Type Type), (CancellationTokenSource CancellationToken, Task? ConsumeTask)> cancellationTokenSources = new();
 	readonly Dictionary<(string Topic, Type type), List<object>> workers = new();
-	
 	public Channel<WorkerError> WorkersExceptions { get; } = Channel.CreateUnbounded<WorkerError> ();
+
+	public Hub () : this (new SemaphoreSlim (1))  { }
+	internal Hub (SemaphoreSlim semaphore)
+	{
+		semaphoreSlim = semaphore;
+	} 
 	
 	void DeliverAtLeastOnce<T> (Channel<Message<T>> channel, IWorker<T> [] workersArray, Message<T> item, TimeSpan? timeout)
 		where T : struct
@@ -125,24 +131,27 @@ void StopConsuming<T> (string topicName) where T : struct
 		if (!cancellationTokenSources.TryGetValue ((topicName, type), out var consumingInfo))
 			return;
 		
-		if (!TryGetChannel<T> (topicName, out var ch))
+		if (!TryGetChannel<T> (topicName, out var topic, out _))
 			return;
 
 		// complete the channels, this wont throw an cancellation exception, it will stop the channels from writing
 		// and the consuming task will finish when it is done with the current message, therefore we can
 		// use that to know when we are done
-		ch.Channel.Writer.Complete ();
+		topic.CloseChannel<T> ();
 		if (consumingInfo.ConsumeTask is not null)
 			await consumingInfo.ConsumeTask;
 
 		// clean behind us
+		topic.RemoveChannel<T> ();
+		workers.Remove ((topicName, type));
 		cancellationTokenSources.Remove ((topicName, type));
 	}
 
-	bool TryGetChannel<T> (string topicName, [NotNullWhen(true)] out TopicInfo<T>? ch) where T : struct
+	bool TryGetChannel<T> (string topicName, [NotNullWhen(true)] out Topic? topic, [NotNullWhen(true)] out TopicInfo<T>? ch) where T : struct
 	{
+		topic = null;
 		ch = null;
-		if (!topics.TryGetValue (topicName, out Topic? topic)) {
+		if (!topics.TryGetValue (topicName, out topic)) {
 			return false;
 		}
 
@@ -173,21 +182,27 @@ public async Task<bool> CreateAsync<T> (string topicName, TopicConfiguration con
 
 		// the topic might already have the channel, in that case, do nothing
 		Type type = typeof (T);
-		if (!topics.TryGetValue (topicName, out Topic? topic)) {
-			topic = new(topicName);
-			topics [topicName] = topic;
-		}
+		await semaphoreSlim.WaitAsync ();
+		try {
+			if (!topics.TryGetValue (topicName, out Topic? topic)) {
+				topic = new(topicName);
+				topics [topicName] = topic;
+			}
 
-		if (!workers.ContainsKey ((topicName, type))) {
-			workers [(topicName, type)] = new(initialWorkers);
-		}
+			if (!workers.ContainsKey ((topicName, type))) {
+				workers [(topicName, type)] = new(initialWorkers);
+			}
 
-		if (topic.TryGetChannel<T> (out _)) {
-			return false;
+			if (topic.TryGetChannel<T> (out _)) {
+				return false;
+			}
+
+			var ch = topic.CreateChannel<T> (configuration);
+			await StartConsuming (topicName, configuration, ch);
+			return true;
+		} finally {
+			semaphoreSlim.Release ( );
 		}
-		var ch = topic.CreateChannel<T> (configuration);
-		await StartConsuming (topicName, configuration, ch);
-		return true;
 	}
 
 	/// <summary>
@@ -258,23 +273,28 @@ public Task<bool> CreateAsync<T> (string topicName, TopicConfiguration configura
 	/// <remarks>Workers can be added to channels that are already being processed. The Hub will pause the consumtion
 	/// of the messages while it adds the worker and will resume the processing after. Producer can be sending
 	/// messages while this operation takes place because messages will be buffered by the channel.</remarks>
-	public Task<bool> RegisterAsync<T> (string topicName, params IWorker<T>[] newWorkers) where T : struct
+	public async Task<bool> RegisterAsync<T> (string topicName, params IWorker<T>[] newWorkers) where T : struct
 	{
 		var type = typeof (T);
-		// we only allow the client to register to an existing topic
-		// in this API we will not create it, there are other APIs for that
-		if (!TryGetChannel<T> (topicName, out var ch))
-			return Task.FromResult(false);
-
-		// do not allow to add more than one worker ig we are in AtMostOnce mode.
-		if (ch.Configuration.Mode == ChannelDeliveryMode.AtMostOnceAsync && workers [(topicName, type)].Count >= 1)
-			return Task.FromResult (false);
-
-		// we will have to stop consuming while we add the new worker
-		// but we do not need to close the channel, the API will buffer
-		StopConsuming<T> (topicName);
-		workers [(topicName, type)].AddRange (newWorkers);
-		return StartConsuming (topicName, ch.Configuration, ch.Channel);
+		await semaphoreSlim.WaitAsync ();
+		try {
+			// we only allow the client to register to an existing topic
+			// in this API we will not create it, there are other APIs for that
+			if (!TryGetChannel<T> (topicName, out _, out var topicInfo))
+				return false;
+
+			// do not allow to add more than one worker ig we are in AtMostOnce mode.
+			if (topicInfo.Configuration.Mode == ChannelDeliveryMode.AtMostOnceAsync && workers [(topicName, type)].Count >= 1)
+				return false;
+
+			// we will have to stop consuming while we add the new worker
+			// but we do not need to close the channel, the API will buffer
+			StopConsuming<T> (topicName);
+			workers [(topicName, type)].AddRange (newWorkers);
+			return await StartConsuming (topicName, topicInfo.Configuration, topicInfo.Channel);
+		} finally {
+			semaphoreSlim.Release ();
+		}
 	}
 
 	/// <summary>
@@ -302,11 +322,11 @@ public Task<bool> RegisterAsync<T> (string topicName, Func<T, CancellationToken,
 	/// (topicName, messageType) combination.</exception>
 	public ValueTask Publish<T> (string topicName, T publishedEvent) where T : struct
 	{
-		if (!TryGetChannel<T> (topicName, out var ch))
+		if (!TryGetChannel<T> (topicName, out _, out var topicInfo))
 			throw new InvalidOperationException (
 				$"Channel with topic {topicName} for event type {typeof(T)} not found");
 		var message = new Message<T> (MessageType.Data, publishedEvent);
-		return ch.Channel.Writer.WriteAsync (message);
+		return topicInfo.Channel.Writer.WriteAsync (message);
 	}
 
 	/// <summary>
@@ -325,17 +345,21 @@ public async Task CloseAllAsync ()
 		// `Task.WhenAll (consumingTasks!);` 
 		//
 		// suppressing the warning is ugly when we do know how to help the compiler ;)
-		
-		var consumingTasks  = from consumeInfo in cancellationTokenSources.Values
-			where consumeInfo.ConsumeTask is not null
-			select consumeInfo.ConsumeTask;
-		// remove the need of a second loop by getting the cancellation tokens cancelled
-		var cancellationTasks = cancellationTokenSources.Values
-			.Select (x => x.CancellationToken.CancelAsync ());
-		
-		// we could do a nested Task.WhenAll but we want to ensure that the cancellation tasks are done before
-		await Task.WhenAll (cancellationTasks);
-		await Task.WhenAll (consumingTasks);
+		await semaphoreSlim.WaitAsync ();
+		try {
+			var consumingTasks = from consumeInfo in cancellationTokenSources.Values
+				where consumeInfo.ConsumeTask is not null
+				select consumeInfo.ConsumeTask;
+			// remove the need of a second loop by getting the cancellation tokens cancelled
+			var cancellationTasks = cancellationTokenSources.Values
+				.Select (x => x.CancellationToken.CancelAsync ());
+
+			// we could do a nested Task.WhenAll but we want to ensure that the cancellation tasks are done before
+			await Task.WhenAll (cancellationTasks);
+			await Task.WhenAll (consumingTasks);
+		} finally {
+			semaphoreSlim.Release ();
+		}
 	}
 
 	/// <summary>
@@ -347,10 +371,15 @@ where consumeInfo.ConsumeTask is not null
 	/// <returns>A task that will be completed once the channel has been flushed.</returns>
 	public async Task<bool> CloseAsync<T> (string topicName) where T : struct
 	{
-		// ensure that the channels does exist, if not, return false
-		if (!TryGetChannel<T> (topicName, out _))
-			return false;
-		await StopConsumingAsync<T> (topicName);
-		return true;
+		await semaphoreSlim.WaitAsync ();
+		try {
+			// ensure that the channels does exist, if not, return false
+			if (!TryGetChannel<T> (topicName, out _, out _))
+				return false;
+			await StopConsumingAsync<T> (topicName);
+			return true;
+		} finally {
+			semaphoreSlim.Release ();
+		}
 	}
 }
diff --git a/Marille/Marille.csproj b/Marille/Marille.csproj
@@ -9,10 +9,16 @@
         <NoWarn>$(NoWarn);CS1591</NoWarn>
     </PropertyGroup>
 
+    <ItemGroup>
+        <AssemblyAttribute Include="System.Runtime.CompilerServices.InternalsVisibleToAttribute">
+            <_Parameter1>Marille.Tests</_Parameter1>
+        </AssemblyAttribute>
+    </ItemGroup>
+
     <PropertyGroup>
         <Title>Marille</Title>
         <PackageId>Marille</PackageId>
-        <Version>0.4.1</Version>
+        <Version>0.4.2</Version>
         <Authors>Manuel de la Peña Saenz</Authors>
         <Owners>Manuel de la Peña Saenz</Owners>
         <Copyright>Manuel de la Peña Saenz</Copyright>
diff --git a/Marille/Topic.cs b/Marille/Topic.cs
@@ -40,7 +40,8 @@ public void CloseChannel<T> () where T : struct
 	}
 
 	public bool ContainsChannel<T> ()
-	{
-		return channels.ContainsKey (typeof (T));
-	}
+		=> channels.ContainsKey (typeof (T));
+
+	public bool RemoveChannel<T> ()
+		=> channels.Remove (typeof (T));
 }

Original file line number	Diff line number	Diff line change
`@@ -40,7 +40,8 @@ public void CloseChannel<T> () where T : struct`
`40`	`40`	`}`
`41`	`41`
`42`	`42`	`public bool ContainsChannel<T> ()`
`43`		`- {`
`44`		`- return channels.ContainsKey (typeof (T));`
`45`		`- }`
	`43`	`+ => channels.ContainsKey (typeof (T));`
	`44`	`+`
	`45`	`+ public bool RemoveChannel<T> ()`
	`46`	`+ => channels.Remove (typeof (T));`
`46`	`47`	`}`