reflex/pynecone/utils/processes.py
2023-03-16 14:59:27 -07:00

123 lines
3.0 KiB
Python

"""Process operations."""
from __future__ import annotations
import contextlib
import os
import signal
import sys
from typing import Optional
from urllib.parse import urlparse
import psutil
from pynecone import constants
from pynecone.config import get_config
from pynecone.utils import console, prerequisites
def kill(pid):
"""Kill a process.
Args:
pid: The process ID.
"""
os.kill(pid, signal.SIGTERM)
def get_num_workers() -> int:
"""Get the number of backend worker processes.
Returns:
The number of backend worker processes.
"""
return 1 if prerequisites.get_redis() is None else (os.cpu_count() or 1) * 2 + 1
def get_api_port() -> int:
"""Get the API port.
Returns:
The API port.
"""
port = urlparse(get_config().api_url).port
if port is None:
port = urlparse(constants.API_URL).port
assert port is not None
return port
def get_process_on_port(port) -> Optional[psutil.Process]:
"""Get the process on the given port.
Args:
port: The port.
Returns:
The process on the given port.
"""
for proc in psutil.process_iter(["pid", "name", "cmdline"]):
try:
for conns in proc.connections(kind="inet"):
if conns.laddr.port == int(port):
return proc
except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
pass
return None
def is_process_on_port(port) -> bool:
"""Check if a process is running on the given port.
Args:
port: The port.
Returns:
Whether a process is running on the given port.
"""
return get_process_on_port(port) is not None
def kill_process_on_port(port):
"""Kill the process on the given port.
Args:
port: The port.
"""
if get_process_on_port(port) is not None:
with contextlib.suppress(psutil.AccessDenied):
get_process_on_port(port).kill() # type: ignore
def change_or_terminate_port(port, _type) -> str:
"""Terminate or change the port.
Args:
port: The port.
_type: The type of the port.
Returns:
The new port or the current one.
"""
console.print(
f"Something is already running on port [bold underline]{port}[/bold underline]. This is the port the {_type} runs on."
)
frontend_action = console.ask("Kill or change it?", choices=["k", "c", "n"])
if frontend_action == "k":
kill_process_on_port(port)
return port
elif frontend_action == "c":
new_port = console.ask("Specify the new port")
# Check if also the new port is used
if is_process_on_port(new_port):
return change_or_terminate_port(new_port, _type)
else:
console.print(
f"The {_type} will run on port [bold underline]{new_port}[/bold underline]."
)
return new_port
else:
console.print("Exiting...")
sys.exit()