From 4163597a082859a11ae780416790dba0e786cf81 Mon Sep 17 00:00:00 2001 From: BennyThink Date: Fri, 7 Jan 2022 18:43:59 +0800 Subject: [PATCH] add worker inspect in ping and download text --- .github/workflows/codeql-analysis.yml | 70 --------------------------- Makefile | 2 +- ytdlbot/constant.py | 18 ++++++- ytdlbot/downloader.py | 3 +- ytdlbot/tasks.py | 13 ++++- ytdlbot/utils.py | 38 ++++++++++++++- ytdlbot/ytdl_bot.py | 15 +++--- 7 files changed, 76 insertions(+), 83 deletions(-) delete mode 100644 .github/workflows/codeql-analysis.yml diff --git a/.github/workflows/codeql-analysis.yml b/.github/workflows/codeql-analysis.yml deleted file mode 100644 index a75e2f3..0000000 --- a/.github/workflows/codeql-analysis.yml +++ /dev/null @@ -1,70 +0,0 @@ -# For most projects, this workflow file will not need changing; you simply need -# to commit it to your repository. -# -# You may wish to alter this file to override the set of languages analyzed, -# or to provide custom queries or build logic. -# -# ******** NOTE ******** -# We have attempted to detect the languages in your repository. Please check -# the `language` matrix defined below to confirm you have the correct set of -# supported CodeQL languages. -# -name: "CodeQL" - -on: - push: - branches: [ master ] - pull_request: - # The branches below must be a subset of the branches above - branches: [ master ] - schedule: - - cron: '38 9 * * 4' - -jobs: - analyze: - name: Analyze - runs-on: ubuntu-latest - permissions: - actions: read - contents: read - security-events: write - - strategy: - fail-fast: false - matrix: - language: [ 'python' ] - # CodeQL supports [ 'cpp', 'csharp', 'go', 'java', 'javascript', 'python', 'ruby' ] - # Learn more about CodeQL language support at https://git.io/codeql-language-support - - steps: - - name: Checkout repository - uses: actions/checkout@v2 - - # Initializes the CodeQL tools for scanning. - - name: Initialize CodeQL - uses: github/codeql-action/init@v1 - with: - languages: ${{ matrix.language }} - # If you wish to specify custom queries, you can do so here or in a config file. - # By default, queries listed here will override any specified in a config file. - # Prefix the list here with "+" to use these queries and those in the config file. - # queries: ./path/to/local/query, your-org/your-repo/queries@main - - # Autobuild attempts to build any compiled languages (C/C++, C#, or Java). - # If this step fails, then you should remove it and run the build manually (see below) - - name: Autobuild - uses: github/codeql-action/autobuild@v1 - - # ℹī¸ Command-line programs to run using the OS shell. - # 📚 https://git.io/JvXDl - - # ✏ī¸ If the Autobuild fails above, remove it and uncomment the following three lines - # and modify them (or add more) to build your code if your project - # uses a compiled language - - #- run: | - # make bootstrap - # make release - - - name: Perform CodeQL Analysis - uses: github/codeql-action/analyze@v1 diff --git a/Makefile b/Makefile index 847d686..785a2d0 100644 --- a/Makefile +++ b/Makefile @@ -22,5 +22,5 @@ upgrade-all-worker: bash upgrade_worker.sh tag: - git tag -a v$(shell date "+%Y-%m-%d") -m v$(shell date "+%Y-%m-%d") + git tag -a v$(shell date "+%Y-%m-%d")_$(shell git rev-parse --short HEAD) -m v$(shell date "+%Y-%m-%d") git push --tags \ No newline at end of file diff --git a/ytdlbot/constant.py b/ytdlbot/constant.py index 870fe32..5e6cbdc 100644 --- a/ytdlbot/constant.py +++ b/ytdlbot/constant.py @@ -9,10 +9,11 @@ __author__ = "Benny " import time -from config import (AFD_LINK, COFFEE_LINK, ENABLE_VIP, EX, MULTIPLY, +from config import (AFD_LINK, COFFEE_LINK, ENABLE_VIP, EX, MULTIPLY, OWNER, REQUIRED_MEMBERSHIP, USD2CNY) from downloader import sizeof_fmt from limit import QUOTA, VIP +from utils import get_queue_stat, get_func_queue class BotText: @@ -109,3 +110,18 @@ Sending format: **{1}** return f"Hello {v[1]}, VIP{v[-2]}â˜ēī¸\n\n" else: return "" + + @staticmethod + def queue_stats(): + _, _, _, stats = get_queue_stat() + return stats + + @staticmethod + def get_receive_link_text(): + reserved = get_func_queue("reserved") + if reserved == 0: + text = "Your task was added to active queue.\nProcessing...\n\n" + else: + text = f"Too many tasks. Your tasks was added to the reserved queue {reserved}." + + return text diff --git a/ytdlbot/downloader.py b/ytdlbot/downloader.py index 66bf8ef..17eabb8 100644 --- a/ytdlbot/downloader.py +++ b/ytdlbot/downloader.py @@ -64,6 +64,7 @@ def download_hook(d: dict, bot_msg): filesize = sizeof_fmt(total) max_size = 2 * 1024 * 1024 * 1024 if total > max_size: + # only for one track, e.g. video. So it's not so accurate raise ValueError(f"\nYour video is too large. " f"{filesize} will exceed Telegram's max limit {sizeof_fmt(max_size)}") @@ -87,7 +88,7 @@ def upload_hook(current, total, bot_msg): def check_quota(file_size, chat_id) -> ("bool", "str"): remain, _, ttl = VIP().check_remaining_quota(chat_id) if file_size > remain: - refresh_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(ttl + time.time())) + refresh_time = current_time(ttl + time.time()) err = f"Quota exceed, you have {sizeof_fmt(remain)} remaining, " \ f"but you want to download a video with {sizeof_fmt(file_size)} in size. \n" \ f"Try again in {ttl} seconds({refresh_time})" diff --git a/ytdlbot/tasks.py b/ytdlbot/tasks.py index 59f64f1..58cc31d 100644 --- a/ytdlbot/tasks.py +++ b/ytdlbot/tasks.py @@ -18,7 +18,7 @@ from pyrogram import idle from pyrogram.types import InlineKeyboardButton, InlineKeyboardMarkup from client_init import create_app -from config import BROKER, ENABLE_CELERY, WORKERS +from config import BROKER, ENABLE_CELERY, OWNER, WORKERS from constant import BotText from db import Redis from downloader import sizeof_fmt, upload_hook, ytdl_download @@ -53,6 +53,14 @@ def download_entrance(bot_msg, client, url): normal_download(bot_msg, client, url) +def get_worker_status(username): + worker_name = os.getenv("WORKER_NAME") + me = celery_client.get_me() + if worker_name and username == OWNER: + return f"Downloaded by {me.mention()}-{worker_name}" + return f"Downloaded by {me.mention()}" + + def normal_download(bot_msg, client, url): chat_id = bot_msg.chat.id temp_dir = tempfile.TemporaryDirectory() @@ -78,7 +86,8 @@ def normal_download(bot_msg, client, url): remain = bot_text.remaining_quota_caption(chat_id) size = sizeof_fmt(os.stat(video_path).st_size) meta = get_metadata(video_path) - cap = f"`{filename}`\n\n{url}\n\nInfo: {meta['width']}x{meta['height']} {size}\n\n{remain}" + worker = get_worker_status(bot_msg.chat.username) + cap = f"`{filename}`\n\n{url}\n\nInfo: {meta['width']}x{meta['height']} {size}\n\n{remain}\n{worker}" settings = get_user_settings(str(chat_id)) if settings[2] == "document": logging.info("Sending as document") diff --git a/ytdlbot/utils.py b/ytdlbot/utils.py index 69cbc80..f19a393 100644 --- a/ytdlbot/utils.py +++ b/ytdlbot/utils.py @@ -15,7 +15,11 @@ import uuid import ffmpeg +from config import ENABLE_CELERY from db import MySQL +from flower_tasks import app + +inspect = app.control.inspect() def apply_log_formatter(): @@ -95,10 +99,40 @@ def get_metadata(video_path): return dict(height=height, width=width, duration=duration, thumb=thumb) -def current_time(): - return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) +def current_time(ts=None): + return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(ts)) def get_revision(): revision = subprocess.check_output("git -C ../ rev-parse --short HEAD".split()).decode("u8").replace("\n", "") return revision + + +def get_func_queue(func) -> int: + count = 0 + data = getattr(inspect, func)() or {} + for _, task in data.items(): + count += len(task) + return count + + +def get_queue_stat() -> (int, int, int, str): + concurrency = 0 + if ENABLE_CELERY is False: + return 0, 0, 0, "" + + stats = inspect.stats() + if stats is None: + err = "No worker is running." + logging.error(err) + return 0, 0, 0, err + + for _, stat in stats.items(): + concurrency += stat["pool"]["max-concurrency"] + + active = get_func_queue("active") + reserved = get_func_queue("reserved") + ping = inspect.ping() + stats = f"concurrency {concurrency}, active {active}, reserved {reserved}.\n\n{ping}" + + return concurrency, active, reserved, stats diff --git a/ytdlbot/ytdl_bot.py b/ytdlbot/ytdl_bot.py index a75a895..7904a13 100644 --- a/ytdlbot/ytdl_bot.py +++ b/ytdlbot/ytdl_bot.py @@ -77,7 +77,10 @@ def start_handler(client: "Client", message: "types.Message"): logging.info("Welcome to youtube-dl bot!") client.send_chat_action(chat_id, "typing") greeting = bot_text.get_vip_greeting(chat_id) - client.send_message(message.chat.id, greeting + bot_text.start + "\n\n" + bot_text.remaining_quota_caption(chat_id)) + quota = bot_text.remaining_quota_caption(chat_id) + text = f"{greeting}{bot_text.start}\n\n{quota}" + + client.send_message(message.chat.id, text) @app.on_message(filters.command(["help"])) @@ -96,7 +99,8 @@ def ping_handler(client: "Client", message: "types.Message"): else: bot_info = get_runtime("ytdlbot_ytdl_1", "YouTube-dl") if message.chat.username == OWNER: - client.send_document(chat_id, Redis().generate_file(), caption=bot_info) + stats = bot_text.queue_stats() + client.send_document(chat_id, Redis().generate_file(), caption=f"{bot_info}\n{stats}") else: client.send_message(chat_id, f"{bot_info}") @@ -156,6 +160,7 @@ def vip_handler(client: "Client", message: "types.Message"): def download_handler(client: "Client", message: "types.Message"): # check remaining quota chat_id = message.chat.id + client.send_chat_action(chat_id, 'typing') Redis().user_count(chat_id) url = re.sub(r'/ytdl\s*', '', message.text) @@ -167,13 +172,11 @@ def download_handler(client: "Client", message: "types.Message"): return Redis().update_metrics("video_request") - bot_msg: typing.Union["types.Message", "typing.Any"] = message.reply_text("Processing", quote=True) + text = bot_text.get_receive_link_text() + bot_msg: typing.Union["types.Message", "typing.Any"] = message.reply_text(text, quote=True) client.send_chat_action(chat_id, 'upload_video') - # temp_dir = tempfile.TemporaryDirectory() download_entrance(bot_msg, client, url) - # temp_dir.cleanup() - @app.on_callback_query(filters.regex(r"document|video")) def send_method_callback(client: "Client", callback_query: types.CallbackQuery):