Merge pull request #15878 from peppy/beatmap-lookup-cache

Cache beatmap metadata lookups used by multiplayer

osu.Game/Database/BeatmapLookupCache.cs

@@ -0,0 +1,149 @@
+// Copyright (c) ppy Pty Ltd <contact@ppy.sh>. Licensed under the MIT Licence.
+// See the LICENCE file in the repository root for full licence text.
+
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using JetBrains.Annotations;
+using osu.Framework.Allocation;
+using osu.Game.Online.API;
+using osu.Game.Online.API.Requests;
+using osu.Game.Online.API.Requests.Responses;
+
+namespace osu.Game.Database
+{
+    // This class is based on `UserLookupCache` which is well tested.
+    // If modifications are to be made here, a base abstract implementation should likely be created and shared between the two.
+    public class BeatmapLookupCache : MemoryCachingComponent<int, APIBeatmap>
+    {
+        [Resolved]
+        private IAPIProvider api { get; set; }
+
+        /// <summary>
+        /// Perform an API lookup on the specified beatmap, populating a <see cref="APIBeatmap"/> model.
+        /// </summary>
+        /// <param name="beatmapId">The beatmap to lookup.</param>
+        /// <param name="token">An optional cancellation token.</param>
+        /// <returns>The populated beatmap, or null if the beatmap does not exist or the request could not be satisfied.</returns>
+        [ItemCanBeNull]
+        public Task<APIBeatmap> GetBeatmapAsync(int beatmapId, CancellationToken token = default) => GetAsync(beatmapId, token);
+
+        /// <summary>
+        /// Perform an API lookup on the specified beatmaps, populating a <see cref="APIBeatmap"/> model.
+        /// </summary>
+        /// <param name="beatmapIds">The beatmaps to lookup.</param>
+        /// <param name="token">An optional cancellation token.</param>
+        /// <returns>The populated beatmaps. May include null results for failed retrievals.</returns>
+        public Task<APIBeatmap[]> GetBeatmapsAsync(int[] beatmapIds, CancellationToken token = default)
+        {
+            var beatmapLookupTasks = new List<Task<APIBeatmap>>();
+
+            foreach (int u in beatmapIds)
+            {
+                beatmapLookupTasks.Add(GetBeatmapAsync(u, token).ContinueWith(task =>
+                {
+                    if (!task.IsCompletedSuccessfully)
+                        return null;
+
+                    return task.Result;
+                }, token));
+            }
+
+            return Task.WhenAll(beatmapLookupTasks);
+        }
+
+        protected override async Task<APIBeatmap> ComputeValueAsync(int lookup, CancellationToken token = default)
+            => await queryBeatmap(lookup).ConfigureAwait(false);
+
+        private readonly Queue<(int id, TaskCompletionSource<APIBeatmap>)> pendingBeatmapTasks = new Queue<(int, TaskCompletionSource<APIBeatmap>)>();
+        private Task pendingRequestTask;
+        private readonly object taskAssignmentLock = new object();
+
+        private Task<APIBeatmap> queryBeatmap(int beatmapId)
+        {
+            lock (taskAssignmentLock)
+            {
+                var tcs = new TaskCompletionSource<APIBeatmap>();
+
+                // Add to the queue.
+                pendingBeatmapTasks.Enqueue((beatmapId, tcs));
+
+                // Create a request task if there's not already one.
+                if (pendingRequestTask == null)
+                    createNewTask();
+
+                return tcs.Task;
+            }
+        }
+
+        private void performLookup()
+        {
+            // contains at most 50 unique beatmap IDs from beatmapTasks, which is used to perform the lookup.
+            var beatmapTasks = new Dictionary<int, List<TaskCompletionSource<APIBeatmap>>>();
+
+            // Grab at most 50 unique beatmap IDs from the queue.
+            lock (taskAssignmentLock)
+            {
+                while (pendingBeatmapTasks.Count > 0 && beatmapTasks.Count < 50)
+                {
+                    (int id, TaskCompletionSource<APIBeatmap> task) next = pendingBeatmapTasks.Dequeue();
+
+                    // Perform a secondary check for existence, in case the beatmap was queried in a previous batch.
+                    if (CheckExists(next.id, out var existing))
+                        next.task.SetResult(existing);
+                    else
+                    {
+                        if (beatmapTasks.TryGetValue(next.id, out var tasks))
+                            tasks.Add(next.task);
+                        else
+                            beatmapTasks[next.id] = new List<TaskCompletionSource<APIBeatmap>> { next.task };
+                    }
+                }
+            }
+
+            if (beatmapTasks.Count == 0)
+                return;
+
+            // Query the beatmaps.
+            var request = new GetBeatmapsRequest(beatmapTasks.Keys.ToArray());
+
+            // rather than queueing, we maintain our own single-threaded request stream.
+            // todo: we probably want retry logic here.
+            api.Perform(request);
+
+            // Create a new request task if there's still more beatmaps to query.
+            lock (taskAssignmentLock)
+            {
+                pendingRequestTask = null;
+                if (pendingBeatmapTasks.Count > 0)
+                    createNewTask();
+            }
+
+            List<APIBeatmap> foundBeatmaps = request.Response?.Beatmaps;
+
+            if (foundBeatmaps != null)
+            {
+                foreach (var beatmap in foundBeatmaps)
+                {
+                    if (beatmapTasks.TryGetValue(beatmap.OnlineID, out var tasks))
+                    {
+                        foreach (var task in tasks)
+                            task.SetResult(beatmap);
+
+                        beatmapTasks.Remove(beatmap.OnlineID);
+                    }
+                }
+            }
+
+            // if any tasks remain which were not satisfied, return null.
+            foreach (var tasks in beatmapTasks.Values)
+            {
+                foreach (var task in tasks)
+                    task.SetResult(null);
+            }
+        }
+
+        private void createNewTask() => pendingRequestTask = Task.Run(performLookup);
+    }
+}

