From 6c350db0973c74f0294cc975c2c2b3aa26bf17a1 Mon Sep 17 00:00:00 2001
From: Dean Herbert <pe@ppy.sh>
Date: Mon, 11 May 2020 21:37:07 +0900
Subject: [PATCH] Add connection flushing support

---
 .../NonVisual/CustomDataDirectoryTest.cs           | 14 ++++++--------
 osu.Game/Database/DatabaseContextFactory.cs        |  8 ++++++++
 osu.Game/OsuGameBase.cs                            |  8 ++++++++
 3 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/osu.Game.Tests/NonVisual/CustomDataDirectoryTest.cs b/osu.Game.Tests/NonVisual/CustomDataDirectoryTest.cs
index 4bce5056ce..ed83ff358c 100644
--- a/osu.Game.Tests/NonVisual/CustomDataDirectoryTest.cs
+++ b/osu.Game.Tests/NonVisual/CustomDataDirectoryTest.cs
@@ -137,7 +137,7 @@ namespace osu.Game.Tests.NonVisual
 
                     Assert.That(storage.GetFullPath("."), Is.EqualTo(defaultStorageLocation));
 
-                    (storage as OsuStorage)?.Migrate(customPath);
+                    osu.Migrate(customPath);
 
                     Assert.That(storage.GetFullPath("."), Is.EqualTo(customPath));
 
@@ -170,19 +170,18 @@ namespace osu.Game.Tests.NonVisual
                 try
                 {
                     var osu = loadOsu(host);
-                    var storage = osu.Dependencies.Get<Storage>();
 
                     string customPath2 = $"{customPath}-2";
 
                     const string database_filename = "client.db";
 
-                    Assert.DoesNotThrow(() => (storage as OsuStorage)?.Migrate(customPath));
+                    Assert.DoesNotThrow(() => osu.Migrate(customPath));
                     Assert.That(File.Exists(Path.Combine(customPath, database_filename)));
 
-                    Assert.DoesNotThrow(() => (storage as OsuStorage)?.Migrate(customPath2));
+                    Assert.DoesNotThrow(() => osu.Migrate(customPath2));
                     Assert.That(File.Exists(Path.Combine(customPath2, database_filename)));
 
-                    Assert.DoesNotThrow(() => (storage as OsuStorage)?.Migrate(customPath));
+                    Assert.DoesNotThrow(() => osu.Migrate(customPath));
                     Assert.That(File.Exists(Path.Combine(customPath, database_filename)));
                 }
                 finally
@@ -200,10 +199,9 @@ namespace osu.Game.Tests.NonVisual
                 try
                 {
                     var osu = loadOsu(host);
-                    var storage = osu.Dependencies.Get<Storage>();
 
-                    Assert.DoesNotThrow(() => (storage as OsuStorage)?.Migrate(customPath));
-                    Assert.Throws<InvalidOperationException>(() => (storage as OsuStorage)?.Migrate(customPath));
+                    Assert.DoesNotThrow(() => osu.Migrate(customPath));
+                    Assert.Throws<InvalidOperationException>(() => osu.Migrate(customPath));
                 }
                 finally
                 {
diff --git a/osu.Game/Database/DatabaseContextFactory.cs b/osu.Game/Database/DatabaseContextFactory.cs
index 1ed5fb3268..1cceb59b11 100644
--- a/osu.Game/Database/DatabaseContextFactory.cs
+++ b/osu.Game/Database/DatabaseContextFactory.cs
@@ -160,5 +160,13 @@ namespace osu.Game.Database
                 }
             }
         }
+
+        public void FlushConnections()
+        {
+            foreach (var context in threadContexts.Values)
+                context.Dispose();
+
+            recycleThreadContexts();
+        }
     }
 }
diff --git a/osu.Game/OsuGameBase.cs b/osu.Game/OsuGameBase.cs
index 8fbde67afe..6282f5cb8b 100644
--- a/osu.Game/OsuGameBase.cs
+++ b/osu.Game/OsuGameBase.cs
@@ -328,6 +328,8 @@ namespace osu.Game
         {
             base.Dispose(isDisposing);
             RulesetStore?.Dispose();
+
+            ContextFactory.FlushConnections();
         }
 
         private class OsuUserInputManager : UserInputManager
@@ -355,5 +357,11 @@ namespace osu.Game
                 public override bool ChangeFocusOnClick => false;
             }
         }
+
+        public void Migrate(string path)
+        {
+            ContextFactory.FlushConnections();
+            (Storage as OsuStorage)?.Migrate(path);
+        }
     }
 }