Implement concurrency for validating process health

Update logger names across the module and README.md

README.md

@@ -58,6 +58,8 @@ pyninja start
 - **WORKERS** - Number of workers for the uvicorn server.
 - **REMOTE_EXECUTION** - Boolean flag to enable remote execution.
 - **API_SECRET** - Secret access key for running commands on server remotely.
+- **DATABASE** - FilePath to store the auth database that handles the authentication errors.
+- **RATE_LIMIT** - List of dictionaries with `max_requests` and `seconds` to apply as rate limit.
 - **APIKEY** - API Key for authentication.
 
 ⚠️ Enabling remote execution can be extremely risky and can be a major security threat. So use **caution** and set the **API_SECRET** to a strong value.

docs/README.html

@@ -95,6 +95,8 @@ <h2>Environment Variables<a class="headerlink" href="#environment-variables" tit
 <li><p><strong>WORKERS</strong> - Number of workers for the uvicorn server.</p></li>
 <li><p><strong>REMOTE_EXECUTION</strong> - Boolean flag to enable remote execution.</p></li>
 <li><p><strong>API_SECRET</strong> - Secret access key for running commands on server remotely.</p></li>
+<li><p><strong>DATABASE</strong> - FilePath to store the auth database that handles the authentication errors.</p></li>
+<li><p><strong>RATE_LIMIT</strong> - List of dictionaries with <code class="docutils literal notranslate"><span class="pre">max_requests</span></code> and <code class="docutils literal notranslate"><span class="pre">seconds</span></code> to apply as rate limit.</p></li>
 <li><p><strong>APIKEY</strong> - API Key for authentication.</p></li>
 </ul>
 <p>⚠️ Enabling remote execution can be extremely risky and can be a major security threat. So use <strong>caution</strong> and set the <strong>API_SECRET</strong> to a strong value.</p>

docs/README.md

@@ -58,6 +58,8 @@ pyninja start
 - **WORKERS** - Number of workers for the uvicorn server.
 - **REMOTE_EXECUTION** - Boolean flag to enable remote execution.
 - **API_SECRET** - Secret access key for running commands on server remotely.
+- **DATABASE** - FilePath to store the auth database that handles the authentication errors.
+- **RATE_LIMIT** - List of dictionaries with `max_requests` and `seconds` to apply as rate limit.
 - **APIKEY** - API Key for authentication.
 
 ⚠️ Enabling remote execution can be extremely risky and can be a major security threat. So use **caution** and set the **API_SECRET** to a strong value.

docs/_sources/README.md.txt

@@ -58,6 +58,8 @@ pyninja start
 - **WORKERS** - Number of workers for the uvicorn server.
 - **REMOTE_EXECUTION** - Boolean flag to enable remote execution.
 - **API_SECRET** - Secret access key for running commands on server remotely.
+- **DATABASE** - FilePath to store the auth database that handles the authentication errors.
+- **RATE_LIMIT** - List of dictionaries with `max_requests` and `seconds` to apply as rate limit.
 - **APIKEY** - API Key for authentication.
 
 ⚠️ Enabling remote execution can be extremely risky and can be a major security threat. So use **caution** and set the **API_SECRET** to a strong value.

pyninja/auth.py

@@ -9,7 +9,7 @@
 
 from pyninja import database, exceptions, models
 
-LOGGER = logging.getLogger("uvicorn.error")
+LOGGER = logging.getLogger("uvicorn.default")
 EPOCH = lambda: int(time.time())  # noqa: E731
 SECURITY = HTTPBearer()
 

pyninja/main.py

@@ -7,14 +7,22 @@
 import pyninja
 from pyninja import models, routers, squire
 
-LOGGER = logging.getLogger("uvicorn.error")
+LOGGER = logging.getLogger("uvicorn.default")
 
 
 def start(**kwargs) -> None:
     """Starter function for the API, which uses uvicorn server as trigger.
 
     Keyword Args:
-        env_file: Filepath for the ``.env`` file.
+        - env_file - Env filepath to load the environment variables.
+        - ninja_host - Hostname for the API server.
+        - ninja_port - Port number for the API server.
+        - workers - Number of workers for the uvicorn server.
+        - remote_execution - Boolean flag to enable remote execution.
+        - api_secret - Secret access key for running commands on server remotely.
+        - database - FilePath to store the auth database that handles the authentication errors.
+        - rate_limit - List of dictionaries with `max_requests` and `seconds` to apply as rate limit.
+        - apikey - API Key for authentication.
     """
     if env_file := kwargs.get("env_file"):
         models.env = squire.env_loader(env_file)

