shaba/openuds
1
0
Fork 1
mirror of https://github.com/dkmstr/openuds.git synced 2024-12-22 13:34:04 +03:00
Code Issues Releases Wiki Activity

Separated processes manager from main uds_tunnel

Browse Source
This commit is contained in:
Adolfo Gómez García 2021-08-05 12:53:44 +02:00
parent c21c0b44ce
commit 1be49a6e0e
2 changed files with 103 additions and 64 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

99
tunnel-server/src/uds_tunnel/processes.py Normal file
View File

@ -0,0 +1,99 @@
import multiprocessing
import logging
import typing
import curio
import psutil
from . import config
if typing.TYPE_CHECKING:
from multiprocessing.connection import Connection
from multiprocessing.managers import Namespace
logger = logging.getLogger(__name__)
class Processes:
"""
This class is used to store the processes that are used by the tunnel.
"""
children: typing.List[
typing.Tuple['Connection', multiprocessing.Process, psutil.Process]
]
process: typing.Callable
cfg: config.ConfigurationType
ns: 'Namespace'
def __init__(self, process: typing.Callable, cfg: config.ConfigurationType, ns: 'Namespace') -> None:
self.children = []
self.process = process # type: ignore
self.cfg = cfg
self.ns = ns
for i in range(cfg.workers):
self.add_child_pid()
def add_child_pid(self):
own_conn, child_conn = multiprocessing.Pipe()
task = multiprocessing.Process(
target=curio.run,
args=(self.process, child_conn, self.cfg, self.ns)
)
task.start()
logger.debug('ADD CHILD PID: %s', task.pid)
self.children.append((own_conn, task, psutil.Process(task.pid)))
def best_child(self) -> 'Connection':
best: typing.Tuple[float, 'Connection'] = (1000.0, self.children[0][0])
missingProcesses: typing.List[int] = []
for i, c in enumerate(self.children):
try:
if c[2].status() == 'zombie': # Bad kill!!
raise psutil.ZombieProcess(c[2].pid)
percent = c[2].cpu_percent()
except (psutil.ZombieProcess, psutil.NoSuchProcess) as e:
# Process is missing...
logger.warning('Missing process found: %s', e.pid)
try:
c[0].close() # Close pipe to missing process
except Exception:
logger.debug('Could not close handle for %s', e.pid)
try:
c[1].kill()
c[1].close()
except Exception:
logger.debug('Could not close process %s', e.pid)
missingProcesses.append(i)
continue
logger.debug('PID %s has %s', c[2].pid, percent)
if percent < best[0]:
best = (percent, c[0])
# If we have a missing process, try to add it back
if missingProcesses:
logger.debug('Regenerating missing processes: %s', len(missingProcesses))
# Regenerate childs and recreate new proceeses for requests...
tmpChilds = [
self.children[i]
for i in range(len(self.children))
if i not in missingProcesses
]
self.children[:] = tmpChilds
# Now add new children
for i in range(len(missingProcesses)):
self.add_child_pid()
return best[1]
def stop(self):
# Try to stop running childs
for i in self.children:
try:
i[2].kill()
except Exception as e:
logger.info('KILLING child %s: %s', i[2], e)

68
tunnel-server/src/udstunnel.py
View File

@ -48,6 +48,7 @@ from uds_tunnel import proxy
from uds_tunnel import consts
from uds_tunnel import message
from uds_tunnel import stats
from uds_tunnel import processes
if typing.TYPE_CHECKING:
from multiprocessing.connection import Connection
@ -217,72 +218,16 @@ def tunnel_main():
signal.signal(signal.SIGINT, stop_signal)
signal.signal(signal.SIGTERM, stop_signal)
# Creates as many processes and pipes as required
child: typing.List[
typing.Tuple['Connection', multiprocessing.Process, psutil.Process]
] = []
stats_collector = stats.GlobalStats()
def add_child_pid():
own_conn, child_conn = multiprocessing.Pipe()
task = multiprocessing.Process(
target=curio.run,
args=(tunnel_proc_async, child_conn, cfg, stats_collector.ns),
)
task.start()
logger.debug('ADD CHILD PID: %s', task.pid)
child.append((own_conn, task, psutil.Process(task.pid)))
for i in range(cfg.workers):
add_child_pid()
def best_child() -> 'Connection':
best: typing.Tuple[float, 'Connection'] = (1000.0, child[0][0])
missingProcesses = []
for i, c in enumerate(child):
try:
if c[2].status() == 'zombie': # Bad kill!!
raise psutil.ZombieProcess(c[2].pid)
percent = c[2].cpu_percent()
except (psutil.ZombieProcess, psutil.NoSuchProcess) as e:
# Process is missing...
logger.warning('Missing process found: %s', e.pid)
try:
c[0].close() # Close pipe to missing process
except Exception:
logger.debug('Could not close handle for %s', e.pid)
try:
c[1].kill()
c[1].close()
except Exception:
logger.debug('Could not close process %s', e.pid)
missingProcesses.append(i)
continue
logger.debug('PID %s has %s', c[2].pid, percent)
if percent < best[0]:
best = (percent, c[0])
if missingProcesses:
logger.debug('Regenerating missing processes: %s', len(missingProcesses))
# Regenerate childs and recreate new proceeses to process requests...
tmpChilds = [child[i] for i in range(len(child)) if i not in missingProcesses]
child[:] = tmpChilds
# Now add new children
for i in range(len(missingProcesses)):
add_child_pid()
return best[1]
prcs = processes.Processes(tunnel_proc_async, cfg, stats_collector.ns)
try:
while not do_stop:
try:
client, addr = sock.accept()
# Select BEST process for sending this new connection
best_child().send(
prcs.best_child().send(
message.Message(message.Command.TUNNEL, (client, addr))
)
del client # Ensure socket is controlled on child process
@ -297,12 +242,7 @@ def tunnel_main():
if sock:
sock.close()
# Try to stop running childs
for i in child:
try:
i[2].kill()
except Exception as e:
logger.info('KILLING child %s: %s', i[2], e)
prcs.stop()
try:
if cfg.pidfile: