sedi-relay/relay/application.py

210 lines
4.6 KiB
Python
Raw Normal View History

2022-11-07 12:54:32 +00:00
import asyncio
import logging
import os
2022-11-20 10:50:14 +00:00
import queue
2022-11-07 12:54:32 +00:00
import signal
2022-11-20 10:50:14 +00:00
import threading
2022-11-07 12:54:32 +00:00
from aiohttp import web
from cachetools import LRUCache
from datetime import datetime, timedelta
from .config import RelayConfig
from .database import RelayDatabase
2022-11-20 10:50:14 +00:00
from .misc import DotDict, check_open_port, request, set_app
2022-11-07 12:54:32 +00:00
from .views import routes
class Application(web.Application):
def __init__(self, cfgpath):
web.Application.__init__(self)
self['starttime'] = None
self['running'] = False
self['is_docker'] = bool(os.environ.get('DOCKER_RUNNING'))
self['config'] = RelayConfig(cfgpath, self['is_docker'])
if not self['config'].load():
self['config'].save()
2022-11-10 17:39:37 +00:00
self['cache'] = DotDict({key: Cache(maxsize=self['config'][key]) for key in self['config'].cachekeys})
2022-11-07 12:54:32 +00:00
self['semaphore'] = asyncio.Semaphore(self['config'].push_limit)
2022-11-20 10:50:14 +00:00
self['workers'] = []
self['last_worker'] = 0
2022-11-07 12:54:32 +00:00
set_app(self)
2022-11-18 19:10:39 +00:00
self['database'] = RelayDatabase(self['config'])
self['database'].load()
self.set_signal_handler()
2022-11-07 12:54:32 +00:00
@property
def cache(self):
return self['cache']
@property
def config(self):
return self['config']
@property
def database(self):
return self['database']
@property
def is_docker(self):
return self['is_docker']
@property
def semaphore(self):
return self['semaphore']
@property
def uptime(self):
if not self['starttime']:
return timedelta(seconds=0)
uptime = datetime.now() - self['starttime']
return timedelta(seconds=uptime.seconds)
2022-11-20 10:50:14 +00:00
def push_message(self, inbox, message):
worker = self['workers'][self['last_worker']]
worker.queue.put((inbox, message))
self['last_worker'] += 1
if self['last_worker'] >= len(self['workers']):
self['last_worker'] = 0
2022-11-07 12:54:32 +00:00
def set_signal_handler(self):
2022-11-16 18:26:47 +00:00
for sig in {'SIGHUP', 'SIGINT', 'SIGQUIT', 'SIGTERM'}:
try:
signal.signal(getattr(signal, sig), self.stop)
# some signals don't exist in windows, so skip them
except AttributeError:
pass
2022-11-07 12:54:32 +00:00
def run(self):
if not check_open_port(self.config.listen, self.config.port):
return logging.error(f'A server is already running on port {self.config.port}')
for route in routes:
if route[1] == '/stats' and logging.DEBUG < logging.root.level:
continue
self.router.add_route(*route)
logging.info(f'Starting webserver at {self.config.host} ({self.config.listen}:{self.config.port})')
asyncio.run(self.handle_run())
def stop(self, *_):
self['running'] = False
async def handle_run(self):
self['running'] = True
2022-11-20 10:50:14 +00:00
if self.config.workers > 0:
for i in range(self.config.workers):
worker = PushWorker(self)
worker.start()
self['workers'].append(worker)
2022-11-07 12:54:32 +00:00
runner = web.AppRunner(self, access_log_format='%{X-Forwarded-For}i "%r" %s %b "%{User-Agent}i"')
await runner.setup()
site = web.TCPSite(runner,
host = self.config.listen,
port = self.config.port,
reuse_address = True
)
await site.start()
self['starttime'] = datetime.now()
while self['running']:
await asyncio.sleep(0.25)
await site.stop()
self['starttime'] = None
self['running'] = False
2022-11-20 10:50:14 +00:00
self['workers'].clear()
2022-11-10 17:39:37 +00:00
class Cache(LRUCache):
def set_maxsize(self, value):
self.__maxsize = int(value)
2022-11-18 21:38:39 +00:00
2022-11-20 10:50:14 +00:00
class PushWorker(threading.Thread):
def __init__(self, app):
threading.Thread.__init__(self)
self.app = app
self.queue = queue.Queue()
def run(self):
asyncio.run(self.handle_queue())
async def handle_queue(self):
while self.app['running']:
try:
inbox, message = self.queue.get(block=True, timeout=0.25)
self.queue.task_done()
await request(inbox, message)
logging.verbose(f'New push from Thread-{threading.get_ident()}')
except queue.Empty:
pass
2022-11-18 21:38:39 +00:00
## Can't sub-class web.Request, so let's just add some properties
def request_actor(self):
try: return self['actor']
except KeyError: pass
def request_instance(self):
try: return self['instance']
except KeyError: pass
2022-11-18 21:38:39 +00:00
def request_message(self):
try: return self['message']
except KeyError: pass
def request_signature(self):
if 'signature' not in self._state:
try: self['signature'] = DotDict.new_from_signature(self.headers['signature'])
except KeyError: return
return self['signature']
setattr(web.Request, 'actor', property(request_actor))
setattr(web.Request, 'instance', property(request_instance))
2022-11-18 21:38:39 +00:00
setattr(web.Request, 'message', property(request_message))
setattr(web.Request, 'signature', property(request_signature))
setattr(web.Request, 'cache', property(lambda self: self.app.cache))
setattr(web.Request, 'config', property(lambda self: self.app.config))
setattr(web.Request, 'database', property(lambda self: self.app.database))
setattr(web.Request, 'semaphore', property(lambda self: self.app.semaphore))