--- /dev/null
+.env
+.idea/
+__pycache__/
+chat/settings_local.py
--- /dev/null
+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
--- /dev/null
+from django.apps import AppConfig
+
+
+class ChatConfig(AppConfig):
+ name = "chat"
--- /dev/null
+"""
+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()
--- /dev/null
+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)
--- /dev/null
+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()
--- /dev/null
+"""
+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
--- /dev/null
+"""
+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),
+]
--- /dev/null
+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)
--- /dev/null
+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
--- /dev/null
+from chat.activity import ActivityManager
+
+
+def on_starting(_):
+ ActivityManager.setup()
+
+
+def on_exit(_):
+ ActivityManager.cleanup()
--- /dev/null
+#!/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()
--- /dev/null
+#!/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
projects / chat / commitdiff