From 5c2e7257394a9a8dcac9bbd787900277f9be7406 Mon Sep 17 00:00:00 2001
From: Alexey Golub <tyrrrz@gmail.com>
Date: Sat, 11 Jan 2020 14:18:04 +0200
Subject: [PATCH] Move partitioning to message renderer facade

---
 .../FacadeMessageRenderer.cs                  | 47 ++++++++++++-------
 .../ExportService.cs                          | 16 ++-----
 2 files changed, 34 insertions(+), 29 deletions(-)

diff --git a/DiscordChatExporter.Core.Rendering/FacadeMessageRenderer.cs b/DiscordChatExporter.Core.Rendering/FacadeMessageRenderer.cs
index 00ad9e1e..22e5033f 100644
--- a/DiscordChatExporter.Core.Rendering/FacadeMessageRenderer.cs
+++ b/DiscordChatExporter.Core.Rendering/FacadeMessageRenderer.cs
@@ -10,16 +10,19 @@ namespace DiscordChatExporter.Core.Rendering
         private readonly string _baseFilePath;
         private readonly ExportFormat _format;
         private readonly RenderContext _context;
+        private readonly int? _partitionLimit;
 
+        private long _renderedMessageCount;
         private int _partitionIndex;
         private TextWriter _writer;
         private IMessageRenderer _innerRenderer;
 
-        public FacadeMessageRenderer(string baseFilePath, ExportFormat format, RenderContext context)
+        public FacadeMessageRenderer(string baseFilePath, ExportFormat format, RenderContext context, int? partitionLimit)
         {
             _baseFilePath = baseFilePath;
             _format = format;
             _context = context;
+            _partitionLimit = partitionLimit;
         }
 
         private void EnsureInnerRendererInitialized()
@@ -35,7 +38,7 @@ namespace DiscordChatExporter.Core.Rendering
             if (!string.IsNullOrWhiteSpace(dirPath))
                 Directory.CreateDirectory(dirPath);
 
-            // Create writer (will be disposed by renderer)
+            // Create writer
             _writer = File.CreateText(filePath);
 
             // Create inner renderer
@@ -61,31 +64,41 @@ namespace DiscordChatExporter.Core.Rendering
             }
         }
 
-        public async Task NextPartitionAsync()
+        private async Task ResetInnerRendererAsync()
         {
-            // Dispose writer and inner renderer
-            await DisposeAsync();
-            _writer = null;
-            _innerRenderer = null;
+            if (_innerRenderer != null)
+            {
+                await _innerRenderer.DisposeAsync();
+                _innerRenderer = null;
+            }
 
-            // Increment partition index
-            _partitionIndex++;
+            if (_writer != null)
+            {
+                await _writer.DisposeAsync();
+                _writer = null;
+            }
         }
 
         public async Task RenderMessageAsync(Message message)
         {
+            // Ensure underlying writer and renderer are initialized
             EnsureInnerRendererInitialized();
+
+            // Render the actual message
             await _innerRenderer.RenderMessageAsync(message);
+
+            // Increment count
+            _renderedMessageCount++;
+
+            // Update partition if necessary
+            if (_partitionLimit != null && _partitionLimit != 0 && _renderedMessageCount % _partitionLimit == 0)
+            {
+                await ResetInnerRendererAsync();
+                _partitionIndex++;
+            }
         }
 
-        public async ValueTask DisposeAsync()
-        {
-            if (_innerRenderer != null)
-                await _innerRenderer.DisposeAsync();
-
-            if (_writer != null)
-                await _writer.DisposeAsync();
-        }
+        public async ValueTask DisposeAsync() => await ResetInnerRendererAsync();
     }
 
     public partial class FacadeMessageRenderer
diff --git a/DiscordChatExporter.Core.Services/ExportService.cs b/DiscordChatExporter.Core.Services/ExportService.cs
index 7e716239..f034cb9a 100644
--- a/DiscordChatExporter.Core.Services/ExportService.cs
+++ b/DiscordChatExporter.Core.Services/ExportService.cs
@@ -38,10 +38,10 @@ namespace DiscordChatExporter.Core.Services
 
             // Create renderer
             var baseFilePath = GetFilePathFromOutputPath(outputPath, format, context);
-            await using var renderer = new FacadeMessageRenderer(baseFilePath, format, context);
+            await using var renderer = new FacadeMessageRenderer(baseFilePath, format, context, partitionLimit);
 
             // Render messages
-            var messageCount = 0L;
+            var renderedAnything = false;
             await foreach (var message in _dataService.GetMessagesAsync(token, channel.Id, after, before, progress))
             {
                 // Add encountered users to the list of mentionable users
@@ -50,19 +50,11 @@ namespace DiscordChatExporter.Core.Services
 
                 // Render message
                 await renderer.RenderMessageAsync(message);
-                messageCount++;
-
-                // Trigger next partition when needed
-                if (partitionLimit != null &&
-                    partitionLimit != 0 &&
-                    messageCount % partitionLimit.Value == 0)
-                {
-                    await renderer.NextPartitionAsync();
-                }
+                renderedAnything = true;
             }
 
             // Throw if no messages were rendered
-            if (messageCount == 0)
+            if (!renderedAnything)
                 throw new DomainException($"Channel [{channel.Name}] contains no messages for specified period");
         }
     }