From ac1534cda2fb51afad43bcc4e4f4fe945a232cde Mon Sep 17 00:00:00 2001
From: Dean Herbert <pe@ppy.sh>
Date: Wed, 28 Apr 2021 15:54:40 +0900
Subject: [PATCH 1/2] Add test covering existing button actually changing to
 `LocallyAvailable` state

---
 osu.Game.Tests/Visual/Online/TestSceneDirectDownloadButton.cs | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/osu.Game.Tests/Visual/Online/TestSceneDirectDownloadButton.cs b/osu.Game.Tests/Visual/Online/TestSceneDirectDownloadButton.cs
index 0c199bfb62..3fc894da0d 100644
--- a/osu.Game.Tests/Visual/Online/TestSceneDirectDownloadButton.cs
+++ b/osu.Game.Tests/Visual/Online/TestSceneDirectDownloadButton.cs
@@ -43,7 +43,10 @@ namespace osu.Game.Tests.Visual.Online
             createButtonWithBeatmap(createSoleily());
             AddAssert("button state not downloaded", () => downloadButton.DownloadState == DownloadState.NotDownloaded);
             AddStep("import soleily", () => beatmaps.Import(TestResources.GetQuickTestBeatmapForImport()));
+
             AddUntilStep("wait for beatmap import", () => beatmaps.GetAllUsableBeatmapSets().Any(b => b.OnlineBeatmapSetID == 241526));
+            AddAssert("button state downloaded", () => downloadButton.DownloadState == DownloadState.LocallyAvailable);
+
             createButtonWithBeatmap(createSoleily());
             AddAssert("button state downloaded", () => downloadButton.DownloadState == DownloadState.LocallyAvailable);
             ensureSoleilyRemoved();

From 05e3a73a7d24acf00431405b96234498f830ddcb Mon Sep 17 00:00:00 2001
From: Dean Herbert <pe@ppy.sh>
Date: Wed, 28 Apr 2021 15:54:58 +0900
Subject: [PATCH 2/2] Fix import cancellation not correctly being forwarded to
 import notification

---
 osu.Game/Database/ArchiveModelManager.cs | 49 +++++++++++++++---------
 1 file changed, 30 insertions(+), 19 deletions(-)

diff --git a/osu.Game/Database/ArchiveModelManager.cs b/osu.Game/Database/ArchiveModelManager.cs
index 6719351530..dbeaebb1cd 100644
--- a/osu.Game/Database/ArchiveModelManager.cs
+++ b/osu.Game/Database/ArchiveModelManager.cs
@@ -156,33 +156,44 @@ namespace osu.Game.Database
 
             bool isLowPriorityImport = tasks.Length > low_priority_import_batch_size;
 
-            await Task.WhenAll(tasks.Select(async task =>
+            try
             {
-                notification.CancellationToken.ThrowIfCancellationRequested();
-
-                try
+                await Task.WhenAll(tasks.Select(async task =>
                 {
-                    var model = await Import(task, isLowPriorityImport, notification.CancellationToken).ConfigureAwait(false);
+                    notification.CancellationToken.ThrowIfCancellationRequested();
 
-                    lock (imported)
+                    try
                     {
-                        if (model != null)
-                            imported.Add(model);
-                        current++;
+                        var model = await Import(task, isLowPriorityImport, notification.CancellationToken).ConfigureAwait(false);
 
-                        notification.Text = $"Imported {current} of {tasks.Length} {HumanisedModelName}s";
-                        notification.Progress = (float)current / tasks.Length;
+                        lock (imported)
+                        {
+                            if (model != null)
+                                imported.Add(model);
+                            current++;
+
+                            notification.Text = $"Imported {current} of {tasks.Length} {HumanisedModelName}s";
+                            notification.Progress = (float)current / tasks.Length;
+                        }
                     }
-                }
-                catch (TaskCanceledException)
+                    catch (TaskCanceledException)
+                    {
+                        throw;
+                    }
+                    catch (Exception e)
+                    {
+                        Logger.Error(e, $@"Could not import ({task})", LoggingTarget.Database);
+                    }
+                })).ConfigureAwait(false);
+            }
+            catch (OperationCanceledException)
+            {
+                if (imported.Count == 0)
                 {
-                    throw;
+                    notification.State = ProgressNotificationState.Cancelled;
+                    return imported;
                 }
-                catch (Exception e)
-                {
-                    Logger.Error(e, $@"Could not import ({task})", LoggingTarget.Database);
-                }
-            })).ConfigureAwait(false);
+            }
 
             if (imported.Count == 0)
             {