mirror of
https://github.com/Cog-Creators/Red-DiscordBot.git
synced 2025-11-21 02:16:09 -05:00
do better with loop cleanup (#3245)
* do better with loop cleanup * changelog * remove redundant line * Do this a bit better than the initial pass * Improve windows support Make some other things coroutines to work with improved design * Wish we'd have done this right from the start... * Update deps surrounding this - see bpo-23057 - neccessary for windows users - nice for consistent support channel info / feature availability * dep issue * Fix tests * duplication plugin py version * actually handle this * Reconfigure some checks with codeclimate, disable pylint for now * style * Is my exasperation showing yet? * handle some stupid stuff * meh * dep changelog
This commit is contained in:
Michael H
GitHub
@@ -3,6 +3,7 @@
|
||||
# Discord Version check
|
||||
|
||||
import asyncio
|
||||
import functools
|
||||
import getpass
|
||||
import json
|
||||
import logging
|
||||
@@ -10,7 +11,9 @@ import os
|
||||
import pip
|
||||
import platform
|
||||
import shutil
|
||||
import signal
|
||||
import sys
|
||||
from argparse import Namespace
|
||||
from copy import deepcopy
|
||||
from pathlib import Path
|
||||
|
||||
@@ -24,17 +27,11 @@ from redbot import _update_event_loop_policy, __version__
|
||||
_update_event_loop_policy()
|
||||
|
||||
import redbot.logging
|
||||
from redbot.core.bot import Red, ExitCodes
|
||||
from redbot.core.cog_manager import CogManagerUI
|
||||
from redbot.core.global_checks import init_global_checks
|
||||
from redbot.core.events import init_events
|
||||
from redbot.core.bot import Red
|
||||
from redbot.core.cli import interactive_config, confirm, parse_cli_flags
|
||||
from redbot.core.core_commands import Core, license_info_command
|
||||
from redbot.setup import get_data_dir, get_name, save_config
|
||||
from redbot.core.dev_commands import Dev
|
||||
from redbot.core import __version__, modlog, bank, data_manager, drivers
|
||||
from redbot.core import data_manager, drivers
|
||||
from redbot.core._sharedlibdeprecation import SharedLibImportWarner
|
||||
from signal import SIGTERM
|
||||
|
||||
|
||||
log = logging.getLogger("red.main")
|
||||
@@ -46,16 +43,6 @@ log = logging.getLogger("red.main")
|
||||
#
|
||||
|
||||
|
||||
async def _get_prefix_and_token(red, indict):
|
||||
"""
|
||||
Again, please blame <@269933075037814786> for this.
|
||||
:param indict:
|
||||
:return:
|
||||
"""
|
||||
indict["token"] = await red._config.token()
|
||||
indict["prefix"] = await red._config.prefix()
|
||||
|
||||
|
||||
def _get_instance_names():
|
||||
with data_manager.config_file.open(encoding="utf-8") as fs:
|
||||
data = json.load(fs)
|
||||
@@ -115,7 +102,7 @@ def debug_info():
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
def edit_instance(red, cli_flags):
|
||||
async def edit_instance(red, cli_flags):
|
||||
no_prompt = cli_flags.no_prompt
|
||||
token = cli_flags.token
|
||||
owner = cli_flags.owner
|
||||
@@ -138,8 +125,8 @@ def edit_instance(red, cli_flags):
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
_edit_token(red, token, no_prompt)
|
||||
_edit_owner(red, owner, no_prompt)
|
||||
await _edit_token(red, token, no_prompt)
|
||||
await _edit_owner(red, owner, no_prompt)
|
||||
|
||||
data = deepcopy(data_manager.basic_config)
|
||||
name = _edit_instance_name(old_name, new_name, confirm_overwrite, no_prompt)
|
||||
@@ -150,7 +137,7 @@ def edit_instance(red, cli_flags):
|
||||
save_config(old_name, {}, remove=True)
|
||||
|
||||
|
||||
def _edit_token(red, token, no_prompt):
|
||||
async def _edit_token(red, token, no_prompt):
|
||||
if token:
|
||||
if not len(token) >= 50:
|
||||
print(
|
||||
@@ -158,13 +145,13 @@ def _edit_token(red, token, no_prompt):
|
||||
" Instance's token will remain unchanged.\n"
|
||||
)
|
||||
return
|
||||
red.loop.run_until_complete(red._config.token.set(token))
|
||||
await red._config.token.set(token)
|
||||
elif not no_prompt and confirm("Would you like to change instance's token?", default=False):
|
||||
interactive_config(red, False, True, print_header=False)
|
||||
await interactive_config(red, False, True, print_header=False)
|
||||
print("Token updated.\n")
|
||||
|
||||
|
||||
def _edit_owner(red, owner, no_prompt):
|
||||
async def _edit_owner(red, owner, no_prompt):
|
||||
if owner:
|
||||
if not (15 <= len(str(owner)) <= 21):
|
||||
print(
|
||||
@@ -172,7 +159,7 @@ def _edit_owner(red, owner, no_prompt):
|
||||
" Instance's owner will remain unchanged."
|
||||
)
|
||||
return
|
||||
red.loop.run_until_complete(red._config.owner.set(owner))
|
||||
await red._config.owner.set(owner)
|
||||
elif not no_prompt and confirm("Would you like to change instance's owner?", default=False):
|
||||
print(
|
||||
"Remember:\n"
|
||||
@@ -188,7 +175,7 @@ def _edit_owner(red, owner, no_prompt):
|
||||
print("That doesn't look like a valid Discord user id.")
|
||||
continue
|
||||
owner_id = int(owner_id)
|
||||
red.loop.run_until_complete(red._config.owner.set(owner_id))
|
||||
await red._config.owner.set(owner_id)
|
||||
print("Owner updated.")
|
||||
break
|
||||
else:
|
||||
@@ -259,14 +246,72 @@ def _copy_data(data):
|
||||
return True
|
||||
|
||||
|
||||
async def sigterm_handler(red, log):
|
||||
log.info("SIGTERM received. Quitting...")
|
||||
await red.shutdown(restart=False)
|
||||
async def run_bot(red: Red, cli_flags: Namespace):
|
||||
|
||||
driver_cls = drivers.get_driver_class()
|
||||
|
||||
await driver_cls.initialize(**data_manager.storage_details())
|
||||
|
||||
redbot.logging.init_logging(
|
||||
level=cli_flags.logging_level, location=data_manager.core_data_path() / "logs"
|
||||
)
|
||||
|
||||
log.debug("====Basic Config====")
|
||||
log.debug("Data Path: %s", data_manager._base_data_path())
|
||||
log.debug("Storage Type: %s", data_manager.storage_type())
|
||||
|
||||
if cli_flags.edit:
|
||||
try:
|
||||
edit_instance(red, cli_flags)
|
||||
except (KeyboardInterrupt, EOFError):
|
||||
print("Aborted!")
|
||||
finally:
|
||||
await driver_cls.teardown()
|
||||
sys.exit(0)
|
||||
|
||||
# lib folder has to be in sys.path before trying to load any 3rd-party cog (GH-3061)
|
||||
# We might want to change handling of requirements in Downloader at later date
|
||||
LIB_PATH = data_manager.cog_data_path(raw_name="Downloader") / "lib"
|
||||
LIB_PATH.mkdir(parents=True, exist_ok=True)
|
||||
if str(LIB_PATH) not in sys.path:
|
||||
sys.path.append(str(LIB_PATH))
|
||||
sys.meta_path.insert(0, SharedLibImportWarner())
|
||||
|
||||
if cli_flags.token:
|
||||
token = cli_flags.token
|
||||
else:
|
||||
token = os.environ.get("RED_TOKEN", None)
|
||||
if not token:
|
||||
token = await red._config.token()
|
||||
|
||||
prefix = cli_flags.prefix or await red._config.prefix()
|
||||
|
||||
if not (token and prefix):
|
||||
if cli_flags.no_prompt is False:
|
||||
new_token = await interactive_config(
|
||||
red, token_set=bool(token), prefix_set=bool(prefix)
|
||||
)
|
||||
if new_token:
|
||||
token = new_token
|
||||
else:
|
||||
log.critical("Token and prefix must be set in order to login.")
|
||||
sys.exit(1)
|
||||
|
||||
if cli_flags.dry_run:
|
||||
await red.http.close()
|
||||
sys.exit(0)
|
||||
try:
|
||||
await red.start(token, bot=True, cli_flags=cli_flags)
|
||||
except discord.LoginFailure:
|
||||
log.critical("This token doesn't seem to be valid.")
|
||||
db_token = await red._config.token()
|
||||
if db_token and not cli_flags.no_prompt:
|
||||
if confirm("\nDo you want to reset the token?"):
|
||||
await red._config.token.set("")
|
||||
print("Token has been reset.")
|
||||
|
||||
|
||||
def main():
|
||||
description = "Red V3"
|
||||
cli_flags = parse_cli_flags(sys.argv[1:])
|
||||
def handle_early_exit_flags(cli_flags: Namespace):
|
||||
if cli_flags.list_instances:
|
||||
list_instances()
|
||||
elif cli_flags.version:
|
||||
@@ -278,111 +323,75 @@ def main():
|
||||
elif not cli_flags.instance_name and (not cli_flags.no_instance or cli_flags.edit):
|
||||
print("Error: No instance name was provided!")
|
||||
sys.exit(1)
|
||||
if cli_flags.no_instance:
|
||||
print(
|
||||
"\033[1m"
|
||||
"Warning: The data will be placed in a temporary folder and removed on next system "
|
||||
"reboot."
|
||||
"\033[0m"
|
||||
)
|
||||
cli_flags.instance_name = "temporary_red"
|
||||
data_manager.create_temp_config()
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
data_manager.load_basic_configuration(cli_flags.instance_name)
|
||||
driver_cls = drivers.get_driver_class()
|
||||
loop.run_until_complete(driver_cls.initialize(**data_manager.storage_details()))
|
||||
redbot.logging.init_logging(
|
||||
level=cli_flags.logging_level, location=data_manager.core_data_path() / "logs"
|
||||
)
|
||||
|
||||
log.debug("====Basic Config====")
|
||||
log.debug("Data Path: %s", data_manager._base_data_path())
|
||||
log.debug("Storage Type: %s", data_manager.storage_type())
|
||||
async def shutdown_handler(red, signal_type=None):
|
||||
if signal_type:
|
||||
log.info("%s received. Quitting...", signal_type)
|
||||
exit_code = 0
|
||||
else:
|
||||
log.info("Shutting down from unhandled exception")
|
||||
exit_code = 1
|
||||
await red.logout()
|
||||
await red.loop.shutdown_asyncgens()
|
||||
pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
|
||||
[task.cancel() for task in pending]
|
||||
await asyncio.gather(*pending, loop=red.loop, return_exceptions=True)
|
||||
sys.exit(exit_code)
|
||||
|
||||
red = Red(
|
||||
cli_flags=cli_flags, description=description, dm_help=None, fetch_offline_members=True
|
||||
)
|
||||
loop.run_until_complete(red._maybe_update_config())
|
||||
|
||||
if cli_flags.edit:
|
||||
try:
|
||||
edit_instance(red, cli_flags)
|
||||
except (KeyboardInterrupt, EOFError):
|
||||
print("Aborted!")
|
||||
finally:
|
||||
loop.run_until_complete(driver_cls.teardown())
|
||||
sys.exit(0)
|
||||
def exception_handler(red, loop, context):
|
||||
msg = context.get("exception", context["message"])
|
||||
if isinstance(msg, KeyboardInterrupt):
|
||||
# Windows support is ugly, I'm sorry
|
||||
logging.error("Received KeyboardInterrupt, treating as interrupt")
|
||||
signal_type = signal.SIGINT
|
||||
else:
|
||||
logging.critical("Caught fatal exception: %s", msg)
|
||||
signal_type = None
|
||||
loop.create_task(shutdown_handler(red, signal_type))
|
||||
|
||||
init_global_checks(red)
|
||||
init_events(red, cli_flags)
|
||||
|
||||
# lib folder has to be in sys.path before trying to load any 3rd-party cog (GH-3061)
|
||||
# We might want to change handling of requirements in Downloader at later date
|
||||
LIB_PATH = data_manager.cog_data_path(raw_name="Downloader") / "lib"
|
||||
LIB_PATH.mkdir(parents=True, exist_ok=True)
|
||||
if str(LIB_PATH) not in sys.path:
|
||||
sys.path.append(str(LIB_PATH))
|
||||
sys.meta_path.insert(0, SharedLibImportWarner())
|
||||
|
||||
red.add_cog(Core(red))
|
||||
red.add_cog(CogManagerUI())
|
||||
red.add_command(license_info_command)
|
||||
if cli_flags.dev:
|
||||
red.add_cog(Dev())
|
||||
# noinspection PyProtectedMember
|
||||
loop.run_until_complete(modlog._init(red))
|
||||
# noinspection PyProtectedMember
|
||||
bank._init()
|
||||
|
||||
if os.name == "posix":
|
||||
loop.add_signal_handler(SIGTERM, lambda: asyncio.ensure_future(sigterm_handler(red, log)))
|
||||
tmp_data = {}
|
||||
loop.run_until_complete(_get_prefix_and_token(red, tmp_data))
|
||||
token = os.environ.get("RED_TOKEN", tmp_data["token"])
|
||||
if cli_flags.token:
|
||||
token = cli_flags.token
|
||||
prefix = cli_flags.prefix or tmp_data["prefix"]
|
||||
if not (token and prefix):
|
||||
if cli_flags.no_prompt is False:
|
||||
new_token = interactive_config(red, token_set=bool(token), prefix_set=bool(prefix))
|
||||
if new_token:
|
||||
token = new_token
|
||||
else:
|
||||
log.critical("Token and prefix must be set in order to login.")
|
||||
sys.exit(1)
|
||||
loop.run_until_complete(_get_prefix_and_token(red, tmp_data))
|
||||
|
||||
if cli_flags.dry_run:
|
||||
loop.run_until_complete(red.http.close())
|
||||
sys.exit(0)
|
||||
def main():
|
||||
cli_flags = parse_cli_flags(sys.argv[1:])
|
||||
handle_early_exit_flags(cli_flags)
|
||||
try:
|
||||
loop.run_until_complete(red.start(token, bot=True, cli_flags=cli_flags))
|
||||
except discord.LoginFailure:
|
||||
log.critical("This token doesn't seem to be valid.")
|
||||
db_token = loop.run_until_complete(red._config.token())
|
||||
if db_token and not cli_flags.no_prompt:
|
||||
if confirm("\nDo you want to reset the token?"):
|
||||
loop.run_until_complete(red._config.token.set(""))
|
||||
print("Token has been reset.")
|
||||
except KeyboardInterrupt:
|
||||
log.info("Keyboard interrupt detected. Quitting...")
|
||||
loop.run_until_complete(red.logout())
|
||||
red._shutdown_mode = ExitCodes.SHUTDOWN
|
||||
except Exception as e:
|
||||
log.critical("Fatal exception", exc_info=e)
|
||||
loop.run_until_complete(red.logout())
|
||||
finally:
|
||||
pending = asyncio.Task.all_tasks(loop=red.loop)
|
||||
gathered = asyncio.gather(*pending, loop=red.loop, return_exceptions=True)
|
||||
gathered.cancel()
|
||||
try:
|
||||
loop.run_until_complete(red.rpc.close())
|
||||
except AttributeError:
|
||||
pass
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
sys.exit(red._shutdown_mode.value)
|
||||
if cli_flags.no_instance:
|
||||
print(
|
||||
"\033[1m"
|
||||
"Warning: The data will be placed in a temporary folder and removed on next system "
|
||||
"reboot."
|
||||
"\033[0m"
|
||||
)
|
||||
cli_flags.instance_name = "temporary_red"
|
||||
data_manager.create_temp_config()
|
||||
|
||||
data_manager.load_basic_configuration(cli_flags.instance_name)
|
||||
|
||||
red = Red(
|
||||
cli_flags=cli_flags, description=description, dm_help=None, fetch_offline_members=True
|
||||
)
|
||||
|
||||
if os.name != "nt":
|
||||
# None of this works on windows, and we have to catch KeyboardInterrupt in a global handler!
|
||||
# At least it's not a redundant handler...
|
||||
signals = (signal.SIGHUP, signal.SIGTERM, signal.SIGINT)
|
||||
for s in signals:
|
||||
loop.add_signal_handler(
|
||||
s, lambda s=s: asyncio.create_task(shutdown_handler(red, s))
|
||||
)
|
||||
|
||||
exc_handler = functools.partial(exception_handler, red)
|
||||
loop.set_exception_handler(exc_handler)
|
||||
# We actually can't use asyncio.run and have graceful cleanup on Windows...
|
||||
loop.create_task(run_bot(red, cli_flags))
|
||||
loop.run_forever()
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
description = "Red V3"
|
||||
main()
|
||||
|
||||
Reference in New Issue
Block a user