Merge pull request #1532 from kuu-26/fix/replyto-fk-violation

fix: FK violation when replying to deleted message (fk_replyto)

Author: SpikeViper

Valour/Server/Workers/PlanetMessageWorker.cs

@@ -1,12 +1,12 @@
-using System.Collections.Concurrent;
+using System.Collections.Concurrent;
 
 namespace Valour.Server.Workers
 {
-    public class PlanetMessageWorker : IHostedService, IDisposable
-    {
-        private readonly IServiceProvider _serviceProvider;
-        private readonly ILogger<PlanetMessageWorker> _logger;
-
+    public class PlanetMessageWorker : IHostedService, IDisposable
+    {
+        private readonly IServiceProvider _serviceProvider;
+        private readonly ILogger<PlanetMessageWorker> _logger;
+
         // A queue of all messages that need to be processed
         private static readonly BlockingCollection<Message> MessageQueue = new(new ConcurrentQueue<Message>());
 
@@ -21,29 +21,29 @@ public class PlanetMessageWorker : IHostedService, IDisposable
 
         // Prevents deleted messages from being staged
         private static readonly ConcurrentDictionary<long, byte> BlockSet = new();
-
-        /// <summary>
-        /// Holds the long-running queue task
+
+        /// <summary>
+        /// Holds the long-running queue task
         /// </summary>
         private static Task _queueTask;
 
         // Timer for executing timed tasks
         private Timer _timer;
         private int _isFlushing;
-        
-        public PlanetMessageWorker(ILogger<PlanetMessageWorker> logger,
-                                   IServiceProvider serviceProvider)
-        {
-            _logger = logger;
-            _serviceProvider = serviceProvider;
-        }
-        
-        public static void AddToQueue(Message message)
-        {
-            if (message.PlanetId is null)
-            {
-                Console.WriteLine("[!!!] Tried to add a message to the planet message queue queue without a planet id");
-                return;
+        
+        public PlanetMessageWorker(ILogger<PlanetMessageWorker> logger,
+                                   IServiceProvider serviceProvider)
+        {
+            _logger = logger;
+            _serviceProvider = serviceProvider;
+        }
+        
+        public static void AddToQueue(Message message)
+        {
+            if (message.PlanetId is null)
+            {
+                Console.WriteLine("[!!!] Tried to add a message to the planet message queue queue without a planet id");
+                return;
             }
 
             QueuedMessages[message.Id] = message;
@@ -72,7 +72,7 @@ public static void RemoveFromQueue(Message message)
                 BlockSet.TryRemove(message.Id, out _);
             }
         }
-
+
         public static Message GetStagedMessage(long messageId)
         {
             StagedMessages.TryGetValue(messageId, out var staged);
@@ -84,28 +84,28 @@ public static Message GetQueuedMessage(long messageId)
             QueuedMessages.TryGetValue(messageId, out var queued);
             return queued;
         }
-        
+        
         public static List<Message> GetStagedMessages(long channelId)
         {
             if (StagedChannelMessages.TryGetValue(channelId, out var stagedQueue))
                 return stagedQueue.ToList();
 
             return new List<Message>();
         }
-        
-        public Task StartAsync(CancellationToken stoppingToken)
-        {
-            _logger.LogInformation("Starting Message Worker");
-
-            // Start the queue task
-            _queueTask = Task.Run(ConsumeMessageQueue, stoppingToken);
-            
-            _timer = new Timer(DoWork, null, TimeSpan.Zero, 
-                TimeSpan.FromSeconds(20));
-
-            return Task.CompletedTask;
-        }
-
+        
+        public Task StartAsync(CancellationToken stoppingToken)
+        {
+            _logger.LogInformation("Starting Message Worker");
+
+            // Start the queue task
+            _queueTask = Task.Run(ConsumeMessageQueue, stoppingToken);
+            
+            _timer = new Timer(DoWork, null, TimeSpan.Zero, 
+                TimeSpan.FromSeconds(20));
+
+            return Task.CompletedTask;
+        }
+
         private async void DoWork(object state)
         {
             if (Interlocked.Exchange(ref _isFlushing, 1) == 1)
@@ -118,8 +118,8 @@ private async void DoWork(object state)
             {
                 // If not, restart it
                 _queueTask = Task.Run(ConsumeMessageQueue);
-                
-                _logger.LogInformation("Planet Message Worker queue task stopped at: {Time}, Restarting queue task", DateTime.UtcNow);
+                
+                _logger.LogInformation("Planet Message Worker queue task stopped at: {Time}, Restarting queue task", DateTime.UtcNow);
             }
 
             // Don't work if there's no staged messages
@@ -129,7 +129,7 @@ private async void DoWork(object state)
             /* Get required services in new scope */
             await using var scope = _serviceProvider.CreateAsyncScope();
             await using var db = scope.ServiceProvider.GetRequiredService<ValourDb>();
-            
+            
             _logger.LogInformation(@"Planet Message Worker running at: {Time}
                                              Queue size: {QueueSize}
                                              Saving {StagedCount} messages to DB", DateTimeOffset.Now, MessageQueue.Count, StagedMessages.Count);
@@ -148,14 +148,14 @@ private async void DoWork(object state)
             {
                 await db.Messages.AddRangeAsync(messages);
                 await db.SaveChangesAsync();
-            } 
-            catch (Exception e)
-            {
-                _logger.LogError(e, "Failed to save messages to database. Falling back to per-message save.");
-                
-                // If we fail to save all messages at once, we can try to save them one by one
-                // We will just dump any messages that fail to save
-                
+            } 
+            catch (Exception e)
+            {
+                _logger.LogError(e, "Failed to save messages to database. Falling back to per-message save.");
+                
+                // If we fail to save all messages at once, we can try to save them one by one
+                // We will just dump any messages that fail to save
+                
                 foreach (var message in messages)
                 {
                     try
@@ -169,8 +169,8 @@ private async void DoWork(object state)
                     {
                         _logger.LogError(ex, "Failed to save message {MessageId} to database. Skipping.", message.Id);
                     }
-                }
-            }
+                }
+            }
             foreach (var staged in stagedSnapshot)
             {
                 StagedMessages.TryRemove(staged.Id, out _);
@@ -183,21 +183,21 @@ private async void DoWork(object state)
                 Volatile.Write(ref _isFlushing, 0);
             }
         }
-
-        /// <summary>
-        /// This task should run forever and consume messages from
-        /// the queue.
-        /// </summary>
-        private async Task ConsumeMessageQueue()
-        {
-            // This scope is long-living, which is usually bad. But it's only used to send events,
-            // and does not insert into the database, so it should be fine.
-            await using var scope = _serviceProvider.CreateAsyncScope();
-            var hubService = scope.ServiceProvider.GetRequiredService<CoreHubService>();
-
-            // This is ONLY READ FROM
-            var dbService = scope.ServiceProvider.GetRequiredService<ValourDb>();
-            
+
+        /// <summary>
+        /// This task should run forever and consume messages from
+        /// the queue.
+        /// </summary>
+        private async Task ConsumeMessageQueue()
+        {
+            // This scope is long-living, which is usually bad. But it's only used to send events,
+            // and does not insert into the database, so it should be fine.
+            await using var scope = _serviceProvider.CreateAsyncScope();
+            var hubService = scope.ServiceProvider.GetRequiredService<CoreHubService>();
+
+            // This is ONLY READ FROM
+            var dbService = scope.ServiceProvider.GetRequiredService<ValourDb>();
+            
             // This is a stream and will run forever
             foreach (var message in MessageQueue.GetConsumingEnumerable())
             {
@@ -208,19 +208,28 @@ private async Task ConsumeMessageQueue()
                     BlockSet.TryRemove(message.Id, out _);
                     continue;
                 }
-
-                message.TimeSent = DateTime.UtcNow;
-                
-                hubService.NotifyChannelStateUpdate(message.PlanetId!.Value, message.ChannelId, message.TimeSent);
-                
-                if (message.ReplyToId is not null && message.ReplyTo is null)
-                {
-                    var replyTo = (await dbService.Messages.FindAsync(message.ReplyToId)).ToModel();
-                    message.ReplyTo = replyTo;
-                }
-                
-                hubService.RelayMessage(message);
-
+
+                message.TimeSent = DateTime.UtcNow;
+                
+                hubService.NotifyChannelStateUpdate(message.PlanetId!.Value, message.ChannelId, message.TimeSent);
+                
+                if (message.ReplyToId is not null && message.ReplyTo is null)
+                {
+                    var replyToDb = await dbService.Messages.FindAsync(message.ReplyToId);
+                    if (replyToDb is not null)
+                    {
+                        message.ReplyTo = replyToDb.ToModel();
+                    }
+                    else
+                    {
+                        // The referenced reply-to message no longer exists (deleted).
+                        // Clear ReplyToId to avoid FK constraint violation (fk_replyto) on insert.
+                        message.ReplyToId = null;
+                    }
+                }
+                
+                hubService.RelayMessage(message);
+
                 // Add message to message staging
                 StagedMessages[message.Id] = message;
 
@@ -248,19 +257,19 @@ private static void RemoveStagedMessageFromChannel(long channelId, long messageI
                 StagedChannelMessages[channelId] = new ConcurrentQueue<Message>(remaining);
             }
         }
-        
-        public Task StopAsync(CancellationToken stoppingToken)
-        {
-            _logger.LogInformation("Message Worker is Stopping");
-
-            _timer?.Change(Timeout.Infinite, 0);
-
-            return Task.CompletedTask;
-        }
-        
-        public void Dispose()
-        {
-            _timer?.Dispose();
-        }
-    }
-}
+        
+        public Task StopAsync(CancellationToken stoppingToken)
+        {
+            _logger.LogInformation("Message Worker is Stopping");
+
+            _timer?.Change(Timeout.Infinite, 0);
+
+            return Task.CompletedTask;
+        }
+        
+        public void Dispose()
+        {
+            _timer?.Dispose();
+        }
+    }
+}