osu.Game/Database/UserLookupCache.cs

@@ -100,6 +100,9 @@ namespace osu.Game.Database
                 }
             }
 
+            if (userTasks.Count == 0)
+                return;
+
             // Query the users.
             var request = new GetUsersRequest(userTasks.Keys.ToArray());
 

osu.Game/Online/API/Requests/GetBeatmapsRequest.cs

@@ -0,0 +1,24 @@
+// Copyright (c) ppy Pty Ltd <contact@ppy.sh>. Licensed under the MIT Licence.
+// See the LICENCE file in the repository root for full licence text.
+
+using System;
+
+namespace osu.Game.Online.API.Requests
+{
+    public class GetBeatmapsRequest : APIRequest<GetBeatmapsResponse>
+    {
+        private readonly int[] beatmapIds;
+
+        private const int max_ids_per_request = 50;
+
+        public GetBeatmapsRequest(int[] beatmapIds)
+        {
+            if (beatmapIds.Length > max_ids_per_request)
+                throw new ArgumentException($"{nameof(GetBeatmapsRequest)} calls only support up to {max_ids_per_request} IDs at once");
+
+            this.beatmapIds = beatmapIds;
+        }
+
+        protected override string Target => "beatmaps/?ids[]=" + string.Join("&ids[]=", beatmapIds);
+    }
+}

osu.Game/Online/API/Requests/GetBeatmapsResponse.cs

@@ -0,0 +1,15 @@
+// Copyright (c) ppy Pty Ltd <contact@ppy.sh>. Licensed under the MIT Licence.
+// See the LICENCE file in the repository root for full licence text.
+
+using System.Collections.Generic;
+using Newtonsoft.Json;
+using osu.Game.Online.API.Requests.Responses;
+
+namespace osu.Game.Online.API.Requests
+{
+    public class GetBeatmapsResponse : ResponseWithCursor
+    {
+        [JsonProperty("beatmaps")]
+        public List<APIBeatmap> Beatmaps;
+    }
+}

osu.Game/Online/Multiplayer/MultiplayerClient.cs

@@ -703,15 +703,7 @@ namespace osu.Game.Online.Multiplayer
 
         private async Task<PlaylistItem> createPlaylistItem(MultiplayerPlaylistItem item)
         {
-            var set = await GetOnlineBeatmapSet(item.BeatmapID).ConfigureAwait(false);
-
-            // The incoming response is deserialised without circular reference handling currently.
-            // Because we require using metadata from this instance, populate the nested beatmaps' sets manually here.
-            foreach (var b in set.Beatmaps)
-                b.BeatmapSet = set;
-
-            var beatmap = set.Beatmaps.Single(b => b.OnlineID == item.BeatmapID);
-            beatmap.Checksum = item.BeatmapChecksum;
+            var apiBeatmap = await GetAPIBeatmap(item.BeatmapID).ConfigureAwait(false);
 
             var ruleset = Rulesets.GetRuleset(item.RulesetID);
             var rulesetInstance = ruleset.CreateInstance();
@@ -720,7 +712,7 @@ namespace osu.Game.Online.Multiplayer
             {
                 ID = item.ID,
                 OwnerID = item.OwnerID,
-                Beatmap = { Value = beatmap },
+                Beatmap = { Value = apiBeatmap },
                 Ruleset = { Value = ruleset },
                 Expired = item.Expired
             };
@@ -732,12 +724,12 @@ namespace osu.Game.Online.Multiplayer
         }
 
         /// <summary>
-        /// Retrieves a <see cref="APIBeatmapSet"/> from an online source.
+        /// Retrieves a <see cref="APIBeatmap"/> from an online source.
         /// </summary>
-        /// <param name="beatmapId">The beatmap set ID.</param>
+        /// <param name="beatmapId">The beatmap ID.</param>
         /// <param name="cancellationToken">A token to cancel the request.</param>
