Use cache in BlockchainController.GetTransactionsAsync (take 2)

WalletWasabi.Backend/Controllers/BlockchainController.cs

@@ -2,11 +2,15 @@
 using Microsoft.Extensions.Caching.Memory;
 using NBitcoin;
 using NBitcoin.RPC;
+using Newtonsoft.Json.Linq;
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.ComponentModel.DataAnnotations;
+using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
+using System.Net.Cache;
 using System.Threading;
 using System.Threading.Tasks;
 using WalletWasabi.Backend.Models;
@@ -32,6 +36,7 @@
 {
 	public static readonly TimeSpan FilterTimeout = TimeSpan.FromMinutes(20);
 	private static readonly MemoryCacheEntryOptions CacheEntryOptions = new() { AbsoluteExpirationRelativeToNow = TimeSpan.FromSeconds(60) };
+	private static MemoryCacheEntryOptions TransactionCacheOptions { get; } = new MemoryCacheEntryOptions { AbsoluteExpirationRelativeToNow = TimeSpan.FromMinutes(1) };
 
 	public BlockchainController(IMemoryCache memoryCache, Global global)
 	{
@@ -42,8 +47,6 @@
 	private IRPCClient RpcClient => Global.RpcClient;
 	private Network Network => Global.Config.Network;
 
-	public static Dictionary<uint256, string> TransactionHexCache { get; } = new();
-	public static object TransactionHexCacheLock { get; } = new();
 	public IdempotencyRequestCache Cache { get; }
 
 	public Global Global { get; }
@@ -140,76 +143,83 @@
 	[HttpGet("transaction-hexes")]
 	[ProducesResponseType(200)]
 	[ProducesResponseType(400)]
-	public async Task<IActionResult> GetTransactionsAsync([FromQuery, Required] IEnumerable<string> transactionIds, CancellationToken cancellationToken)
+	public async Task<IActionResult> GetTransactionsAsync([FromQuery, Required] List<string> transactionIds, CancellationToken cancellationToken)
 	{
-		var maxTxToRequest = 10;
-		if (transactionIds.Count() > maxTxToRequest)
+		const int MaxTxToRequest = 10;
+		int requestCount = transactionIds.Count;
+
+		if (requestCount > MaxTxToRequest)
 		{
-			return BadRequest($"Maximum {maxTxToRequest} transactions can be requested.");
+			return BadRequest($"Maximum {MaxTxToRequest} transactions can be requested.");
 		}
 
-		var parsedIds = new List<uint256>();
+		uint256[] parsedTxIds;
+
+		// Make sure TXIDs are not malformed.
 		try
 		{
-			// Remove duplicates, do not use Distinct(), order is not guaranteed.
-			foreach (var txid in transactionIds.Select(x => new uint256(x)))
-			{
-				if (!parsedIds.Contains(txid))
-				{
-					parsedIds.Add(txid);
-				}
-			}
+			parsedTxIds = transactionIds.Select(x => new uint256(x)).ToArray();
 		}
 		catch
 		{
 			return BadRequest("Invalid transaction Ids.");
 		}
 
+		Dictionary<uint256, TaskCompletionSource<Transaction>> txIdsRetrieve = [];
+		TaskCompletionSource<Transaction>[] txsCompletionSources = new TaskCompletionSource<Transaction>[requestCount];
+
 		try
 		{
-			var hexes = new Dictionary<uint256, string>();
-			List<uint256> missingTxs = new();
-			lock (TransactionHexCacheLock)
+			// Get task completion sources for transactions. They are either new (no one else is getting that transaction right now) or existing
+			// and then some other caller needs the same transaction so we can use the existing task completion source.
+			for (int i = 0; i < requestCount; i++)
 			{
-				foreach (var txid in parsedIds)
+				uint256 txId = parsedTxIds[i];
+				string cacheKey = $"{nameof(BlockchainController)}#{txId}";
+
+				if (Cache.TryAddKey(cacheKey, TransactionCacheOptions, out TaskCompletionSource<Transaction> tcs))
 				{
-					if (TransactionHexCache.TryGetValue(txid, out string? hex))
-					{
-						hexes.Add(txid, hex);
-					}
-					else
-					{
-						missingTxs.Add(txid);
-					}
+					txIdsRetrieve.Add(txId, tcs);
 				}
+
+				txsCompletionSources[i] = tcs;
 			}
 
-			if (missingTxs.Count != 0)
+			if (txIdsRetrieve.Count > 0)
 			{
-				foreach (var tx in await RpcClient.GetRawTransactionsAsync(missingTxs, cancellationToken))
+				// Ask to get missing transactions over RPC.
+				IEnumerable<Transaction> txs = await RpcClient.GetRawTransactionsAsync(txIdsRetrieve.Keys, cancellationToken).ConfigureAwait(false);
+				Dictionary<uint256, Transaction> rpcBatch = txs.ToDictionary(x => x.GetHash(), x => x);
+
+				foreach (KeyValuePair<uint256, Transaction> kvp in rpcBatch)
 				{
-					string hex = tx.ToHex();
-					hexes.Add(tx.GetHash(), hex);
-
-					lock (TransactionHexCacheLock)
-					{
-						if (TransactionHexCache.TryAdd(tx.GetHash(), hex) && TransactionHexCache.Count >= 1000)
-						{
-							TransactionHexCache.Remove(TransactionHexCache.Keys.First());
-						}
-					}
+					_ = txIdsRetrieve[kvp.Key].TrySetResult(kvp.Value);
 				}
 			}
 
-			// Order hexes according to the order of the query.
-			var orderedResult = parsedIds.Where(x => hexes.ContainsKey(x)).Select(x => hexes[x]);
-			return Ok(orderedResult);
+			Transaction[] result = new Transaction[requestCount];
+
+			// Add missing transactions to the result array.
+			for (int i = 0; i < requestCount; i++)
+			{
+				uint256 txId = parsedTxIds[i];
+				result[i] = await txsCompletionSources[i].Task.ConfigureAwait(false);
+			}
 		}
-		catch (Exception ex)
+		finally
 		{
-			Logger.LogDebug(ex);
-			return BadRequest(ex.Message);
+			if (txIdsRetrieve.Count > 0)
+			{
+				// It's necessary to always set a result to the task completion sources. Otherwise, cache can get corrupted.
+				Exception ex = new InvalidOperationException("Failed to get the transaction.");
+				foreach (TaskCompletionSource<Transaction> tcs in txIdsRetrieve.Values)
+				{
+					_ = tcs.TrySetException(ex);
+				}
+			}
 		}
+
+		return Ok(txsCompletionSources);
 	}
 
 	/// <summary>

WalletWasabi/Cache/IdempotencyRequestCache.cs

@@ -1,7 +1,6 @@
 using Microsoft.Extensions.Caching.Memory;
 using System.Threading;
 using System.Threading.Tasks;
-using WalletWasabi.Extensions;
 
 namespace WalletWasabi.Cache;
 
@@ -28,6 +27,29 @@ public IdempotencyRequestCache(IMemoryCache cache)
 	/// <remarks>Guarded by <see cref="ResponseCacheLock"/>.</remarks>
 	private IMemoryCache ResponseCache { get; }
 
+	/// <summary>
+	/// Tries to add the cache key to cache to avoid other callers to add such a key in parallel.
+	/// </summary>
+	/// <returns><c>true</c> if the key was added to the cache, <c>false</c> otherwise.</returns>
+	/// <remarks>Caller is responsible to ALWAYS set a result to <paramref name="responseTcs"/> even if an exception is thrown.</remarks>
+	public bool TryAddKey<TRequest, TResponse>(TRequest cacheKey, MemoryCacheEntryOptions options, out TaskCompletionSource<TResponse> responseTcs)
+		where TRequest : notnull
+	{
+		lock (ResponseCacheLock)
+		{
+			if (!ResponseCache.TryGetValue(cacheKey, out TaskCompletionSource<TResponse>? tcs))
+			{
+				responseTcs = new();
+				ResponseCache.Set(cacheKey, responseTcs, options);
+
+				return true;
+			}
+
+			responseTcs = tcs!;
+			return false;
+		}
+	}
+
 	/// <typeparam name="TRequest">
 	/// <see langword="record"/>s are preferred as <see cref="object.GetHashCode"/>
 	/// and <see cref="object.Equals(object?)"/> are generated for <see langword="record"/> types automatically.
@@ -47,18 +69,7 @@ public IdempotencyRequestCache(IMemoryCache cache)
 	public async Task<TResponse> GetCachedResponseAsync<TRequest, TResponse>(TRequest request, ProcessRequestDelegateAsync<TRequest, TResponse> action, MemoryCacheEntryOptions options, CancellationToken cancellationToken)
 		where TRequest : notnull
 	{
-		bool callAction = false;
-		TaskCompletionSource<TResponse>? responseTcs;
-
-		lock (ResponseCacheLock)
-		{
-			if (!ResponseCache.TryGetValue(request, out responseTcs))
-			{
-				callAction = true;
-				responseTcs = new();
-				ResponseCache.Set(request, responseTcs, options);
-			}
-		}
+		bool callAction = TryAddKey(request, options, out TaskCompletionSource<TResponse>? responseTcs);
 
 		if (callAction)
 		{