StrongReferenceMessenger.cs 40 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858
  1. // Licensed to the .NET Foundation under one or more agreements.
  2. // The .NET Foundation licenses this file to you under the MIT license.
  3. // See the LICENSE file in the project root for more information.
  4. using System;
  5. using System.Buffers;
  6. using System.Collections.Generic;
  7. using System.Diagnostics.CodeAnalysis;
  8. using System.Runtime.CompilerServices;
  9. using System.Threading;
  10. using CommunityToolkit.Mvvm.Messaging.Internals;
  11. namespace CommunityToolkit.Mvvm.Messaging;
  12. /// <summary>
  13. /// A class providing a reference implementation for the <see cref="IMessenger"/> interface.
  14. /// </summary>
  15. /// <remarks>
  16. /// This <see cref="IMessenger"/> implementation uses strong references to track the registered
  17. /// recipients, so it is necessary to manually unregister them when they're no longer needed.
  18. /// </remarks>
  19. public sealed class StrongReferenceMessenger : IMessenger
  20. {
  21. // This messenger uses the following logic to link stored instances together:
  22. // --------------------------------------------------------------------------------------------------------
  23. // Dictionary2<Recipient, HashSet<IMapping>> recipientsMap;
  24. // | \________________[*]IDictionary2<Recipient, IDictionary2<TToken>>
  25. // | \_______________[*]IDictionary2<Recipient, object?> /
  26. // | \_________/_________/___ /
  27. // |\ _(recipients registrations)_/ / \ /
  28. // | \__________________ / _____(channel registrations)_____/______\____/
  29. // | \ / / __________________________/ \
  30. // | / / / \
  31. // | Dictionary2<Recipient, object?> mapping = Mapping________________\
  32. // | __________________/ / | / \
  33. // |/ / | / \
  34. // Dictionary2<Recipient, Dictionary2<TToken, object?>> mapping = Mapping<TToken>____________\
  35. // / / / /
  36. // ___(Type2.TToken)____/ / / /
  37. // /________________(Type2.TMessage)_______/_______/__/
  38. // / ________________________________/
  39. // / /
  40. // Dictionary2<Type2, IMapping> typesMap;
  41. // --------------------------------------------------------------------------------------------------------
  42. // Each combination of <TMessage, TToken> results in a concrete Mapping type (if TToken is Unit) or Mapping<Token> type,
  43. // which holds the references from registered recipients to handlers. Mapping is used when the default channel is being
  44. // requested, as in that case there will only ever be up to a handler per recipient, per message type. In that case,
  45. // each recipient will only track the message dispatcher (stored as an object?, see notes below), instead of a dictionary
  46. // mapping each TToken value to the corresponding dispatcher for that recipient. When a custom channel is used, the
  47. // dispatchers are stored in a <TToken, object?> dictionary, so that each recipient can have up to one registered handler
  48. // for a given token, for each message type. Note that the registered dispatchers are only stored as object references, as
  49. // they can either be null or a MessageHandlerDispatcher.For<TRecipient, TMessage> instance.
  50. //
  51. // The first case happens if the handler was registered through an IRecipient<TMessage> instance, while the second one is
  52. // used to wrap input MessageHandler<TRecipient, TMessage> instances. The MessageHandlerDispatcher.For<TRecipient, TMessage>
  53. // instances will just be cast to MessageHandlerDispatcher when invoking it. This allows users to retain type information on
  54. // each registered recipient, instead of having to manually cast each recipient to the right type within the handler
  55. // (additionally, using double dispatch here avoids the need to alias delegate types). The type conversion is guaranteed to be
  56. // respected due to how the messenger type itself works - as registered handlers are always invoked on their respective recipients.
  57. //
  58. // Each mapping is stored in the types map, which associates each pair of concrete types to its mapping instance. Mapping instances
  59. // are exposed as IMapping items, as each will be a closed type over a different combination of TMessage and TToken generic type
  60. // parameters (or just of TMessage, for the default channel). Each existing recipient is also stored in the main recipients map,
  61. // along with a set of all the existing (dictionaries of) handlers for that recipient (for all message types and token types, if any).
  62. //
  63. // A recipient is stored in the main map as long as it has at least one registered handler in any of the existing mappings for every
  64. // message/token type combination. The shared map is used to access the set of all registered handlers for a given recipient, without
  65. // having to know in advance the type of message or token being used for the registration, and without having to use reflection. This
  66. // is the same approach used in the types map, as we expose saved items as IMapping values too.
  67. //
  68. // Note that each mapping stored in the associated set for each recipient also indirectly implements either IDictionary2<Recipient, Token>
  69. // or IDictionary2<Recipient>, with any token type currently in use by that recipient (or none, if using the default channel). This allows
  70. // to retrieve the type-closed mappings of registered handlers with a given token type, for any message type, for every receiver, again
  71. // without having to use reflection. This shared map is used to unregister messages from a given recipients either unconditionally, by
  72. // message type, by token, or for a specific pair of message type and token value.
  73. /// <summary>
  74. /// The collection of currently registered recipients, with a link to their linked message receivers.
  75. /// </summary>
  76. /// <remarks>
  77. /// This collection is used to allow reflection-free access to all the existing
  78. /// registered recipients from <see cref="UnregisterAll"/> and other methods in this type,
  79. /// so that all the existing handlers can be removed without having to dynamically create
  80. /// the generic types for the containers of the various dictionaries mapping the handlers.
  81. /// </remarks>
  82. private readonly Dictionary2<Recipient, HashSet<IMapping>> recipientsMap = new();
  83. /// <summary>
  84. /// The <see cref="Mapping"/> and <see cref="Mapping{TToken}"/> instance for types combination.
  85. /// </summary>
  86. /// <remarks>
  87. /// The values are just of type <see cref="IDictionary2{T}"/> as we don't know the type parameters in advance.
  88. /// Each method relies on <see cref="GetOrAddMapping{TMessage,TToken}"/> to get the type-safe instance of the
  89. /// <see cref="Mapping"/> or <see cref="Mapping{TToken}"/> class for each pair of generic arguments in use.
  90. /// </remarks>
  91. private readonly Dictionary2<Type2, IMapping> typesMap = new();
  92. /// <summary>
  93. /// Gets the default <see cref="StrongReferenceMessenger"/> instance.
  94. /// </summary>
  95. public static StrongReferenceMessenger Default { get; } = new();
  96. /// <inheritdoc/>
  97. public bool IsRegistered<TMessage, TToken>(object recipient, TToken token)
  98. where TMessage : class
  99. where TToken : IEquatable<TToken>
  100. {
  101. ArgumentNullException.ThrowIfNull(recipient);
  102. ArgumentNullException.For<TToken>.ThrowIfNull(token);
  103. lock (this.recipientsMap)
  104. {
  105. if (typeof(TToken) == typeof(Unit))
  106. {
  107. if (!TryGetMapping<TMessage>(out Mapping? mapping))
  108. {
  109. return false;
  110. }
  111. Recipient key = new(recipient);
  112. return mapping.ContainsKey(key);
  113. }
  114. else
  115. {
  116. if (!TryGetMapping<TMessage, TToken>(out Mapping<TToken>? mapping))
  117. {
  118. return false;
  119. }
  120. Recipient key = new(recipient);
  121. return
  122. mapping.TryGetValue(key, out Dictionary2<TToken, object?>? handlers) &&
  123. handlers.ContainsKey(token);
  124. }
  125. }
  126. }
  127. /// <inheritdoc/>
  128. public void Register<TRecipient, TMessage, TToken>(TRecipient recipient, TToken token, MessageHandler<TRecipient, TMessage> handler)
  129. where TRecipient : class
  130. where TMessage : class
  131. where TToken : IEquatable<TToken>
  132. {
  133. ArgumentNullException.ThrowIfNull(recipient);
  134. ArgumentNullException.For<TToken>.ThrowIfNull(token);
  135. ArgumentNullException.ThrowIfNull(handler);
  136. Register<TMessage, TToken>(recipient, token, new MessageHandlerDispatcher.For<TRecipient, TMessage>(handler));
  137. }
  138. /// <inheritdoc cref="WeakReferenceMessenger.Register{TMessage, TToken}(IRecipient{TMessage}, TToken)"/>
  139. internal void Register<TMessage, TToken>(IRecipient<TMessage> recipient, TToken token)
  140. where TMessage : class
  141. where TToken : IEquatable<TToken>
  142. {
  143. Register<TMessage, TToken>(recipient, token, null);
  144. }
  145. /// <summary>
  146. /// Registers a recipient for a given type of message.
  147. /// </summary>
  148. /// <typeparam name="TMessage">The type of message to receive.</typeparam>
  149. /// <typeparam name="TToken">The type of token to use to pick the messages to receive.</typeparam>
  150. /// <param name="recipient">The recipient that will receive the messages.</param>
  151. /// <param name="token">A token used to determine the receiving channel to use.</param>
  152. /// <param name="dispatcher">The input <see cref="MessageHandlerDispatcher"/> instance to register, or null.</param>
  153. /// <exception cref="InvalidOperationException">Thrown when trying to register the same message twice.</exception>
  154. private void Register<TMessage, TToken>(object recipient, TToken token, MessageHandlerDispatcher? dispatcher)
  155. where TMessage : class
  156. where TToken : IEquatable<TToken>
  157. {
  158. lock (this.recipientsMap)
  159. {
  160. Recipient key = new(recipient);
  161. IMapping mapping;
  162. // Fast path for unit tokens
  163. if (typeof(TToken) == typeof(Unit))
  164. {
  165. // Get the <TMessage> registration list for this recipient
  166. Mapping underlyingMapping = GetOrAddMapping<TMessage>();
  167. ref object? registeredHandler = ref underlyingMapping.GetOrAddValueRef(key);
  168. if (registeredHandler is not null)
  169. {
  170. ThrowInvalidOperationExceptionForDuplicateRegistration();
  171. }
  172. // Store the input handler
  173. registeredHandler = dispatcher;
  174. mapping = underlyingMapping;
  175. }
  176. else
  177. {
  178. // Get the <TMessage, TToken> registration list for this recipient
  179. Mapping<TToken> underlyingMapping = GetOrAddMapping<TMessage, TToken>();
  180. ref Dictionary2<TToken, object?>? map = ref underlyingMapping.GetOrAddValueRef(key);
  181. map ??= new Dictionary2<TToken, object?>();
  182. // Add the new registration entry
  183. ref object? registeredHandler = ref map.GetOrAddValueRef(token);
  184. if (registeredHandler is not null)
  185. {
  186. ThrowInvalidOperationExceptionForDuplicateRegistration();
  187. }
  188. registeredHandler = dispatcher;
  189. mapping = underlyingMapping;
  190. }
  191. // Make sure this registration map is tracked for the current recipient
  192. ref HashSet<IMapping>? set = ref this.recipientsMap.GetOrAddValueRef(key);
  193. set ??= new HashSet<IMapping>();
  194. _ = set.Add(mapping);
  195. }
  196. }
  197. /// <inheritdoc/>
  198. public void UnregisterAll(object recipient)
  199. {
  200. ArgumentNullException.ThrowIfNull(recipient);
  201. lock (this.recipientsMap)
  202. {
  203. // If the recipient has no registered messages at all, ignore
  204. Recipient key = new(recipient);
  205. if (!this.recipientsMap.TryGetValue(key, out HashSet<IMapping>? set))
  206. {
  207. return;
  208. }
  209. // Removes all the lists of registered handlers for the recipient
  210. foreach (IMapping mapping in set)
  211. {
  212. if (mapping.TryRemove(key) &&
  213. mapping.Count == 0)
  214. {
  215. // Maps here are really of type Mapping<,> and with unknown type arguments.
  216. // If after removing the current recipient a given map becomes empty, it means
  217. // that there are no registered recipients at all for a given pair of message
  218. // and token types. In that case, we also remove the map from the types map.
  219. // The reason for keeping a key in each mapping is that removing items from a
  220. // dictionary (a hashed collection) only costs O(1) in the best case, while
  221. // if we had tried to iterate the whole dictionary every time we would have
  222. // paid an O(n) minimum cost for each single remove operation.
  223. _ = this.typesMap.TryRemove(mapping.TypeArguments);
  224. }
  225. }
  226. // Remove the associated set in the recipients map
  227. _ = this.recipientsMap.TryRemove(key);
  228. }
  229. }
  230. /// <inheritdoc/>
  231. public void UnregisterAll<TToken>(object recipient, TToken token)
  232. where TToken : IEquatable<TToken>
  233. {
  234. ArgumentNullException.ThrowIfNull(recipient);
  235. ArgumentNullException.For<TToken>.ThrowIfNull(token);
  236. // This method is never called with the unit type, so this path is not implemented. This
  237. // exception should not ever be thrown, it's here just to double check for regressions in
  238. // case a bug was introduced that caused this path to somehow be invoked with the Unit type.
  239. // This type is internal, so consumers of the library would never be able to pass it here,
  240. // and there are (and shouldn't be) any APIs publicly exposed from the library that would
  241. // cause this path to be taken either. When using the default channel, only UnregisterAll(object)
  242. // is supported, which would just unregister all recipients regardless of the selected channel.
  243. if (typeof(TToken) == typeof(Unit))
  244. {
  245. throw new NotImplementedException();
  246. }
  247. bool lockTaken = false;
  248. object[]? maps = null;
  249. int i = 0;
  250. // We use an explicit try/finally block here instead of the lock syntax so that we can use a single
  251. // one both to release the lock and to clear the rented buffer and return it to the pool. The reason
  252. // why we're declaring the buffer here and clearing and returning it in this outer finally block is
  253. // that doing so doesn't require the lock to be kept, and releasing it before performing this last
  254. // step reduces the total time spent while the lock is acquired, which in turn reduces the lock
  255. // contention in multi-threaded scenarios where this method is invoked concurrently.
  256. try
  257. {
  258. Monitor.Enter(this.recipientsMap, ref lockTaken);
  259. // Get the shared set of mappings for the recipient, if present
  260. Recipient key = new(recipient);
  261. if (!this.recipientsMap.TryGetValue(key, out HashSet<IMapping>? set))
  262. {
  263. return;
  264. }
  265. // Copy the candidate mappings for the target recipient to a local array, as we can't modify the
  266. // contents of the set while iterating it. The rented buffer is oversized and will also include
  267. // mappings for handlers of messages that are registered through a different token. Note that
  268. // we're using just an object array to minimize the number of total rented buffers, that would
  269. // just remain in the shared pool unused, other than when they are rented here. Instead, we're
  270. // using a type that would possibly also be used by the users of the library, which increases
  271. // the opportunities to reuse existing buffers for both. When we need to reference an item
  272. // stored in the buffer with the type we know it will have, we use Unsafe.As<T> to avoid the
  273. // expensive type check in the cast, since we already know the assignment will be valid.
  274. maps = ArrayPool<object>.Shared.Rent(set.Count);
  275. foreach (IMapping item in set)
  276. {
  277. // Select all mappings using the same token type
  278. if (item is IDictionary2<Recipient, IDictionary2<TToken>> mapping)
  279. {
  280. maps[i++] = mapping;
  281. }
  282. }
  283. // Iterate through all the local maps. These are all the currently
  284. // existing maps of handlers for messages of any given type, with a token
  285. // of the current type, for the target recipient. We heavily rely on
  286. // interfaces here to be able to iterate through all the available mappings
  287. // without having to know the concrete type in advance, and without having
  288. // to deal with reflection: we can just check if the type of the closed interface
  289. // matches with the token type currently in use, and operate on those instances.
  290. foreach (object obj in maps.AsSpan(0, i))
  291. {
  292. IDictionary2<Recipient, IDictionary2<TToken>>? handlersMap = Unsafe.As<IDictionary2<Recipient, IDictionary2<TToken>>>(obj);
  293. // We don't need whether or not the map contains the recipient, as the
  294. // sequence of maps has already been copied from the set containing all
  295. // the mappings for the target recipients: it is guaranteed to be here.
  296. IDictionary2<TToken> holder = handlersMap[key];
  297. // Try to remove the registered handler for the input token,
  298. // for the current message type (unknown from here).
  299. if (holder.TryRemove(token) &&
  300. holder.Count == 0)
  301. {
  302. // If the map is empty, remove the recipient entirely from its container
  303. _ = handlersMap.TryRemove(key);
  304. IMapping mapping = Unsafe.As<IMapping>(handlersMap);
  305. // This recipient has no registrations left for this combination of token
  306. // and message type, so this mapping can be removed from its associated set.
  307. _ = set.Remove(mapping);
  308. // If the resulting set is empty, then this means that there are no more handlers
  309. // left for this recipient for any message or token type, so the recipient can also
  310. // be removed from the map of all existing recipients with at least one handler.
  311. if (set.Count == 0)
  312. {
  313. _ = this.recipientsMap.TryRemove(key);
  314. }
  315. // If no handlers are left at all for any recipient, across all message types and token
  316. // types, remove the set of mappings entirely for the current recipient, and remove the
  317. // strong reference to it as well. This is the same situation that would've been achieved
  318. // by just calling UnregisterAll(recipient).
  319. if (handlersMap.Count == 0)
  320. {
  321. _ = this.typesMap.TryRemove(mapping.TypeArguments);
  322. }
  323. }
  324. }
  325. }
  326. finally
  327. {
  328. // Release the lock, if we did acquire it
  329. if (lockTaken)
  330. {
  331. Monitor.Exit(this.recipientsMap);
  332. }
  333. // If we got to renting the array of maps, return it to the shared pool.
  334. // Remove references to avoid leaks coming from the shared memory pool.
  335. // We manually create a span and clear it as a small optimization, as
  336. // arrays rented from the pool can be larger than the requested size.
  337. if (maps is not null)
  338. {
  339. maps.AsSpan(0, i).Clear();
  340. ArrayPool<object>.Shared.Return(maps);
  341. }
  342. }
  343. }
  344. /// <inheritdoc/>
  345. public void Unregister<TMessage, TToken>(object recipient, TToken token)
  346. where TMessage : class
  347. where TToken : IEquatable<TToken>
  348. {
  349. ArgumentNullException.ThrowIfNull(recipient);
  350. ArgumentNullException.For<TToken>.ThrowIfNull(token);
  351. lock (this.recipientsMap)
  352. {
  353. if (typeof(TToken) == typeof(Unit))
  354. {
  355. // Get the registration list, if available
  356. if (!TryGetMapping<TMessage>(out Mapping? mapping))
  357. {
  358. return;
  359. }
  360. Recipient key = new(recipient);
  361. // Remove the handler (there can only be one for the unit type)
  362. if (!mapping.TryRemove(key))
  363. {
  364. return;
  365. }
  366. // Remove the map entirely from this container, and remove the link to the map itself to
  367. // the current mapping between existing registered recipients (or entire recipients too).
  368. // This is the same as below, except for the unit type there can only be one handler, so
  369. // removing it already implies the target recipient has no remaining handlers left.
  370. _ = mapping.TryRemove(key);
  371. // If there are no handlers left at all for this type combination, drop it
  372. if (mapping.Count == 0)
  373. {
  374. _ = this.typesMap.TryRemove(mapping.TypeArguments);
  375. }
  376. HashSet<IMapping> set = this.recipientsMap[key];
  377. // The current mapping no longer has any handlers left for this recipient.
  378. // Remove it and then also remove the recipient if this was the last handler.
  379. // Again, this is the same as below, except with the assumption of the unit type.
  380. _ = set.Remove(mapping);
  381. if (set.Count == 0)
  382. {
  383. _ = this.recipientsMap.TryRemove(key);
  384. }
  385. }
  386. else
  387. {
  388. // Get the registration list, if available
  389. if (!TryGetMapping<TMessage, TToken>(out Mapping<TToken>? mapping))
  390. {
  391. return;
  392. }
  393. Recipient key = new(recipient);
  394. if (!mapping.TryGetValue(key, out Dictionary2<TToken, object?>? dictionary))
  395. {
  396. return;
  397. }
  398. // Remove the target handler
  399. if (dictionary.TryRemove(token) &&
  400. dictionary.Count == 0)
  401. {
  402. // If the map is empty, it means that the current recipient has no remaining
  403. // registered handlers for the current <TMessage, TToken> combination, regardless,
  404. // of the specific token value (ie. the channel used to receive messages of that type).
  405. // We can remove the map entirely from this container, and remove the link to the map itself
  406. // to the current mapping between existing registered recipients (or entire recipients too).
  407. _ = mapping.TryRemove(key);
  408. // If there are no handlers left at all for this type combination, drop it
  409. if (mapping.Count == 0)
  410. {
  411. _ = this.typesMap.TryRemove(mapping.TypeArguments);
  412. }
  413. HashSet<IMapping> set = this.recipientsMap[key];
  414. // The current mapping no longer has any handlers left for this recipient
  415. _ = set.Remove(mapping);
  416. // If the current recipients has no handlers left at all, remove it
  417. if (set.Count == 0)
  418. {
  419. _ = this.recipientsMap.TryRemove(key);
  420. }
  421. }
  422. }
  423. }
  424. }
  425. /// <inheritdoc/>
  426. public TMessage Send<TMessage, TToken>(TMessage message, TToken token)
  427. where TMessage : class
  428. where TToken : IEquatable<TToken>
  429. {
  430. ArgumentNullException.ThrowIfNull(message);
  431. ArgumentNullException.For<TToken>.ThrowIfNull(token);
  432. object?[] rentedArray;
  433. Span<object?> pairs;
  434. int i = 0;
  435. lock (this.recipientsMap)
  436. {
  437. if (typeof(TToken) == typeof(Unit))
  438. {
  439. // Check whether there are any registered recipients
  440. if (!TryGetMapping<TMessage>(out Mapping? mapping))
  441. {
  442. goto End;
  443. }
  444. // Check the number of remaining handlers, see below
  445. int totalHandlersCount = mapping.Count;
  446. if (totalHandlersCount == 0)
  447. {
  448. goto End;
  449. }
  450. pairs = rentedArray = ArrayPool<object?>.Shared.Rent(2 * totalHandlersCount);
  451. // Same logic as below, except here we're only traversing one handler per recipient
  452. Dictionary2<Recipient, object?>.Enumerator mappingEnumerator = mapping.GetEnumerator();
  453. while (mappingEnumerator.MoveNext())
  454. {
  455. pairs[2 * i] = mappingEnumerator.GetValue();
  456. pairs[(2 * i) + 1] = mappingEnumerator.GetKey().Target;
  457. i++;
  458. }
  459. }
  460. else
  461. {
  462. // Check whether there are any registered recipients
  463. if (!TryGetMapping<TMessage, TToken>(out Mapping<TToken>? mapping))
  464. {
  465. goto End;
  466. }
  467. // We need to make a local copy of the currently registered handlers, since users might
  468. // try to unregister (or register) new handlers from inside one of the currently existing
  469. // handlers. We can use memory pooling to reuse arrays, to minimize the average memory
  470. // usage. In practice, we usually just need to pay the small overhead of copying the items.
  471. // The current mapping contains all the currently registered recipients and handlers for
  472. // the <TMessage, TToken> combination in use. In the worst case scenario, all recipients
  473. // will have a registered handler with a token matching the input one, meaning that we could
  474. // have at worst a number of pending handlers to invoke equal to the total number of recipient
  475. // in the mapping. This relies on the fact that tokens are unique, and that there is only
  476. // one handler associated with a given token. We can use this upper bound as the requested
  477. // size for each array rented from the pool, which guarantees that we'll have enough space.
  478. int totalHandlersCount = mapping.Count;
  479. if (totalHandlersCount == 0)
  480. {
  481. goto End;
  482. }
  483. // Rent the array and also assign it to a span, which will be used to access values.
  484. // We're doing this to avoid the array covariance checks slowdown in the loops below.
  485. pairs = rentedArray = ArrayPool<object?>.Shared.Rent(2 * totalHandlersCount);
  486. // Copy the handlers to the local collection.
  487. // The array is oversized at this point, since it also includes
  488. // handlers for different tokens. We can reuse the same variable
  489. // to count the number of matching handlers to invoke later on.
  490. // This will be the array slice with valid handler in the rented buffer.
  491. Dictionary2<Recipient, Dictionary2<TToken, object?>>.Enumerator mappingEnumerator = mapping.GetEnumerator();
  492. // Explicit enumerator usage here as we're using a custom one
  493. // that doesn't expose the single standard Current property.
  494. while (mappingEnumerator.MoveNext())
  495. {
  496. // Pick the target handler, if the token is a match for the recipient
  497. if (mappingEnumerator.GetValue().TryGetValue(token, out object? handler))
  498. {
  499. // This span access should always guaranteed to be valid due to the size of the
  500. // array being set according to the current total number of registered handlers,
  501. // which will always be greater or equal than the ones matching the previous test.
  502. // We're still using a checked span accesses here though to make sure an out of
  503. // bounds write can never happen even if an error was present in the logic above.
  504. pairs[2 * i] = handler;
  505. pairs[(2 * i) + 1] = mappingEnumerator.GetKey().Target;
  506. i++;
  507. }
  508. }
  509. }
  510. }
  511. try
  512. {
  513. // The core broadcasting logic is the same as the weak reference messenger one
  514. WeakReferenceMessenger.SendAll(pairs, i, message);
  515. }
  516. finally
  517. {
  518. // As before, we also need to clear it first to avoid having potentially long
  519. // lasting memory leaks due to leftover references being stored in the pool.
  520. Array.Clear(rentedArray, 0, 2 * i);
  521. ArrayPool<object?>.Shared.Return(rentedArray);
  522. }
  523. End:
  524. return message;
  525. }
  526. /// <inheritdoc/>
  527. void IMessenger.Cleanup()
  528. {
  529. // The current implementation doesn't require any kind of cleanup operation, as
  530. // all the internal data structures are already kept in sync whenever a recipient
  531. // is added or removed. This method is implemented through an explicit interface
  532. // implementation so that developers using this type directly will not see it in
  533. // the API surface (as it wouldn't be useful anyway, since it's a no-op here).
  534. }
  535. /// <inheritdoc/>
  536. public void Reset()
  537. {
  538. lock (this.recipientsMap)
  539. {
  540. this.recipientsMap.Clear();
  541. this.typesMap.Clear();
  542. }
  543. }
  544. /// <summary>
  545. /// Tries to get the <see cref="Mapping"/> instance of currently
  546. /// registered recipients for the input <typeparamref name="TMessage"/> type.
  547. /// </summary>
  548. /// <typeparam name="TMessage">The type of message to send.</typeparam>
  549. /// <param name="mapping">The resulting <see cref="Mapping"/> instance, if found.</param>
  550. /// <returns>Whether or not the required <see cref="Mapping"/> instance was found.</returns>
  551. [MethodImpl(MethodImplOptions.AggressiveInlining)]
  552. private bool TryGetMapping<TMessage>([NotNullWhen(true)] out Mapping? mapping)
  553. where TMessage : class
  554. {
  555. Type2 key = new(typeof(TMessage), typeof(Unit));
  556. if (this.typesMap.TryGetValue(key, out IMapping? target))
  557. {
  558. // This method and the ones below are the only ones handling values in the types map,
  559. // and here we are sure that the object reference we have points to an instance of the
  560. // right type. Using an unsafe cast skips two conditional branches and is faster.
  561. mapping = Unsafe.As<Mapping>(target);
  562. return true;
  563. }
  564. mapping = null;
  565. return false;
  566. }
  567. /// <summary>
  568. /// Tries to get the <see cref="Mapping{TToken}"/> instance of currently registered recipients
  569. /// for the combination of types <typeparamref name="TMessage"/> and <typeparamref name="TToken"/>.
  570. /// </summary>
  571. /// <typeparam name="TMessage">The type of message to send.</typeparam>
  572. /// <typeparam name="TToken">The type of token to identify what channel to use to send the message.</typeparam>
  573. /// <param name="mapping">The resulting <see cref="Mapping{TToken}"/> instance, if found.</param>
  574. /// <returns>Whether or not the required <see cref="Mapping{TToken}"/> instance was found.</returns>
  575. [MethodImpl(MethodImplOptions.AggressiveInlining)]
  576. private bool TryGetMapping<TMessage, TToken>([NotNullWhen(true)] out Mapping<TToken>? mapping)
  577. where TMessage : class
  578. where TToken : IEquatable<TToken>
  579. {
  580. Type2 key = new(typeof(TMessage), typeof(TToken));
  581. if (this.typesMap.TryGetValue(key, out IMapping? target))
  582. {
  583. mapping = Unsafe.As<Mapping<TToken>>(target);
  584. return true;
  585. }
  586. mapping = null;
  587. return false;
  588. }
  589. /// <summary>
  590. /// Gets the <see cref="Mapping"/> instance of currently
  591. /// registered recipients for the input <typeparamref name="TMessage"/> type.
  592. /// </summary>
  593. /// <typeparam name="TMessage">The type of message to send.</typeparam>
  594. /// <returns>A <see cref="Mapping"/> instance with the requested type arguments.</returns>
  595. [MethodImpl(MethodImplOptions.AggressiveInlining)]
  596. private Mapping GetOrAddMapping<TMessage>()
  597. where TMessage : class
  598. {
  599. Type2 key = new(typeof(TMessage), typeof(Unit));
  600. ref IMapping? target = ref this.typesMap.GetOrAddValueRef(key);
  601. target ??= Mapping.Create<TMessage>();
  602. return Unsafe.As<Mapping>(target);
  603. }
  604. /// <summary>
  605. /// Gets the <see cref="Mapping{TToken}"/> instance of currently registered recipients
  606. /// for the combination of types <typeparamref name="TMessage"/> and <typeparamref name="TToken"/>.
  607. /// </summary>
  608. /// <typeparam name="TMessage">The type of message to send.</typeparam>
  609. /// <typeparam name="TToken">The type of token to identify what channel to use to send the message.</typeparam>
  610. /// <returns>A <see cref="Mapping{TToken}"/> instance with the requested type arguments.</returns>
  611. [MethodImpl(MethodImplOptions.AggressiveInlining)]
  612. private Mapping<TToken> GetOrAddMapping<TMessage, TToken>()
  613. where TMessage : class
  614. where TToken : IEquatable<TToken>
  615. {
  616. Type2 key = new(typeof(TMessage), typeof(TToken));
  617. ref IMapping? target = ref this.typesMap.GetOrAddValueRef(key);
  618. target ??= Mapping<TToken>.Create<TMessage>();
  619. return Unsafe.As<Mapping<TToken>>(target);
  620. }
  621. /// <summary>
  622. /// A mapping type representing a link to recipients and their view of handlers per communication channel.
  623. /// </summary>
  624. /// <remarks>
  625. /// This type is a specialization of <see cref="Mapping{TToken}"/> for <see cref="Unit"/> tokens.
  626. /// </remarks>
  627. private sealed class Mapping : Dictionary2<Recipient, object?>, IMapping
  628. {
  629. /// <summary>
  630. /// Initializes a new instance of the <see cref="Mapping"/> class.
  631. /// </summary>
  632. /// <param name="messageType">The message type being used.</param>
  633. private Mapping(Type messageType)
  634. {
  635. TypeArguments = new Type2(messageType, typeof(Unit));
  636. }
  637. /// <summary>
  638. /// Creates a new instance of the <see cref="Mapping"/> class.
  639. /// </summary>
  640. /// <typeparam name="TMessage">The type of message to receive.</typeparam>
  641. /// <returns>A new <see cref="Mapping"/> instance.</returns>
  642. [MethodImpl(MethodImplOptions.AggressiveInlining)]
  643. public static Mapping Create<TMessage>()
  644. where TMessage : class
  645. {
  646. return new(typeof(TMessage));
  647. }
  648. /// <inheritdoc/>
  649. public Type2 TypeArguments { get; }
  650. }
  651. /// <summary>
  652. /// A mapping type representing a link to recipients and their view of handlers per communication channel.
  653. /// </summary>
  654. /// <typeparam name="TToken">The type of token to use to pick the messages to receive.</typeparam>
  655. /// <remarks>
  656. /// This type is defined for simplicity and as a workaround for the lack of support for using type aliases
  657. /// over open generic types in C# (using type aliases can only be used for concrete, closed types).
  658. /// </remarks>
  659. private sealed class Mapping<TToken> : Dictionary2<Recipient, Dictionary2<TToken, object?>>, IMapping
  660. where TToken : IEquatable<TToken>
  661. {
  662. /// <summary>
  663. /// Initializes a new instance of the <see cref="Mapping{TToken}"/> class.
  664. /// </summary>
  665. /// <param name="messageType">The message type being used.</param>
  666. private Mapping(Type messageType)
  667. {
  668. TypeArguments = new Type2(messageType, typeof(TToken));
  669. }
  670. /// <summary>
  671. /// Creates a new instance of the <see cref="Mapping{TToken}"/> class.
  672. /// </summary>
  673. /// <typeparam name="TMessage">The type of message to receive.</typeparam>
  674. /// <returns>A new <see cref="Mapping{TToken}"/> instance.</returns>
  675. [MethodImpl(MethodImplOptions.AggressiveInlining)]
  676. public static Mapping<TToken> Create<TMessage>()
  677. where TMessage : class
  678. {
  679. return new(typeof(TMessage));
  680. }
  681. /// <inheritdoc/>
  682. public Type2 TypeArguments { get; }
  683. }
  684. /// <summary>
  685. /// An interface for the <see cref="Mapping"/> and <see cref="Mapping{TToken}"/> types which allows to retrieve
  686. /// the type arguments from a given generic instance without having any prior knowledge about those arguments.
  687. /// </summary>
  688. private interface IMapping : IDictionary2<Recipient>
  689. {
  690. /// <summary>
  691. /// Gets the <see cref="Type2"/> instance representing the current type arguments.
  692. /// </summary>
  693. Type2 TypeArguments { get; }
  694. }
  695. /// <summary>
  696. /// A simple type representing a recipient.
  697. /// </summary>
  698. /// <remarks>
  699. /// This type is used to enable fast indexing in each mapping dictionary,
  700. /// since it acts as an external override for the <see cref="GetHashCode"/> and
  701. /// <see cref="Equals(object?)"/> methods for arbitrary objects, removing both
  702. /// the virtual call and preventing instances overriding those methods in this context.
  703. /// Using this type guarantees that all the equality operations are always only done
  704. /// based on reference equality for each registered recipient, regardless of its type.
  705. /// </remarks>
  706. private readonly struct Recipient : IEquatable<Recipient>
  707. {
  708. /// <summary>
  709. /// The registered recipient.
  710. /// </summary>
  711. public readonly object Target;
  712. /// <summary>
  713. /// Initializes a new instance of the <see cref="Recipient"/> struct.
  714. /// </summary>
  715. /// <param name="target">The target recipient instance.</param>
  716. [MethodImpl(MethodImplOptions.AggressiveInlining)]
  717. public Recipient(object target)
  718. {
  719. this.Target = target;
  720. }
  721. /// <inheritdoc/>
  722. [MethodImpl(MethodImplOptions.AggressiveInlining)]
  723. public bool Equals(Recipient other)
  724. {
  725. return ReferenceEquals(this.Target, other.Target);
  726. }
  727. /// <inheritdoc/>
  728. public override bool Equals(object? obj)
  729. {
  730. return obj is Recipient other && Equals(other);
  731. }
  732. /// <inheritdoc/>
  733. [MethodImpl(MethodImplOptions.AggressiveInlining)]
  734. public override int GetHashCode()
  735. {
  736. return RuntimeHelpers.GetHashCode(this.Target);
  737. }
  738. }
  739. /// <summary>
  740. /// Throws an <see cref="InvalidOperationException"/> when trying to add a duplicate handler.
  741. /// </summary>
  742. private static void ThrowInvalidOperationExceptionForDuplicateRegistration()
  743. {
  744. throw new InvalidOperationException("The target recipient has already subscribed to the target message.");
  745. }
  746. }