add worker inspect in ping and download text

.github/workflows/codeql-analysis.yml

@@ -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

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

ytdlbot/constant.py

@@ -9,10 +9,11 @@ __author__ = "Benny <benny.think@gmail.com>"
 
 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

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})"

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")

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

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):