pyninja/models.py

@@ -4,14 +4,7 @@
 import sqlite3
 from typing import Dict, List, Set, Tuple
 
-from pydantic import (
-    BaseModel,
-    Field,
-    FilePath,
-    PositiveFloat,
-    PositiveInt,
-    field_validator,
-)
+from pydantic import BaseModel, Field, FilePath, PositiveInt, field_validator
 from pydantic_settings import BaseSettings
 
 
@@ -60,7 +53,7 @@ class Payload(BaseModel):
     """
 
     command: str
-    timeout: PositiveInt | PositiveFloat = 3
+    timeout: PositiveInt = 3
 
 
 class ServiceStatus(BaseModel):

pyninja/process.py

@@ -1,54 +1,76 @@
 import logging
-from collections.abc import Generator
-from typing import Dict
+import os
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Dict, List
 
 import psutil
+from pydantic import PositiveInt
 
-LOGGER = logging.getLogger("uvicorn.error")
+LOGGER = logging.getLogger("uvicorn.default")
 
 
-def get_process_status(process_name: str) -> Generator[Dict[str, int]]:
+def get_process_status(
+    process_name: str, cpu_interval: PositiveInt
+) -> List[Dict[str, int | float | str | bool]]:
     """Get process information by name.
 
     Args:
         process_name: Name of the process.
+        cpu_interval: CPU interval to get the CPU performance.
 
-    Yields:
-        Generator[Dict[str, int]]:
-        Yields the process metrics as a dictionary of key-value pairs.
+    Returns:
+        List[Dict[str, int | float | str | bool]]:
+        Returns a list of performance report for each process hosting the given process name.
     """
-    # todo: implement concurrency
-    for proc in psutil.process_iter(["pid", "name"]):
-        if proc.name().lower() == process_name.lower():
-            process = psutil.Process(proc.pid)
-            process._name = process_name
-            try:
-                perf_report = get_performance(process)
-                LOGGER.info({f"{process_name} [{process.pid}]": perf_report})
-                perf_report["pname"] = process_name
-                perf_report["zombie"] = False
-                yield perf_report
-            except psutil.ZombieProcess as warn:
-                LOGGER.warning(warn)
-                yield {"zombie": True, "process_name": process_name}
-
-
-def get_performance(process: psutil.Process) -> Dict[str, int | float]:
-    """Checks performance by monitoring CPU utilization, number of threads and open files.
+    result = []
+    futures = {}
+    executor = ThreadPoolExecutor(max_workers=os.cpu_count())
+    with executor:
+        for proc in psutil.process_iter(["pid", "name"]):
+            if proc.name().lower() == process_name.lower():
+                future = executor.submit(
+                    get_performance, process=proc, cpu_interval=cpu_interval
+                )
+                futures[future] = proc.name()
+    for future in as_completed(futures):
+        if future.exception():
+            LOGGER.error(
+                "Thread processing for '%s' received an exception: %s",
+                futures[future],
+                future.exception(),
+            )
+        else:
+            result.append(future.result())
+    return result
+
+
+def get_performance(
+    process: psutil.Process, cpu_interval: PositiveInt
+) -> Dict[str, int | float | str | bool]:
+    """Checks process performance by monitoring CPU utilization, number of threads and open files.
 
     Args:
         process: Process object.
+        cpu_interval: CPU interval to get the CPU performance.
 
     Returns:
-        Dict[str, int]:
+        Dict[str, int | float | str | bool]:
         Returns the process metrics as key-value pairs.
     """
