From 90f1e30ef279329b2a3b2cf63224323a2b717d0c Mon Sep 17 00:00:00 2001
From: Kopper <kopper@w.on-t.work>
Date: Sat, 26 Oct 2024 03:50:32 +0300
Subject: [PATCH] [backend] Fix reply backfill limiting

---
 .../Core/Queues/BackfillQueue.cs              | 41 +++++++++++++------
 .../Core/Services/NoteService.cs              |  2 +
 2 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/Iceshrimp.Backend/Core/Queues/BackfillQueue.cs b/Iceshrimp.Backend/Core/Queues/BackfillQueue.cs
index 792ac10c..9bfb8936 100644
--- a/Iceshrimp.Backend/Core/Queues/BackfillQueue.cs
+++ b/Iceshrimp.Backend/Core/Queues/BackfillQueue.cs
@@ -72,7 +72,12 @@ public class BackfillQueue(int parallelism)
 				continue;
 			}
 
-			logger.LogTrace("Backfilling {collection} (remaining limit {limit})", current.RepliesCollection, backfillLimit);
+			if (--backfillLimit <= 0)
+			{
+				logger.LogDebug("Reached backfill limit");
+				break;
+			}
+			logger.LogTrace("Backfilling collection {collection} (remaining limit {limit})", current.RepliesCollection, backfillLimit);
 			
 			await db.Notes
 			        .Where(n => n.Id == current.Id)
@@ -83,20 +88,32 @@ public class BackfillQueue(int parallelism)
 			                                             .Where(p => p.Id != null)
 														 .WithCancellation(token))
 			{
-				if (--backfillLimit <= 0) 
+				logger.LogTrace("Backfilling note {note} (remaining limit {limit})", asNote.Id, backfillLimit);
+
+				try
 				{
-					logger.LogDebug("Reached backfill limit");
-					toBackfill.Clear();
-					break;
+					var note = await noteSvc.ResolveNoteAsync(asNote.Id!, asNote as ASNote, user, clearHistory: true,
+					                                          forceRefresh: false);
+
+					backfillLimit -= Math.Max(noteSvc.NotesFetched, 1);
+					if (backfillLimit <= 0)
+					{
+						logger.LogDebug("Reached backfill limit");
+						toBackfill.Clear();
+						break;
+					}
+
+					if (note is { UserHost: not null, RepliesCollection: not null, RepliesCount: < MaxRepliesPerNote } &&
+					    note.CreatedAt <= DateTime.UtcNow - cfg.NewNoteDelayTimeSpan &&
+					    (note.RepliesFetchedAt == null ||
+					     note.RepliesFetchedAt <= DateTime.UtcNow - cfg.RefreshAfterTimeSpan))
+					{
+						toBackfill.Enqueue(new BackfillData(note.Id, note.RepliesCollection!));
+					}
 				}
-				
-				var note = await noteSvc.ResolveNoteAsync(asNote.Id!, asNote as ASNote, user, clearHistory: true, forceRefresh: false);
-				
-				if (note is { UserHost: not null, RepliesCollection: not null, RepliesCount: < MaxRepliesPerNote }
-				    && note.CreatedAt <= DateTime.UtcNow - cfg.NewNoteDelayTimeSpan 
-				    && (note.RepliesFetchedAt == null || note.RepliesFetchedAt <= DateTime.UtcNow - cfg.RefreshAfterTimeSpan))
+				catch (Exception e)
 				{
-					toBackfill.Enqueue(new BackfillData(note.Id, note.RepliesCollection!));
+					logger.LogWarning(e, "Failed to backfill {note}", asNote.Id);
 				}
 			}
 		}
diff --git a/Iceshrimp.Backend/Core/Services/NoteService.cs b/Iceshrimp.Backend/Core/Services/NoteService.cs
index 1fa27bed..dee1eb56 100644
--- a/Iceshrimp.Backend/Core/Services/NoteService.cs
+++ b/Iceshrimp.Backend/Core/Services/NoteService.cs
@@ -55,6 +55,8 @@ public class NoteService(
 
 	private readonly List<string> _resolverHistory = [];
 	private          int          _recursionLimit  = DefaultRecursionLimit;
+	internal         int          NotesFetched => DefaultRecursionLimit - _recursionLimit;
+
 
 	public class NoteCreationData
 	{