diff --git a/cartographer/CHANGELOG.md b/cartographer/CHANGELOG.md index 09fd49252..9a4218b3c 100644 --- a/cartographer/CHANGELOG.md +++ b/cartographer/CHANGELOG.md @@ -1,5 +1,12 @@ # Cartographer Changelog +## v2.3.0 + +- **Fix**: Backing up a server with many message attachments could exhaust the host's memory and get the bot OOM-killed. Every attachment was downloaded, base64-encoded, and held on the model until the whole guild had been serialized — roughly 4x the attachment size resident at once, before the JSON document was built on top. The v2.2.0 per-file check (`attachment.size < guild.filesize_limit`) only rejects individual oversized files and puts no ceiling on the total, so a media-heavy server produced a multi-GB backup and a peak RSS several times larger. Restoring such a backup had the same problem in reverse, since the whole file was read into a string and parsed. +- **Change**: Backups are now written as `.zip` archives (`backup.json` plus an `attachments/` folder) instead of a single `.json` file. Attachments stream into the archive as they download and are read back out one at a time, so memory stays flat regardless of how much media a server has. Nothing is skipped — a large backup simply takes longer. Backups are also considerably smaller, since base64 inflated every file by a third and the JSON is now compressed. +- **Fix**: Restored messages came back in reverse order. `channel.history(limit=...)` yields newest-first, and the restore replays the stored list in order, so a channel was rebuilt backwards. Messages are now stored oldest-first; the same most-recent N messages are still the ones backed up. +- **Note**: Existing `.json` backups are still readable and restorable; `load_backup` picks the format by extension. New backups are always `.zip`. + ## v2.2.0 - **New**: Message attachments are now backed up and restored. Files are stored base64 inside the backup (`FileBackup`), capped at the guild's upload size limit so oversized files can't balloon the backup or fail re-upload. diff --git a/cartographer/common/formatting.py b/cartographer/common/formatting.py index ef3a3a764..2cdc67ed0 100644 --- a/cartographer/common/formatting.py +++ b/cartographer/common/formatting.py @@ -5,13 +5,13 @@ from redbot.core.i18n import Translator from redbot.core.utils.chat_formatting import humanize_number -from .serializers import GuildBackup +from .serializers import load_backup _ = Translator("Cartographer", __file__) def backup_str(filepath: Path) -> str: - backup = GuildBackup.model_validate_json(filepath.read_text(encoding="utf-8")) + backup = load_backup(filepath) total_messages = sum(len(channel.messages) for channel in backup.text_channels) voice_messages = sum(len(channel.messages) for channel in backup.voice_channels) txt = _( diff --git a/cartographer/common/models.py b/cartographer/common/models.py index b6e751726..a9b063179 100644 --- a/cartographer/common/models.py +++ b/cartographer/common/models.py @@ -3,6 +3,7 @@ import asyncio import logging import os +import zipfile from datetime import datetime, timedelta from pathlib import Path @@ -11,7 +12,7 @@ from redbot.core.i18n import Translator from . import Base -from .serializers import GuildBackup +from .serializers import BACKUP_MEMBER, AttachmentStore, GuildBackup log = logging.getLogger("red.vrt.cartographer.models") _ = Translator("Cartographer", __file__) @@ -96,25 +97,37 @@ async def backup( backup_emojis: bool = True, backup_stickers: bool = True, ) -> None: - backup_obj = await GuildBackup.serialize( - guild=guild, - limit=limit, - backup_members=backup_members, - backup_roles=backup_roles, - backup_emojis=backup_emojis, - backup_stickers=backup_stickers, - ) - dump = await asyncio.to_thread(backup_obj.model_dump_json) backup_dir = backups_dir / str(guild.id) backup_dir.mkdir(parents=True, exist_ok=True) # Clean the guild name to make it filename safe guild_name = "".join(c for c in guild.name if c.isalnum()) - backup_file = backup_dir / f"{guild_name}_{int(datetime.now().timestamp())}.json" - with open(backup_file, "w", encoding="utf-8") as f: - f.write(dump) - f.flush() - os.fsync(f.fileno()) + backup_file = backup_dir / f"{guild_name}_{int(datetime.now().timestamp())}.zip" + + # The archive is opened *before* serializing so attachments can be streamed into + # it as they download. Holding them on the model until the end is what used to + # push memory into the gigabytes on servers with a lot of media. + try: + with open(backup_file, "wb") as f: + with zipfile.ZipFile(f, "w", zipfile.ZIP_DEFLATED) as archive: + backup_obj = await GuildBackup.serialize( + guild=guild, + limit=limit, + backup_members=backup_members, + backup_roles=backup_roles, + backup_emojis=backup_emojis, + backup_stickers=backup_stickers, + store=AttachmentStore(archive), + ) + dump = await asyncio.to_thread(backup_obj.model_dump_json) + await asyncio.to_thread(archive.writestr, BACKUP_MEMBER, dump) + f.flush() + os.fsync(f.fileno()) + except BaseException: + # Serializing now happens with the file already open, so a failure would + # otherwise leave a truncated archive behind for the menu to trip over. + backup_file.unlink(missing_ok=True) + raise if hasattr(os, "O_DIRECTORY"): fd = os.open(backup_file.parent, os.O_DIRECTORY) diff --git a/cartographer/common/serializers.py b/cartographer/common/serializers.py index bc1b08572..b365fd6a0 100644 --- a/cartographer/common/serializers.py +++ b/cartographer/common/serializers.py @@ -4,8 +4,10 @@ import base64 import logging import typing as t +import zipfile from datetime import datetime, timezone from io import BytesIO, StringIO +from pathlib import Path from time import perf_counter import aiohttp @@ -40,6 +42,10 @@ VOICE = t.Union[discord.VoiceChannel, discord.StageChannel] GuildChannels = t.Union[VOICE, discord.ForumChannel, discord.TextChannel, discord.CategoryChannel] +# Layout of a backup archive +BACKUP_MEMBER = "backup.json" +ATTACHMENT_DIR = "attachments" + class Role(Base): id: int @@ -396,19 +402,50 @@ async def restore( return category +class AttachmentStore: + """Streams message attachments into the backup archive, one at a time. + + Attachments used to be base64-encoded onto the model and kept there until the + entire guild had been serialized, so a media-heavy server held gigabytes in + memory at once (roughly 4x the attachment size, before the JSON document was + even built). Writing each file into the archive as it arrives keeps the + high-water mark at a single attachment no matter how many there are. + """ + + def __init__(self, archive: zipfile.ZipFile) -> None: + self.archive = archive + self.count = 0 + + async def add(self, attachment: discord.Attachment) -> str: + """Download an attachment straight into the archive, returning its member name.""" + self.count += 1 + # Prefixed with a counter so duplicate filenames can't collide + safe_name = Path(attachment.filename).name + member = f"{ATTACHMENT_DIR}/{self.count}_{safe_name}" + data = await attachment.read() + # Attachments are already-compressed media as a rule, so storing beats deflating + await asyncio.to_thread(self.archive.writestr, member, data, zipfile.ZIP_STORED) + return member + + class FileBackup(Base): filename: str - filebytes: str # base64 encoded file + filebytes: str = "" # base64 encoded file, only set by backups made before v2.3.0 + stored_name: str | None = None # member name within the backup archive @classmethod - async def serialize(cls, attachment: discord.Attachment) -> FileBackup: - return cls( - filename=attachment.filename, - filebytes=base64.b64encode(await attachment.read()).decode(), - ) - - async def restore(self) -> discord.File: - return discord.File(BytesIO(base64.b64decode(self.filebytes)), filename=self.filename) + async def serialize(cls, attachment: discord.Attachment, store: AttachmentStore) -> FileBackup: + return cls(filename=attachment.filename, stored_name=await store.add(attachment)) + + async def restore(self, archive: zipfile.ZipFile | None = None) -> discord.File: + if self.stored_name is not None: + if archive is None: + raise ValueError(f"{self.filename} lives in a backup archive, but no archive was opened") + data = await asyncio.to_thread(archive.read, self.stored_name) + else: + # Legacy backup with the file inlined as base64 + data = await asyncio.to_thread(base64.b64decode, self.filebytes) + return discord.File(BytesIO(data), filename=self.filename) class MessageBackup(Base): @@ -421,15 +458,20 @@ class MessageBackup(Base): avatar_url: str @classmethod - async def serialize(cls, message: discord.Message) -> MessageBackup: + async def serialize(cls, message: discord.Message, store: AttachmentStore | None = None) -> MessageBackup: # Files larger than the guild's upload limit can't be re-uploaded on restore anyway max_size = message.guild.filesize_limit if message.guild else 8388608 + files: list[FileBackup] = [] + if store is not None: + for attachment in message.attachments: + if attachment.size < max_size: + files.append(await FileBackup.serialize(attachment, store)) return cls( channel_id=message.channel.id, channel_name=message.channel.name, content=message.content[:2000] if message.content else None, embeds=[i.to_dict() for i in message.embeds], - files=[await FileBackup.serialize(i) for i in message.attachments if i.size < max_size], + files=files, username=message.author.display_name, avatar_url=message.author.display_avatar.url, ) @@ -437,18 +479,26 @@ async def serialize(cls, message: discord.Message) -> MessageBackup: async def embed_objects(self) -> list[discord.Embed]: return [discord.Embed.from_dict(i) for i in self.embeds] - async def attachment_objects(self) -> list[discord.File]: - return [await i.restore() for i in self.files] + async def attachment_objects(self, archive: zipfile.ZipFile | None = None) -> list[discord.File]: + return [await i.restore(archive) for i in self.files] async def restore_channel_messages( - channel: discord.TextChannel | discord.VoiceChannel, messages: list[MessageBackup] + channel: discord.TextChannel | discord.VoiceChannel, + messages: list[MessageBackup], + archive_path: Path | None = None, ) -> None: + # This runs as a detached task, so it owns its own handle on the archive rather + # than borrowing the caller's, which is long gone by the time messages restore. + needs_archive = any(f.stored_name for m in messages for f in m.files) + archive: zipfile.ZipFile | None = None try: + if needs_archive and archive_path is not None and archive_path.suffix == ".zip": + archive = await asyncio.to_thread(zipfile.ZipFile, archive_path) hook = await channel.create_webhook(name=_("Cartographer Restore"), reason=_("Restoring messages from backup")) for message in messages: embeds = await message.embed_objects() - files = await message.attachment_objects() + files = await message.attachment_objects(archive) if not any([embeds, files, message.content]): continue await hook.send( @@ -461,6 +511,9 @@ async def restore_channel_messages( await asyncio.sleep(1) except Exception as e: log.exception("Failed to restore messages for channel %s", channel.name, exc_info=e) + finally: + if archive is not None: + await asyncio.to_thread(archive.close) class TextChannel(ChannelBase): @@ -489,14 +542,19 @@ def is_match(self, channel: discord.TextChannel, check_category: bool = False) - return all(matches) and super().is_match(channel) @classmethod - async def serialize(cls, channel: discord.TextChannel, limit: int = 0) -> TextChannel: + async def serialize( + cls, channel: discord.TextChannel, limit: int = 0, store: AttachmentStore | None = None + ) -> TextChannel: messages: list[MessageBackup] = [] if limit: try: async for message in channel.history(limit=limit): - messages.append(await MessageBackup.serialize(message)) + messages.append(await MessageBackup.serialize(message, store)) except discord.HTTPException as e: log.warning("Failed to fetch messages for text channel %s: %s", channel.name, e) + # history() returns newest first, but restore replays the list in order, + # which put the channel back in reverse. Store oldest first instead. + messages.reverse() return cls( id=channel.id, name=channel.name, @@ -525,6 +583,7 @@ async def restore( missing_overwrites: dict[str, list[str]] | None = None, only_missing: bool = False, restore_category: bool = True, + archive_path: Path | None = None, ) -> discord.TextChannel: existing: discord.TextChannel | None = guild.get_channel(self.id) if not existing: @@ -578,7 +637,7 @@ async def restore( ) self.id = channel.id if self.messages: - asyncio.create_task(restore_channel_messages(channel, self.messages)) + asyncio.create_task(restore_channel_messages(channel, self.messages, archive_path)) return channel @@ -812,14 +871,17 @@ def is_match(self, channel: discord.VoiceChannel, check_category: bool = False) return all(matches) and super().is_match(channel) @classmethod - async def serialize(cls, channel: VOICE, limit: int = 0) -> VoiceChannel: + async def serialize(cls, channel: VOICE, limit: int = 0, store: AttachmentStore | None = None) -> VoiceChannel: messages: list[MessageBackup] = [] if limit: try: async for message in channel.history(limit=limit): - messages.append(await MessageBackup.serialize(message)) + messages.append(await MessageBackup.serialize(message, store)) except discord.HTTPException as e: log.warning("Failed to fetch messages for voice channel %s: %s", channel.name, e) + # history() returns newest first, but restore replays the list in order, + # which put the channel back in reverse. Store oldest first instead. + messages.reverse() kwargs = { "id": channel.id, "name": channel.name, @@ -850,6 +912,7 @@ async def restore( missing_overwrites: dict[str, list[str]] | None = None, only_missing: bool = False, restore_category: bool = True, + archive_path: Path | None = None, ) -> discord.VoiceChannel: existing: discord.VoiceChannel | None = guild.get_channel(self.id) if not existing: @@ -899,7 +962,7 @@ async def restore( await channel.edit(slowmode_delay=self.slowmode_delay) if self.messages: - asyncio.create_task(restore_channel_messages(channel, self.messages)) + asyncio.create_task(restore_channel_messages(channel, self.messages, archive_path)) return channel @@ -1066,6 +1129,7 @@ async def serialize( backup_roles: bool = True, backup_emojis: bool = True, backup_stickers: bool = True, + store: AttachmentStore | None = None, ) -> GuildBackup: banner = await guild.banner.read() if guild.banner else None icon = await guild.icon.read() if guild.icon else None @@ -1088,9 +1152,9 @@ async def serialize( indexes[channel.id] = index index += 1 if isinstance(channel, discord.TextChannel): - text_channels.append(await TextChannel.serialize(channel, limit)) + text_channels.append(await TextChannel.serialize(channel, limit, store)) elif isinstance(channel, (discord.VoiceChannel, discord.StageChannel)): - voice_channels.append(await VoiceChannel.serialize(channel, limit)) + voice_channels.append(await VoiceChannel.serialize(channel, limit, store)) elif isinstance(channel, discord.ForumChannel): forums.append(await ForumChannel.serialize(channel)) else: @@ -1143,6 +1207,7 @@ async def restore( target_guild: discord.Guild, ctx: discord.TextChannel, options: RestoreOptions | None = None, + archive_path: Path | None = None, ) -> str: """Restore a guild backup to a target guild. @@ -1150,6 +1215,7 @@ async def restore( target_guild: The guild to restore to ctx: The channel to send status updates to options: Granular restore options. If None, restores everything with delete_unmatched=True (legacy behavior) + archive_path: Path to the backup archive, needed to pull message attachments back out of it """ # Import here to avoid circular import from .models import RestoreOptions @@ -1433,12 +1499,15 @@ def get_status_embed(description: str, complete: bool = False) -> discord.Embed: # Cant restore forums until community is enabled, try again after settings are restored skipped_forums.append(channel) continue + # Only message-bearing channels need the archive to pull attachments back out + extra = {"archive_path": archive_path} if isinstance(channel, (TextChannel, VoiceChannel)) else {} await channel.restore( target_guild, results, missing_overwrites, only_missing=options.only_missing, restore_category=options.categories, + **extra, ) # ---------------------------- REMAINING SETTINGS ---------------------------- @@ -1550,3 +1619,15 @@ def get_status_embed(description: str, complete: bool = False) -> discord.Embed: await message.edit(embed=get_status_embed(_("Restoration complete!"), complete=True)) return results.getvalue() + + +def load_backup(path: Path) -> GuildBackup: + """Read a backup from disk. + + Backups are zip archives as of v2.3.0; plain .json files written by earlier + versions are still loaded so existing backups keep working. + """ + if path.suffix == ".zip": + with zipfile.ZipFile(path) as archive: + return GuildBackup.model_validate_json(archive.read(BACKUP_MEMBER)) + return GuildBackup.model_validate_json(path.read_text(encoding="utf-8")) diff --git a/cartographer/common/views.py b/cartographer/common/views.py index de43264f2..7c67c40c4 100644 --- a/cartographer/common/views.py +++ b/cartographer/common/views.py @@ -11,7 +11,7 @@ from .formatting import backup_str, humanize_size from .models import DB, GuildSettings, RestoreOptions -from .serializers import GuildBackup +from .serializers import GuildBackup, load_backup log = logging.getLogger("red.vrt.cartographer.views") _ = Translator("Cartographer", __file__) @@ -576,8 +576,7 @@ async def restore(self, interaction: discord.Interaction, button: discord.Button # Load the backup first to show options self.page %= len(self.backups) backup_file = self.backups[self.page] - text = await asyncio.to_thread(backup_file.read_text, encoding="utf-8") - backup: GuildBackup = await asyncio.to_thread(GuildBackup.model_validate_json, text) + backup: GuildBackup = await asyncio.to_thread(load_backup, backup_file) # Show the restore options view options_view = RestoreOptionsView( @@ -599,7 +598,9 @@ async def restore(self, interaction: discord.Interaction, button: discord.Button await interaction.followup.send(txt, ephemeral=True) async with self.ctx.typing(): - results = await backup.restore(self.guild, interaction.channel, options_view.options) + results = await backup.restore( + self.guild, interaction.channel, options_view.options, archive_path=backup_file + ) if results: txt = _("The following errors occurred while restoring the backup") await interaction.channel.send(txt, file=text_to_file(results, "restore_results.txt")) diff --git a/cartographer/main.py b/cartographer/main.py index b000e8810..6e1ca027e 100644 --- a/cartographer/main.py +++ b/cartographer/main.py @@ -13,7 +13,7 @@ from .common.formatting import humanize_size from .common.models import DB -from .common.serializers import GuildBackup +from .common.serializers import load_backup from .common.views import BackupMenu log = logging.getLogger("red.vrt.cartographer") @@ -45,7 +45,7 @@ class Cartographer(commands.Cog): """ __author__ = "[vertyco](https://github.com/vertyco/vrt-cogs)" - __version__ = "2.2.0" + __version__ = "2.3.0" def __init__(self, bot: Red): super().__init__() @@ -263,10 +263,8 @@ async def restore_server_latest(self, ctx: commands.Context, confirm: bool = Fal txt = _("There are no backups for this guild!") return await ctx.send(txt) latest = backup_files[-1] - backup = await asyncio.to_thread( - lambda: GuildBackup.model_validate_json(latest.read_text(encoding="utf-8")) - ) - results = await backup.restore(ctx.guild, ctx.channel) + backup = await asyncio.to_thread(load_backup, latest) + results = await backup.restore(ctx.guild, ctx.channel, archive_path=latest) await ctx.send(_("Server restore is complete!")) if results: txt = _("The following errors occurred while restoring the backup")