Skip to content

Commit

Permalink
refacto: start is async
Browse files Browse the repository at this point in the history 
not tested but may speed up start and fix issues when shutting down the bot while it being unresponsive
  • Loading branch information
@EDM115
EDM115 committed Dec 9, 2024
1 parent f803e5a commit 7aef751
Show file tree
Hide file tree
Showing 2 changed files with 55 additions and 78 deletions.
95 changes: 43 additions & 52 deletions unzipbot/__main__.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
import asyncio
import os
import shutil
import signal
import time

Expand All @@ -13,7 +12,7 @@
from .helpers.start import (
check_logs,
dl_thumbs,
removal,
remove_expired_tasks,
set_boot_time,
start_cron_jobs,
)
Expand All @@ -23,36 +22,14 @@
messages = Messages(lang_fetcher=get_lang)


def recurse_delete(dir):
for items in os.listdir(dir):
if os.path.isdir(os.path.join(dir, items)):
shutil.rmtree(os.path.join(dir, items), ignore_errors=True)
else:
os.remove(os.path.join(dir, items))


def handler_stop_signals(signum, frame):
LOGGER.info(
messages.get(
"main",
"RECEIVED_STOP_SIGNAL",
None,
signal.Signals(signum).name,
signum,
frame,
)
)
shutdown_bot()


signal.signal(signal.SIGINT, handler_stop_signals)
signal.signal(signal.SIGTERM, handler_stop_signals)


async def async_shutdown_bot():
stoptime = time.strftime("%Y/%m/%d - %H:%M:%S")
LOGGER.info(messages.get("main", "STOP_TXT", None, stoptime))

tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
[task.cancel() for task in tasks]
await asyncio.gather(*tasks, return_exceptions=True)

