Merge pull request #227 from bdach/queue-processor-scoping

Allow specifying queue name & scope of calculation in queue processor
ppy · Dec 13, 2023 · 6fbd071 · 6fbd071
2 parents f6f0e85 + 3a07737
commit 6fbd071
Show file tree

Hide file tree

Showing 5 changed files with 59 additions and 30 deletions.
diff --git a/osu.Server.DifficultyCalculator/Commands/CalculatorCommand.cs b/osu.Server.DifficultyCalculator/Commands/CalculatorCommand.cs
@@ -42,7 +42,7 @@ public abstract class CalculatorCommand : CommandBase
         public bool DryRun { get; set; }
 
         [Option(CommandOptionType.SingleValue, Template = "--processing-mode", Description = "The mode in which to process beatmaps.")]
-        public ProcessingModes ProcessingMode { get; set; } = ProcessingModes.All;
+        public ProcessingMode ProcessingMode { get; set; } = ProcessingMode.All;
 
         private int[] threadBeatmapIds = null!;
         private IReporter reporter = null!;
@@ -92,20 +92,7 @@ public void OnExecute(CommandLineApplication app, IConsole console)
                             // ensure the correct online id is set
                             beatmap.BeatmapInfo.OnlineID = beatmapId;
 
-                            switch (ProcessingMode)
-                            {
-                                case ProcessingModes.All:
-                                    calc.ProcessAll(beatmap);
-                                    break;
-
-                                case ProcessingModes.Difficulty:
-                                    calc.ProcessDifficulty(beatmap);
-                                    break;
-
-                                case ProcessingModes.ScoreAttributes:
-                                    calc.ProcessLegacyAttributes(beatmap);
-                                    break;
-                            }
+                            calc.Process(beatmap, ProcessingMode);
 
                             reporter.Verbose($"Difficulty updated for beatmap {beatmapId}.");
                         }
@@ -168,12 +155,5 @@ protected string CombineSqlConditions(params string?[] conditions)
         }
 
         protected abstract IEnumerable<int> GetBeatmaps();
-
-        public enum ProcessingModes
-        {
-            All,
-            Difficulty,
-            ScoreAttributes
-        }
     }
 }
diff --git a/osu.Server.DifficultyCalculator/Commands/ProcessingMode.cs b/osu.Server.DifficultyCalculator/Commands/ProcessingMode.cs
@@ -0,0 +1,12 @@
+// Copyright (c) ppy Pty Ltd <[email protected]>. Licensed under the MIT Licence.
+// See the LICENCE file in the repository root for full licence text.
+
+namespace osu.Server.DifficultyCalculator.Commands
+{
+    public enum ProcessingMode
+    {
+        All,
+        Difficulty,
+        ScoreAttributes
+    }
+}
diff --git a/osu.Server.DifficultyCalculator/ServerDifficultyCalculator.cs b/osu.Server.DifficultyCalculator/ServerDifficultyCalculator.cs
@@ -13,6 +13,7 @@
 using osu.Game.Rulesets;
 using osu.Game.Rulesets.Mods;
 using osu.Game.Rulesets.Scoring.Legacy;
+using osu.Server.DifficultyCalculator.Commands;
 
 namespace osu.Server.DifficultyCalculator
 {
@@ -40,10 +41,26 @@ public ServerDifficultyCalculator(int[]? rulesetIds = null, bool processConverts
             }
         }
 
-        public void ProcessAll(WorkingBeatmap beatmap)
+        public void Process(WorkingBeatmap beatmap, ProcessingMode mode)
         {
-            ProcessDifficulty(beatmap);
-            ProcessLegacyAttributes(beatmap);
+            switch (mode)
+            {
+                case ProcessingMode.All:
+                    ProcessDifficulty(beatmap);
+                    ProcessLegacyAttributes(beatmap);
+                    break;
+
+                case ProcessingMode.Difficulty:
+                    ProcessDifficulty(beatmap);
+                    break;
+
+                case ProcessingMode.ScoreAttributes:
+                    ProcessLegacyAttributes(beatmap);
+                    break;
+
+                default:
+                    throw new ArgumentOutOfRangeException(nameof(mode), mode, "Unsupported processing mode supplied");
+            }
         }
 
         public void ProcessDifficulty(WorkingBeatmap beatmap) => run(beatmap, processDifficulty);

diff --git a/osu.Server.Queues.BeatmapProcessor/BeatmapProcessor.cs b/osu.Server.Queues.BeatmapProcessor/BeatmapProcessor.cs
@@ -3,21 +3,24 @@
 
 using Dapper;
 using osu.Server.DifficultyCalculator;
+using osu.Server.DifficultyCalculator.Commands;
 using osu.Server.QueueProcessor;
 
 namespace osu.Server.Queues.BeatmapProcessor
 {
     internal class BeatmapProcessor : QueueProcessor<BeatmapItem>
     {
+        private readonly ProcessingMode processingMode;
         private readonly ServerDifficultyCalculator calculator;
 
-        public BeatmapProcessor()
+        public BeatmapProcessor(ProcessingMode processingMode, string queueName)
             : base(new QueueConfiguration
             {
-                InputQueueName = "beatmap",
+                InputQueueName = queueName,
                 MaxInFlightItems = 4,
             })
         {
+            this.processingMode = processingMode;
             calculator = new ServerDifficultyCalculator(new[] { 0, 1, 2, 3 });
         }
 
@@ -34,7 +37,7 @@ protected override void ProcessResult(BeatmapItem item)
                     // ensure the correct online id is set
                     working.BeatmapInfo.OnlineID = (int)beatmapId;
 
-                    calculator.ProcessAll(working);
+                    calculator.Process(working, processingMode);
                 }
             }
         }

diff --git a/osu.Server.Queues.BeatmapProcessor/Program.cs b/osu.Server.Queues.BeatmapProcessor/Program.cs
@@ -1,13 +1,30 @@
 // Copyright (c) ppy Pty Ltd <[email protected]>. Licensed under the MIT Licence.
 // See the LICENCE file in the repository root for full licence text.
 
+using JetBrains.Annotations;
+using McMaster.Extensions.CommandLineUtils;
+using osu.Server.DifficultyCalculator.Commands;
+
 namespace osu.Server.Queues.BeatmapProcessor
 {
+    [Command]
     public class Program
     {
-        public static void Main(string[] args)
+        [Argument(0, "mode", "The target mode to process the beatmaps from the queue in.")]
+        [UsedImplicitly]
+        private ProcessingMode processingMode { get; } = ProcessingMode.All;
+
+        [Argument(1, "queue-name", "The name of the queue to watch. The `osu-queue:` prefix must be omitted.")]
+        [UsedImplicitly]
+        private string queueName { get; } = "beatmap";
+
+        public static void Main(string[] args) => CommandLineApplication.Execute<Program>(args);
+
+        [UsedImplicitly]
+        public int OnExecute(CommandLineApplication app)
         {
-            new BeatmapProcessor().Run();
+            new BeatmapProcessor(processingMode, queueName).Run();
+            return 0;
         }
     }
 }