Allow specifying queue name & scope of calculation in queue processor

ppy · Dec 12, 2023 · 3a07737 · 3a07737
1 parent a8f2cc7
commit 3a07737
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 5 deletions.
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;
         }
     }
 }