try:
await unzipbot_client.send_message(
chat_id=Config.LOGS_CHANNEL,
Expand All @@ -71,21 +48,33 @@ async def async_shutdown_bot():
except Exception as e:
LOGGER.error(messages.get("main", "ERROR_SHUTDOWN_MSG", None, e))
finally:
await unzipbot_client.stop()
LOGGER.info(messages.get("main", "BOT_STOPPED"))
await unzipbot_client.stop(block=False)


def shutdown_bot():
try:
loop = asyncio.get_running_loop()
loop.create_task(async_shutdown_bot())
except RuntimeError:
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
loop.run_until_complete(async_shutdown_bot())
def handle_stop_signals(signum, frame):
LOGGER.info(
messages.get(
"main",
"RECEIVED_STOP_SIGNAL",
None,
signal.Signals(signum).name,
signum,
frame,
)
)
loop = asyncio.get_event_loop()
loop.create_task(async_shutdown_bot())


if __name__ == "__main__":
def setup_signal_handlers():
loop = asyncio.get_event_loop()

for sig in (signal.SIGINT, signal.SIGTERM):
loop.add_signal_handler(sig, lambda s=sig: handle_stop_signals(s, None))


async def main():
try:
os.makedirs(Config.DOWNLOAD_LOCATION, exist_ok=True)
os.makedirs(Config.THUMB_LOCATION, exist_ok=True)
Expand All @@ -97,40 +86,42 @@ def shutdown_bot():
lock_f.close()

LOGGER.info(messages.get("main", "STARTING_BOT"))
unzipbot_client.start()
await unzipbot_client.start()
starttime = time.strftime("%Y/%m/%d - %H:%M:%S")
unzipbot_client.send_message(
await unzipbot_client.send_message(
chat_id=Config.LOGS_CHANNEL,
text=messages.get("main", "START_TXT", None, starttime),
)
set_boot_time()
await set_boot_time()
LOGGER.info(messages.get("main", "CHECK_LOG"))

if check_logs():
if await check_logs():
LOGGER.info(messages.get("main", "LOG_CHECKED"))
removal(True)
dl_thumbs()
start_cron_jobs()
# clean previous downloads on volumes
recurse_delete(Config.DOWNLOAD_LOCATION)
os.makedirs(Config.DOWNLOAD_LOCATION, exist_ok=True)
setup_signal_handlers()
await remove_expired_tasks(True)
await dl_thumbs()
await start_cron_jobs()
os.remove(Config.LOCKFILE)
LOGGER.info(messages.get("main", "BOT_RUNNING"))
idle()
await idle()
else:
try:
unzipbot_client.send_message(
await unzipbot_client.send_message(
chat_id=Config.BOT_OWNER,
text=messages.get("main", "WRONG_LOG", None, Config.LOGS_CHANNEL),
)
except:
pass

os.remove(Config.LOCKFILE)
shutdown_bot()
await async_shutdown_bot()
except Exception as e:
LOGGER.error(messages.get("main", "ERROR_MAIN_LOOP", None, e))
finally:
if os.path.exists(Config.LOCKFILE):
os.remove(Config.LOCKFILE)
shutdown_bot()
await async_shutdown_bot()


if __name__ == "__main__":
unzipbot_client.run(main())
38 changes: 12 additions & 26 deletions unzipbot/helpers/start.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,10 +41,10 @@ def get_size(doc_f):
messages = Messages(lang_fetcher=get_lang)


def check_logs():
async def check_logs():
try:
if Config.LOGS_CHANNEL:
c_info = unzipbot_client.get_chat(chat_id=Config.LOGS_CHANNEL)
c_info = await unzipbot_client.get_chat(chat_id=Config.LOGS_CHANNEL)

if c_info.type in (enums.ChatType.PRIVATE, enums.ChatType.BOT):
LOGGER.error(messages.get("start", "PRIVATE_CHAT"))
Expand All @@ -62,10 +62,8 @@ def check_logs():
return False


def dl_thumbs():
loop = asyncio.get_event_loop()
coroutine = get_thumb_users()
thumbs = loop.run_until_complete(coroutine)
async def dl_thumbs():
thumbs = await get_thumb_users()
i = 0
maxthumbs = len(thumbs)
LOGGER.info(messages.get("start", "DL_THUMBS", None, maxthumbs))
Expand All @@ -76,21 +74,19 @@ def dl_thumbs():
if not os.path.exists(file_path):
if thumb.get("url") is None and thumb.get("file_id") is not None:
try:
unzipbot_client.download_media(
await unzipbot_client.download_media(
message=thumb.get("file_id"),
file_name=file_path,
)
except:
# Here we could encounter 400 FILE_REFERENCE_EXPIRED
# A possible fix is to retrieve the message again with chat ID + message ID to get a refreshed file reference
unzipbot_client.send_message(
await unzipbot_client.send_message(
thumb.get("_id"),
messages.get("start", "MISSING_THUMB", thumb.get("_id")),
)
elif thumb.get("url") is not None and thumb.get("file_id") is None:
loop2 = asyncio.get_event_loop()
coroutine2 = download(thumb.get("url"), file_path)
loop2.run_until_complete(coroutine2)
await download(thumb.get("url"), file_path)

if get_size(file_path) in (0, -1):
os.remove(file_path)
Expand All @@ -101,13 +97,7 @@ def dl_thumbs():
LOGGER.info(messages.get("start", "DOWNLOADED_THUMBS", None, i, maxthumbs))


def set_boot_time():
loop = asyncio.get_event_loop()
coroutine = check_boot()
loop.run_until_complete(coroutine)


async def check_boot():
async def set_boot_time():
boot = await get_boot()
await set_old_boot(boot)
await set_boot(boottime)
Expand Down Expand Up @@ -158,12 +148,6 @@ async def warn_users():
await clear_ongoing_tasks()


def removal(firststart=False):
loop = asyncio.get_event_loop()
loop.create_task(remove_expired_tasks(firststart))
loop.run_until_complete(asyncio.sleep(0))


async def remove_expired_tasks(firststart=False):
ongoing_tasks = await get_ongoing_tasks()
await clear_cancel_tasks()
Expand All @@ -175,11 +159,13 @@ async def remove_expired_tasks(firststart=False):
shutil.rmtree(Config.DOWNLOAD_LOCATION)
except:
pass

os.makedirs(Config.DOWNLOAD_LOCATION, exist_ok=True)
else:
for task in ongoing_tasks:
user_id = task.get("user_id")

if not user_id == Config.BOT_OWNER:
if user_id != Config.BOT_OWNER:
current_time = time()
start_time = task.get("start_time")
task_type = task.get("type")
Expand Down Expand Up @@ -224,5 +210,5 @@ async def scheduled_remove_expired_tasks():
await remove_expired_tasks()


def start_cron_jobs():
async def start_cron_jobs():
scheduled_remove_expired_tasks.start()

1 comment on commit 7aef751

@EDM115
Copy link
Owner Author

@EDM115 EDM115 commented on 7aef751 Dec 18, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more. 

2024-12-18 17:03:59,518 - INFO - root - MainThread - Stop signal received (SIGTERM). Exiting...
2024-12-18 17:03:59,520 - INFO - unzipbot - MainThread - ℹ️ The bot goes sleeping at `2024/12/18 - 17:03:59` 😴
2024-12-18 17:03:59,521 - INFO - unzipbot - MainThread - Received stop signal (SIGTERM, 15, None). Exiting...
2024-12-18 17:03:59,523 - INFO - unzipbot - MainThread - ℹ️ The bot goes sleeping at `2024/12/18 - 17:03:59` 😴
Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/app/unzipbot/__main__.py", line 127, in <module>
    unzipbot_client.run(main())
  File "/venv/lib/python3.12/site-packages/pyrogram/methods/utilities/run.py", line 77, in run
    run(coroutine)
  File "/usr/local/lib/python3.12/asyncio/base_events.py", line 686, in run_until_complete
    return future.result()
           ^^^^^^^^^^^^^^^
  File "/app/unzipbot/__main__.py", line 123, in main
    await async_shutdown_bot()
  File "/app/unzipbot/__main__.py", line 31, in async_shutdown_bot
    await asyncio.gather(*tasks, return_exceptions=True)
asyncio.exceptions.CancelledError

Please sign in to comment.