Support celery appids with queue workers (#332)

* Support celery app id with queue workers
* Support short prefix args and add test case for multiple queues

Author: Maor Kleinberger

.gitignore

@@ -145,3 +145,6 @@ cython_debug/
 # Resources
 gprofiler/resources/*
 !gprofiler/resources/flamegraph
+
+# IDEs
+.idea

gprofiler/metadata/application_identifiers.py

@@ -18,7 +18,6 @@
 
 _logger = get_logger_adapter(__name__)
 
-
 _PYTHON_BIN_RE = re.compile(r"^python([23](\.\d{1,2})?)?$")
 
 
@@ -47,7 +46,9 @@ def _is_python_bin(bin_name: str) -> bool:
     return _PYTHON_BIN_RE.match(os.path.basename(bin_name)) is not None
 
 
-def _get_cli_arg_by_name(args: List[str], arg_name: str, check_for_equals_arg: bool = False) -> str:
+def _get_cli_arg_by_name(
+    args: List[str], arg_name: str, check_for_equals_arg: bool = False, check_for_short_prefix_arg: bool = False
+) -> str:
     if arg_name in args:
         return args[args.index(arg_name) + 1]
 
@@ -57,6 +58,11 @@ def _get_cli_arg_by_name(args: List[str], arg_name: str, check_for_equals_arg: b
             if arg_key == arg_name:
                 return arg_val
 
+    if check_for_short_prefix_arg:
+        for arg in args:
+            if arg.startswith(arg_name):
+                return arg[len(arg_name) :]
+
     return _NON_AVAILABLE_ARG
 
 
@@ -201,9 +207,17 @@ def get_application_name(self, process: Process) -> Optional[str]:
         if not self.is_celery_process(process):
             return None
 
-        app_name = _get_cli_arg_by_name(process.cmdline(), "-A") or _get_cli_arg_by_name(
-            process.cmdline(), "--app", check_for_equals_arg=True
-        )
+        app_name = _get_cli_arg_by_name(
+            process.cmdline(), "-A", check_for_short_prefix_arg=True
+        ) or _get_cli_arg_by_name(process.cmdline(), "--app", check_for_equals_arg=True)
+        if app_name is _NON_AVAILABLE_ARG:
+            queue_name = _get_cli_arg_by_name(
+                process.cmdline(), "-Q", check_for_short_prefix_arg=True
+            ) or _get_cli_arg_by_name(process.cmdline(), "--queues", check_for_equals_arg=True)
+            # TODO: One worker can handle multiple queues, it could be useful to encode that into the app id.
+            if queue_name is not _NON_AVAILABLE_ARG:
+                # The queue handler routing is defined in the directory where the worker is run
+                return f"celery queue: {queue_name} ({process.cwd()})"
         if app_name is _NON_AVAILABLE_ARG:
             _logger.warning(
                 f"{self.__class__.__name__}: Couldn't find positional argument -A or --app for application indication",

tests/test_appids.py

@@ -81,14 +81,17 @@ def get_uwsgi_config(process: Process, config_file: str) -> TextIO:
     assert "uwsgi: my.ini" == get_application_name(process_with_cmdline(["uwsgi", "a", "b", "--ini", "my.ini"]))
 
 
-def test_celery() -> None:
+def test_celery_with_app() -> None:
     # celery -A
     assert f"celery: app1 ({PROCESS_CWD}/app1.py)" == get_application_name(
         process_with_cmdline(["celery", "a", "b", "-A", "app1"])
     )
     assert "celery: /path/to/app1 (/path/to/app1.py)" == get_application_name(
         process_with_cmdline(["celery", "a", "b", "-A", "/path/to/app1"])
     )
+    assert "celery: /path/to/app1 (/path/to/app1.py)" == get_application_name(
+        process_with_cmdline(["celery", "a", "b", "-A/path/to/app1"])
+    )
     # python celery -A
     assert f"celery: app1 ({PROCESS_CWD}/app1.py)" == get_application_name(
         process_with_cmdline(["python", "/path/to/celery", "a", "b", "-A", "app1"])
@@ -110,7 +113,36 @@ def test_celery() -> None:
     assert "celery: /path/to/app3 (/path/to/app3.py)" == get_application_name(
         process_with_cmdline(["celery", "a", "b", "--app=/path/to/app3"])
     )
-    # No app
+
+
+def test_celery_with_queue() -> None:
+    # celery -Q queue
+    assert f"celery queue: qqq ({PROCESS_CWD})" == get_application_name(
+        process_with_cmdline(["celery", "a", "b", "-Q", "qqq"])
+    )
+    # celery -Qqueue
+    assert f"celery queue: qqq ({PROCESS_CWD})" == get_application_name(
+        process_with_cmdline(["celery", "a", "b", "-Qqqq"])
+    )
+    # python celery -Q queue
+    assert f"celery queue: qqq ({PROCESS_CWD})" == get_application_name(
+        process_with_cmdline(["python", "/path/to/celery", "a", "b", "-Q", "qqq"])
+    )
+    # --queues queue
+    assert f"celery queue: qqq ({PROCESS_CWD})" == get_application_name(
+        process_with_cmdline(["celery", "a", "b", "--queues", "qqq"])
+    )
+    # --queues=queue
+    assert f"celery queue: qqq ({PROCESS_CWD})" == get_application_name(
+        process_with_cmdline(["celery", "a", "b", "--queues=qqq"])
+    )
+    # multiple queues
+    assert f"celery queue: qqq,www ({PROCESS_CWD})" == get_application_name(
+        process_with_cmdline(["celery", "a", "b", "-Q", "qqq,www"])
+    )
+
+
+def test_celery_without_app() -> None:
     assert get_application_name(process_with_cmdline(["celery", "a", "b"])) is None