-        /// <returns>The <see cref="APIBeatmapSet"/> retrieval task.</returns>
-        protected abstract Task<APIBeatmapSet> GetOnlineBeatmapSet(int beatmapId, CancellationToken cancellationToken = default);
+        /// <returns>The <see cref="APIBeatmap"/> retrieval task.</returns>
+        protected abstract Task<APIBeatmap> GetAPIBeatmap(int beatmapId, CancellationToken cancellationToken = default);
 
         /// <summary>
         /// For the provided user ID, update whether the user is included in <see cref="CurrentMatchPlayingUserIds"/>.

osu.Game/Online/Multiplayer/OnlineMultiplayerClient.cs

@@ -9,8 +9,8 @@ using System.Threading.Tasks;
 using Microsoft.AspNetCore.SignalR.Client;
 using osu.Framework.Allocation;
 using osu.Framework.Bindables;
+using osu.Game.Database;
 using osu.Game.Online.API;
-using osu.Game.Online.API.Requests;
 using osu.Game.Online.API.Requests.Responses;
 using osu.Game.Online.Rooms;
 
@@ -29,6 +29,9 @@ namespace osu.Game.Online.Multiplayer
 
         private HubConnection? connection => connector?.CurrentConnection;
 
+        [Resolved]
+        private BeatmapLookupCache beatmapLookupCache { get; set; } = null!;
+
         public OnlineMultiplayerClient(EndpointConfiguration endpoints)
         {
             endpoint = endpoints.MultiplayerEndpointUrl;
@@ -159,27 +162,9 @@ namespace osu.Game.Online.Multiplayer
             return connection.InvokeAsync(nameof(IMultiplayerServer.AddPlaylistItem), item);
         }
 
-        protected override Task<APIBeatmapSet> GetOnlineBeatmapSet(int beatmapId, CancellationToken cancellationToken = default)
+        protected override Task<APIBeatmap> GetAPIBeatmap(int beatmapId, CancellationToken cancellationToken = default)
         {
-            var tcs = new TaskCompletionSource<APIBeatmapSet>();
-            var req = new GetBeatmapSetRequest(beatmapId, BeatmapSetLookupType.BeatmapId);
-
-            req.Success += res =>
-            {
-                if (cancellationToken.IsCancellationRequested)
-                {
-                    tcs.SetCanceled();
-                    return;
-                }
-
-                tcs.SetResult(res);
-            };
-
-            req.Failure += e => tcs.SetException(e);
-
-            API.Queue(req);
-
-            return tcs.Task;
+            return beatmapLookupCache.GetBeatmapAsync(beatmapId, cancellationToken);
         }
 
         protected override void Dispose(bool isDisposing)

osu.Game/OsuGameBase.cs

@@ -142,6 +142,7 @@ namespace osu.Game
         private BeatmapDifficultyCache difficultyCache;
 
         private UserLookupCache userCache;
+        private BeatmapLookupCache beatmapCache;
 
         private FileStore fileStore;
 
@@ -265,6 +266,9 @@ namespace osu.Game
             dependencies.Cache(userCache = new UserLookupCache());
             AddInternal(userCache);
 
+            dependencies.Cache(beatmapCache = new BeatmapLookupCache());
+            AddInternal(beatmapCache);
+
             var scorePerformanceManager = new ScorePerformanceCache();
             dependencies.Cache(scorePerformanceManager);
             AddInternal(scorePerformanceManager);

osu.Game/Tests/Visual/Multiplayer/TestMultiplayerClient.cs

@@ -336,7 +336,7 @@ namespace osu.Game.Tests.Visual.Multiplayer
 
         public override Task AddPlaylistItem(MultiplayerPlaylistItem item) => AddUserPlaylistItem(api.LocalUser.Value.OnlineID, item);
 
-        protected override Task<APIBeatmapSet> GetOnlineBeatmapSet(int beatmapId, CancellationToken cancellationToken = default)
+        protected override Task<APIBeatmap> GetAPIBeatmap(int beatmapId, CancellationToken cancellationToken = default)
         {
             IBeatmapSetInfo? set = roomManager.ServerSideRooms.SelectMany(r => r.Playlist)
                                               .FirstOrDefault(p => p.BeatmapID == beatmapId)?.Beatmap.Value.BeatmapSet
@@ -345,13 +345,12 @@ namespace osu.Game.Tests.Visual.Multiplayer
             if (set == null)
                 throw new InvalidOperationException("Beatmap not found.");
 
-            var apiSet = new APIBeatmapSet
+            return Task.FromResult(new APIBeatmap
             {
-                OnlineID = set.OnlineID,
-                Beatmaps = set.Beatmaps.Select(b => new APIBeatmap { OnlineID = b.OnlineID }).ToArray(),
-            };
-
-            return Task.FromResult(apiSet);
+                BeatmapSet = new APIBeatmapSet { OnlineID = set.OnlineID },
+                OnlineID = beatmapId,
+                Checksum = set.Beatmaps.First(b => b.OnlineID == beatmapId).MD5Hash
+            });
         }
 
         private async Task changeMatchType(MatchType type)