From ee4af2b1f398facca9ca86b75ea0e1078b192648 Mon Sep 17 00:00:00 2001 From: mar77i Date: Tue, 22 Sep 2026 14:20:08 +0200 Subject: [PATCH] Initial commit --- .gitignore | 4 + chat/__init__.py | 0 chat/activity.py | 171 ++++++++++++++++++++++++++ chat/apps.py | 5 + chat/asgi.py | 37 ++++++ chat/lifespan.py | 102 +++++++++++++++ chat/management/__init__.py | 0 chat/management/commands/__init__.py | 0 chat/management/commands/runserver.py | 34 +++++ chat/settings.py | 136 ++++++++++++++++++++ chat/urls.py | 22 ++++ chat/utils.py | 34 +++++ chat/websocket.py | 67 ++++++++++ gunicorn.conf.py | 9 ++ manage.py | 22 ++++ setup_venv.sh | 21 ++++ 16 files changed, 664 insertions(+) create mode 100644 .gitignore create mode 100644 chat/__init__.py create mode 100644 chat/activity.py create mode 100644 chat/apps.py create mode 100644 chat/asgi.py create mode 100644 chat/lifespan.py create mode 100644 chat/management/__init__.py create mode 100644 chat/management/commands/__init__.py create mode 100644 chat/management/commands/runserver.py create mode 100644 chat/settings.py create mode 100644 chat/urls.py create mode 100644 chat/utils.py create mode 100644 chat/websocket.py create mode 100644 gunicorn.conf.py create mode 100755 manage.py create mode 100644 setup_venv.sh diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..e67a52e --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +.env +.idea/ +__pycache__/ +chat/settings_local.py diff --git a/chat/__init__.py b/chat/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/chat/activity.py b/chat/activity.py new file mode 100644 index 0000000..93bf0e1 --- /dev/null +++ b/chat/activity.py @@ -0,0 +1,171 @@ +import os +import sqlite3 +import sys +from contextlib import asynccontextmanager +from pathlib import Path +from tempfile import gettempdir, NamedTemporaryFile + + +import aiosqlite + + +def get_tempdir() -> Path: + tmpdir = Path("/dev/shm") + if not tmpdir.exists(): + tmpdir = Path(gettempdir()) + assert tmpdir.exists() + return tmpdir + + +class ActivityManager: + tmp_ctx: NamedTemporaryFile + ENVIRON_KEY = "_ACTIVITY_CTX" + worker_id: int + + @classmethod + def setup(cls): + cls.tmp_ctx = NamedTemporaryFile( + suffix=".db", prefix="activity", dir=get_tempdir(), delete_on_close=False + ) + cls.tmp_ctx.__enter__().close() + os.environ[cls.ENVIRON_KEY] = tmp_path = cls.tmp_ctx.name + with sqlite3.connect(tmp_path) as db: + db.execute("PRAGMA journal_mode = WAL") + db.executescript( + """ + CREATE TABLE worker (id INTEGER PRIMARY KEY, pid INTEGER) STRICT; + CREATE TABLE websocket ( + id INTEGER PRIMARY KEY, + worker_id INTEGER, + user_id INTEGER, + FOREIGN KEY(worker_id) REFERENCES worker(id) + ) STRICT; + CREATE TABLE user_status ( + id INTEGER PRIMARY KEY, + user_id INTEGER UNIQUE, + status TEXT + ) STRICT; + """ + ) + db.commit() + + @classmethod + def cleanup(cls): + tmp_path = cls.tmp_ctx.name + for suffix in ("-shm", "-wal"): + Path(f"{tmp_path}{suffix}").unlink(missing_ok=True) + cls.tmp_ctx.__exit__(*sys.exc_info()) + + def __init__(self): + self.db = None + + @asynccontextmanager + async def db_block(self): + try: + yield + await self.db.commit() + except Exception as original_error: + try: + await self.db.rollback() + except Exception as rollback_error: + raise ExceptionGroup( + "Transaction and rollback both failed", + [original_error, rollback_error], + ) from original_error + raise + + async def insert_returning_id(self, sql, params): + async with await self.db.execute(sql, params) as cursor: + return cursor.lastrowid + + async def execute_fetchall(self, sql, params, cursor=None): + if cursor is None: + async with await self.db.execute(sql, params) as cursor: + return await cursor.fetchall() + await cursor.execute(sql, params) + return await cursor.fetchall() + + async def __aenter__(self): + self.db = await aiosqlite.connect(os.environ[self.ENVIRON_KEY]) + self.db.row_factory = aiosqlite.Row + await self.db.execute("PRAGMA synchronous=NORMAL") + await self.db.execute("PRAGMA foreign_keys=ON") + await self.db.execute("PRAGMA busy_timeout=5000") + async with self.db_block(): + self.worker_id = await self.insert_returning_id( + "INSERT INTO worker (pid) VALUES (?)", (os.getpid(),) + ) + return self + + async def register_websocket(self, user_id=0): + async with self.db_block(): + websocket_id = await self.insert_returning_id( + "INSERT INTO websocket (worker_id, user_id) VALUES (?, ?)", + (self.worker_id, user_id), + ) + if user_id: + async with await self.db.cursor() as cursor: + await cursor.execute( + ( + "INSERT OR IGNORE INTO user_status (user_id, status) " + "VALUES (?, ?)" + ), + (user_id, "online"), + ) + return websocket_id + + async def update_websocket_status(self, websocket_id, new_status): + async with self.db_block(), await self.db.cursor() as cursor: + await cursor.execute( + ( + "UPDATE user_status SET status=? WHERE user_id=(" + " SELECT user_id FROM websocket WHERE worker_id=? AND id=?" + ")" + ), + (new_status, self.worker_id, websocket_id), + ) + + async def get_websocket_status(self, user_ids): + p = ", ".join("?" for _ in range(len(user_ids))) + return await self.execute_fetchall( + f"SELECT user_id, status FROM user_status WHERE user_id IN ({p})", + user_ids, + ) + + async def cleanup_user(self, cursor, user_id): + if not await self.execute_fetchall( + "SELECT id FROM websocket WHERE user_id=?", + (user_id,), + cursor, + ): + await cursor.execute( + "DELETE FROM user_status WHERE user_id=?", + (user_id,) + ) + + async def unregister_websocket(self, websocket_id, user_id): + async with self.db_block(), await self.db.cursor() as cursor: + await cursor.execute("DELETE FROM websocket WHERE id=?", (websocket_id,)) + if user_id: + await self.cleanup_user(cursor, user_id) + + async def __aexit__(self, et, exc, tb): + try: + if not self.worker_id: + return + async with self.db_block(), await self.db.cursor() as cursor: + for row in await self.execute_fetchall( + "SELECT user_id FROM websocket WHERE worker_id=?", + (self.worker_id,), + cursor, + ): + await self.cleanup_user(cursor, row[0]) + await cursor.execute( + "DELETE FROM websocket WHERE worker_id=?", (self.worker_id,) + ) + await cursor.execute( + "DELETE FROM worker WHERE id=?", (self.worker_id,) + ) + finally: + await self.db.close() + self.db = None diff --git a/chat/apps.py b/chat/apps.py new file mode 100644 index 0000000..194dfde --- /dev/null +++ b/chat/apps.py @@ -0,0 +1,5 @@ +from django.apps import AppConfig + + +class ChatConfig(AppConfig): + name = "chat" diff --git a/chat/asgi.py b/chat/asgi.py new file mode 100644 index 0000000..91bfbb8 --- /dev/null +++ b/chat/asgi.py @@ -0,0 +1,37 @@ +""" +ASGI config for chat project. + +It exposes the ASGI callable as a module-level variable named ``application``. + +For more information on this file, see +https://docs.djangoproject.com/en/6.1/howto/deployment/asgi/ +""" + +import os + +from django import setup +from django.core.handlers.asgi import ASGIHandler + +from .lifespan import LifespanManager +from .websocket import Websocket + +os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'chat.settings') + + +class ChatASGIHandler(ASGIHandler): + async def __call__(self, scope, receive, send): + match scope["type"]: + case "lifespan": + await LifespanManager().serve(receive, send) + case "websocket": + await Websocket().serve(receive, send) + case "http": + await super().__call__(scope, receive, send) + + +def get_asgi_application(): + setup(set_prefix=False) + return ChatASGIHandler() + + +application = get_asgi_application() diff --git a/chat/lifespan.py b/chat/lifespan.py new file mode 100644 index 0000000..9d409c7 --- /dev/null +++ b/chat/lifespan.py @@ -0,0 +1,102 @@ +from asyncio import CancelledError, Queue, sleep +from asyncpg import create_pool +from contextlib import AsyncExitStack +from typing import TypedDict + +from django.conf import settings + +from .activity import ActivityManager +from .utils import CancellingTaskGroup +from .websocket import StopSend + + +class LifespanResponse(TypedDict): + type: str | None + + +class LifespanManager: + queue = Queue() + websockets = {} + + def __init__(self): + self.activity_manager = ActivityManager() + + @staticmethod + async def send_pending_msg(send, pending_msg): + if pending_msg["type"] is None: + return + await send(pending_msg) + pending_msg["type"] = None + + async def serve(self, receive, send): + pending_msg: LifespanResponse = {"type": None} + try: + async with AsyncExitStack() as exit_stack: + while (await receive())["type"] != "lifespan.startup": + pass + pending_msg["type"] = "lifespan.startup.complete" + await exit_stack.enter_async_context(self.activity_manager) + await self.setup(exit_stack) + await self.send_pending_msg(send, pending_msg) + while msg := await receive(): + if msg["type"] == "lifespan.shutdown": + break + elif msg["type"] == "lifespan.startup": + pending_msg["type"] = "lifespan.startup.failed" + raise ValueError("Unexpectedly repeating lifespan.startup") + pending_msg["type"] = "lifespan.shutdown.complete" + except CancelledError: + raise + except Exception: + pending_msg_type = pending_msg["type"] + if isinstance(pending_msg_type, str): + pending_msg["type"] = pending_msg_type.replace(".complete", ".failed") + raise + finally: + await self.send_pending_msg(send, pending_msg) + + async def setup(self, exit_stack): + creds = settings.DATABASES["default"] + pool = await exit_stack.enter_async_context( + create_pool( + host=creds["HOST"], + port=creds["PORT"], + user=creds["USER"], + password=creds["PASSWORD"], + database=creds["NAME"], + ) + ) + task_group = await exit_stack.enter_async_context(CancellingTaskGroup()) + create_task( + self.consume_websocket_queue(pool) + ) + task_group.create_task(self.consume_channel(pool, "chat_channel")) + + async def consume_websocket_queue(self, pool): + # websockets -> pg rows (+ triggers) and activity + async with pool.acquire() as con: + print("websocket queue") + try: + while item := await self.queue.get(): + websocket = self.websockets[item["websocket"]] + match item["type"]: + case "register": + websocket.websocket_id = ( + await self.activity_manager.register_websocket() + ) + case "unregister": + await self.activity_manager.unregister_websocket( + websocket.websocket_id, websocket.user_id + ) + websocket.send_queue.put(StopSend) + except CancelledError: + print("cancelled consume_websocket_queue") + raise + finally: + print("websocket queue done") + + async def consume_channel(self, pool, channel_name): + # pg channel -> websockets + async with pool.acquire() as con: + while True: + await sleep(1) diff --git a/chat/management/__init__.py b/chat/management/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/chat/management/commands/__init__.py b/chat/management/commands/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/chat/management/commands/runserver.py b/chat/management/commands/runserver.py new file mode 100644 index 0000000..10a926f --- /dev/null +++ b/chat/management/commands/runserver.py @@ -0,0 +1,34 @@ +import sys +from contextlib import contextmanager +from typing import Any + +from django.core.management import BaseCommand, CommandParser +from gunicorn.app.wsgiapp import WSGIApplication + +from chat.utils import replace_attr + + +class Command(BaseCommand): + def add_arguments(self, parser: CommandParser) -> None: + parser.add_argument( + "gunicorn_args", + nargs="*", + help="Arguments passed directly to gunicorn", + ) + + def handle(self, *args: Any, **options: Any) -> str | None: + argv = [ + "gunicorn", + "chat.asgi:application", + "--bind", + "127.0.0.1:8000", + "-k", + "uvicorn_worker.UvicornWorker", + "-w", + "3", + "-t", + "60", + *options["gunicorn_args"], + ] + with replace_attr(sys, "argv", argv): + WSGIApplication().run() diff --git a/chat/settings.py b/chat/settings.py new file mode 100644 index 0000000..d72d144 --- /dev/null +++ b/chat/settings.py @@ -0,0 +1,136 @@ +""" +Django settings for chat project. + +Generated by 'django-admin startproject' using Django 6.1. + +For more information on this file, see +https://docs.djangoproject.com/en/6.1/topics/settings/ + +For the full list of settings and their values, see +https://docs.djangoproject.com/en/6.1/ref/settings/ +""" + +import os +from pathlib import Path + +BASE_DIR = Path(__file__).resolve().parents[1] +SECRET_KEY = os.getenv( + "SECRET_KEY", + "django-insecure-i*_h#mx)#md&_v6+hsuin4qxx%a9c0z#234h6!_%tcykoo1gyv", +) +DEBUG = os.getenv("DEBUG", "1").lower() in ("1", "t", "true") +ALLOWED_HOSTS = os.getenv("ALLOWED_HOSTS", "localhost").split(",") + +INSTALLED_APPS = [ + "chat", + "django.contrib.admin", + "django.contrib.auth", + "django.contrib.contenttypes", + "django.contrib.sessions", + "django.contrib.messages", + "django.contrib.staticfiles", +] + +MIDDLEWARE = [ + "django.middleware.security.SecurityMiddleware", + "django.contrib.sessions.middleware.SessionMiddleware", + "django.middleware.common.CommonMiddleware", + "django.middleware.csrf.CsrfViewMiddleware", + "django.contrib.auth.middleware.AuthenticationMiddleware", + "django.contrib.messages.middleware.MessageMiddleware", + "django.middleware.clickjacking.XFrameOptionsMiddleware", +] + +ROOT_URLCONF = "chat.urls" + +TEMPLATES = [ + { + "BACKEND": "django.template.backends.django.DjangoTemplates", + "DIRS": [], + "APP_DIRS": True, + "OPTIONS": { + "context_processors": [ + "django.template.context_processors.request", + "django.contrib.auth.context_processors.auth", + "django.contrib.messages.context_processors.messages", + ], + }, + }, +] + +ASGI_APPLICATION = "chat.asgi.application" + +DATABASES = { + "default": { + "ENGINE": "django.db.backends.postgresql", + "NAME": "chat", + "HOST": os.environ.get("PGHOST", "localhost"), + "PORT": os.environ.get("PGPORT", "5432"), + "USER": os.environ.get("PGUSER", "chat"), + "PASSWORD": os.environ.get("PGPASS", ""), + } +} + +AUTH_PASSWORD_VALIDATORS = [ + { + "NAME": ( + "django.contrib.auth.password_validation.UserAttributeSimilarityValidator" + ), + }, + {"NAME": "django.contrib.auth.password_validation.MinimumLengthValidator"}, + {"NAME": "django.contrib.auth.password_validation.CommonPasswordValidator"}, + {"NAME": "django.contrib.auth.password_validation.NumericPasswordValidator"}, + { + "NAME": "django.contrib.auth.password_validation.MinimumLengthValidator", + "OPTIONS": {"min_length": 9}, + }, + {"NAME":"user.password_validation.ComplexityValidator"}, +] + +# Internationalization +# https://docs.djangoproject.com/en/6.1/topics/i18n/ + +LANGUAGE_CODE = "en-us" +TIME_ZONE = "UTC" +USE_I18N = True +USE_TZ = True +LOCALE_PATHS = [BASE_DIR / "locales"] + +STATIC_URL = "static/" +STATIC_ROOT = BASE_DIR / "static" +MEDIA_URL = "media/" +MEDIA_ROOT = BASE_DIR / "media" +LOGIN_URL = "/user/login/" + +_use_tls = os.getenv("SMTP_CONN") == "SMTP_SSL" +MAILERS = { + "default": { + "BACKEND": "django.core.mail.backends.smtp.EmailBackend", + "OPTIONS": { + "host": os.getenv("SMTP_HOST") or "localhost", + "port": int(os.getenv("SMTP_PORT") or "25" if _use_tls else "587"), + "use_tls": _use_tls, + "username": os.getenv("SMTP_USER") or "", + "password": os.getenv("SMTP_PASSWORD") or "", + }, + }, +} + +DEFAULT_FROM_EMAIL = os.getenv("DEFAULT_FROM_EMAIL", f"noreply@{ALLOWED_HOSTS[0]}") + +LOGGING = { + "version": 1, + "disable_existing_loggers": False, + "handlers": { + "console": {"class": "logging.StreamHandler"}, + }, + "root": { + "handlers": ["console"], + "level": "INFO", + }, +} + +try: + from .settings_local import * +except ImportError: + pass diff --git a/chat/urls.py b/chat/urls.py new file mode 100644 index 0000000..69ffe9a --- /dev/null +++ b/chat/urls.py @@ -0,0 +1,22 @@ +""" +URL configuration for chat project. + +The `urlpatterns` list routes URLs to views. For more information please see: + https://docs.djangoproject.com/en/6.1/topics/http/urls/ +Examples: +Function views + 1. Add an import: from my_app import views + 2. Add a URL to urlpatterns: path('', views.home, name='home') +Class-based views + 1. Add an import: from other_app.views import Home + 2. Add a URL to urlpatterns: path('', Home.as_view(), name='home') +Including another URLconf + 1. Import the include() function: from django.urls import include, path + 2. Add a URL to urlpatterns: path('blog/', include('blog.urls')) +""" +from django.contrib import admin +from django.urls import path + +urlpatterns = [ + path('admin/', admin.site.urls), +] diff --git a/chat/utils.py b/chat/utils.py new file mode 100644 index 0000000..4b34279 --- /dev/null +++ b/chat/utils.py @@ -0,0 +1,34 @@ +from asyncio import TaskGroup +from contextlib import contextmanager + + +class CancellingTaskGroup: + def __init__(self): + self.tg = TaskGroup() + self.tasks = [] + + async def __aenter__(self): + await self.tg.__aenter__() + return self + + def create_task(self, task): + self.tasks.append(self.tg.create_task(task)) + + async def __aexit__(self, et, exc, tb): + for task in self.tasks: + task.cancel() + result = await self.tg.__aexit__(et, exc, tb) + return result + +Unset = object() + + +@contextmanager +def replace_attr(obj, attr, new_value): + old_value = getattr(obj, attr, Unset) + setattr(obj, attr, new_value) + yield + if old_value is Unset: + delattr(obj, attr) + else: + setattr(obj, attr, old_value) diff --git a/chat/websocket.py b/chat/websocket.py new file mode 100644 index 0000000..2568912 --- /dev/null +++ b/chat/websocket.py @@ -0,0 +1,67 @@ +import json +from asyncio import CancelledError, Queue, create_task, shield, wait_for +from contextlib import suppress + +from django.contrib.auth.models import AnonymousUser + +StopSend = object() + + +class Websocket: + # keep it dumb, we'd like to run many of these concurrently + def __init__(self): + self.websocket_id = -1 + self.user = AnonymousUser + self.send_queue = Queue() + + async def serve(self, send, receive): + from .lifespan import LifespanManager + + disconnected = False + send_task = None + try: + await send({"type": "websocket.accept"}) + send_task = create_task(self.serve_send_queue(send)) + LifespanManager.websockets[id(self)] = self + await LifespanManager.queue.put({"type": "register", "websocket": id(self)}) + while True: + event = await receive() + if event["type"] == "websocket.receive": + text = event.get("text") + if text is not None: + await self.send_queue.put(json.loads(text)) + elif event["type"] == "websocket.disconnect": + disconnected = True + break + except CancelledError: + if disconnected: + raise + if send_task: + send_task.cancel() + with suppress(CancelledError): + await send_task + with suppress(Exception): + await wait_for( + shield(send({"type": "websocket.close", "code": 1013})), 10 + ) + raise + finally: + await LifespanManager.queue.put( + {"type": "unregister", "websocket": id(self)} + ) + if send_task: + await send_task + + async def serve_send_queue(self, send): + while True: + item = await self.send_queue.get() + if isinstance(item, dict): + await send( + { + "type": "websocket.send", + "text": json.dumps(item, ensure_ascii=False), + } + ) + self.send_queue.task_done() + if item is StopSend: + return diff --git a/gunicorn.conf.py b/gunicorn.conf.py new file mode 100644 index 0000000..3a3f292 --- /dev/null +++ b/gunicorn.conf.py @@ -0,0 +1,9 @@ +from chat.activity import ActivityManager + + +def on_starting(_): + ActivityManager.setup() + + +def on_exit(_): + ActivityManager.cleanup() diff --git a/manage.py b/manage.py new file mode 100755 index 0000000..60eb248 --- /dev/null +++ b/manage.py @@ -0,0 +1,22 @@ +#!/usr/bin/env python +"""Django's command-line utility for administrative tasks.""" +import os +import sys + + +def main(): + """Run administrative tasks.""" + os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'chat.settings') + try: + from django.core.management import execute_from_command_line + except ImportError as exc: + raise ImportError( + "Couldn't import Django. Are you sure it's installed and " + "available on your PYTHONPATH environment variable? Did you " + "forget to activate a virtual environment?" + ) from exc + execute_from_command_line(sys.argv) + + +if __name__ == '__main__': + main() diff --git a/setup_venv.sh b/setup_venv.sh new file mode 100644 index 0000000..f2dcd6e --- /dev/null +++ b/setup_venv.sh @@ -0,0 +1,21 @@ +#!/usr/bin/env bash + +set -a +PYTHON="${PYTHON:-python}" +"${PYTHON}" -m venv .venv + +if [[ -r .env ]]; then + . .env +fi + +CFLAGS="$( + python -c 'import sysconfig; print(sysconfig.get_config_var("CFLAGS"))' +) -O2" +set +a +pip_args=(-U --no-binary :all:) +# avoid attempting network connections when the network is measurably down +if [[ -z "$(ip route show default)" ]]; then + pip_args+=(--no-index --no-build-isolation) +fi +PIP_RETRIES=2 PIP_TIMEOUT=2 ".venv/bin/${PYTHON%%*/}" -m pip install "${pip_args[@]}" \ + aiosqlite asyncpg django django-pgtrigger gunicorn pip 'psycopg[c]' uvicorn-worker -- 2.55.0