-    cpu = process.cpu_percent(interval=0.5)
-    threads = process.num_threads()
-    open_files = len(process.open_files())
-    return {
-        "cpu": cpu,
-        "threads": threads,
-        "open_files": open_files,
-        "pid": process.pid.real,
-    }
+    try:
+        cpu = process.cpu_percent(interval=cpu_interval)
+        threads = process.num_threads()
+        open_files = len(process.open_files())
+        perf_report = {
+            "pid": process.pid.real,
+            "pname": process.name(),
+            "cpu": cpu,
+            "threads": threads,
+            "open_files": open_files,
+            "zombie": False,
+        }
+        LOGGER.info({f"{process.name()} [{process.pid}]": perf_report})
+    except psutil.ZombieProcess as warn:
+        LOGGER.warning(warn)
+        perf_report = {"zombie": True, "process_name": process.name()}
+    return perf_report

pyninja/routers.py

@@ -7,10 +7,11 @@
 from fastapi.responses import RedirectResponse
 from fastapi.routing import APIRoute
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
+from pydantic import PositiveFloat, PositiveInt
 
 from pyninja import auth, exceptions, models, process, rate_limit, service, squire
 
-LOGGER = logging.getLogger("uvicorn.error")
+LOGGER = logging.getLogger("uvicorn.default")
 security = HTTPBearer()
 
 
@@ -24,15 +25,20 @@ async def run_command(
 
     **Args:**
 
+        request: Reference to the FastAPI request object.
         payload: Payload received as request body.
+        apikey: API Key to authenticate the request.
+        token: API secret to authenticate the request.
 
     **Raises:**
 
         APIResponse:
         Raises the HTTPStatus object with a status code and detail as response.
     """
     await auth.level_2(request, apikey, token)
-    LOGGER.info("Requested command: '%s'", payload.command)
+    LOGGER.info(
+        "Requested command: '%s' with timeout: %ds", payload.command, payload.timeout
+    )
     try:
         response = squire.process_command(payload.command, payload.timeout)
     except subprocess.TimeoutExpired as warn:
@@ -46,21 +52,25 @@ async def run_command(
 async def process_status(
     request: Request,
     process_name: str,
+    cpu_interval: PositiveInt | PositiveFloat = 1,
     apikey: HTTPAuthorizationCredentials = Depends(security),
 ):
     """**API function to monitor a process.**
 
     **Args:**
 
+        request: Reference to the FastAPI request object.
         process_name: Name of the process to check status.
+        cpu_interval: Interval in seconds to get the CPU usage.
+        apikey: API Key to authenticate the request.
 
     **Raises:**
 
         APIResponse:
         Raises the HTTPStatus object with a status code and detail as response.
     """
     await auth.level_1(request, apikey)
-    if response := list(process.get_process_status(process_name)):
+    if response := process.get_process_status(process_name, cpu_interval):
         raise exceptions.APIResponse(status_code=HTTPStatus.OK.real, detail=response)
     LOGGER.error("%s: 404 - No such process", process_name)
     raise exceptions.APIResponse(
@@ -77,7 +87,9 @@ async def service_status(
 
     **Args:**
 
+        request: Reference to the FastAPI request object.
         service_name: Name of the service to check status.
+        apikey: API Key to authenticate the request.
 
     **Raises:**
 
@@ -114,12 +126,12 @@ def get_all_routes() -> List[APIRoute]:
         List[APIRoute]:
         Returns the routes as a list of APIRoute objects.
     """
-    APIRoute(path="/", endpoint=docs, methods=["GET"], include_in_schema=False)
     dependencies = [
         Depends(dependency=rate_limit.RateLimiter(each_rate_limit).init)
         for each_rate_limit in models.env.rate_limit
     ]
     routes = [
+        APIRoute(path="/", endpoint=docs, methods=["GET"], include_in_schema=False),
         APIRoute(
             path="/service-status",
             endpoint=service_status,

pyninja/service.py

@@ -14,14 +14,14 @@
 
 
 def get_service_status(service_name: str) -> models.ServiceStatus:
-    """Get service status.
+    """Get service status by name.
 
     Args:
-        service_name (str): Name of the service.
+        service_name: Name of the service.
 
     Returns:
         ServiceStatus:
-        Returns an instance of the ServiceStatus.
+        Returns an instance of the ServiceStatus object.
     """
     running = models.ServiceStatus(
         status_code=HTTPStatus.OK.real,

pyninja/squire.py

@@ -10,7 +10,7 @@
 
 from pyninja.models import EnvConfig
 
-LOGGER = logging.getLogger("uvicorn.error")
+LOGGER = logging.getLogger("uvicorn.default")
 
 
 def process_command(