From cb538de3151d8e3f64bc9a8c60da33da0db289d0 Mon Sep 17 00:00:00 2001 From: mar77i Date: Sun, 27 Sep 2026 12:51:54 +0200 Subject: [PATCH] integrate activity tracking on postgres --- activity/__init__.py | 0 activity/apps.py | 5 ++ activity/migrations/0001_initial.py | 36 +++++++++ activity/migrations/0002_initial.py | 33 +++++++++ activity/migrations/__init__.py | 0 activity/models.py | 46 ++++++++++++ chat/async_db.py | 55 ++++++++++++++ chat/lifespan.py | 111 ++++++++++++++++------------ chat/settings.py | 3 + chat/websocket.py | 38 +++++----- gunicorn.conf.py | 6 +- setup_venv.sh | 6 +- user/__init__.py | 0 user/apps.py | 5 ++ user/migrations/0001_initial.py | 44 +++++++++++ user/migrations/__init__.py | 0 user/models.py | 5 ++ 17 files changed, 323 insertions(+), 70 deletions(-) create mode 100644 activity/__init__.py create mode 100644 activity/apps.py create mode 100644 activity/migrations/0001_initial.py create mode 100644 activity/migrations/0002_initial.py create mode 100644 activity/migrations/__init__.py create mode 100644 activity/models.py create mode 100644 chat/async_db.py create mode 100644 user/__init__.py create mode 100644 user/apps.py create mode 100644 user/migrations/0001_initial.py create mode 100644 user/migrations/__init__.py create mode 100644 user/models.py diff --git a/activity/__init__.py b/activity/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/activity/apps.py b/activity/apps.py new file mode 100644 index 0000000..34cfcc0 --- /dev/null +++ b/activity/apps.py @@ -0,0 +1,5 @@ +from django.apps import AppConfig + + +class ActivityConfig(AppConfig): + name = 'activity' diff --git a/activity/migrations/0001_initial.py b/activity/migrations/0001_initial.py new file mode 100644 index 0000000..cfea456 --- /dev/null +++ b/activity/migrations/0001_initial.py @@ -0,0 +1,36 @@ +# Generated by Django 6.1.1 on 2026-09-27 10:49 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + initial = True + + dependencies = [ + ] + + operations = [ + migrations.CreateModel( + name='Activity', + fields=[ + ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), + ('status', models.CharField(max_length=256)), + ], + ), + migrations.CreateModel( + name='Websocket', + fields=[ + ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), + ], + ), + migrations.CreateModel( + name='Worker', + fields=[ + ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), + ('hostname', models.CharField(max_length=256)), + ('pid', models.IntegerField()), + ('starttime', models.IntegerField()), + ], + ), + ] diff --git a/activity/migrations/0002_initial.py b/activity/migrations/0002_initial.py new file mode 100644 index 0000000..2905784 --- /dev/null +++ b/activity/migrations/0002_initial.py @@ -0,0 +1,33 @@ +# Generated by Django 6.1.1 on 2026-09-27 10:49 + +import django.db.models.deletion +from django.conf import settings +from django.db import migrations, models + + +class Migration(migrations.Migration): + + initial = True + + dependencies = [ + ('activity', '0001_initial'), + migrations.swappable_dependency(settings.AUTH_USER_MODEL), + ] + + operations = [ + migrations.AddField( + model_name='websocket', + name='user', + field=models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to=settings.AUTH_USER_MODEL), + ), + migrations.AddField( + model_name='activity', + name='websockets', + field=models.ManyToManyField(related_name='activity', to='activity.websocket'), + ), + migrations.AddField( + model_name='websocket', + name='worker', + field=models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='activity.worker'), + ), + ] diff --git a/activity/migrations/__init__.py b/activity/migrations/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/activity/models.py b/activity/models.py new file mode 100644 index 0000000..3c9c7fd --- /dev/null +++ b/activity/models.py @@ -0,0 +1,46 @@ +from pathlib import Path +from socket import gethostname + +from django.db.models import ( + CASCADE, CharField, ForeignKey, IntegerField, ManyToManyField, Model +) + +from user.models import User + + +class Worker(Model): + hostname = CharField(max_length=256) + pid = IntegerField() + starttime = IntegerField() + + @staticmethod + def get_starttime(pid: int) -> int | None: + path = Path(f"/proc/{pid}/stat") + if not path.parent.exists(): + return None + line = path.read_bytes() + pos = line.rfind(b") ") + if pos == -1: + raise ValueError(f"Malformed /proc/{pid}/stat") + return int(line[pos + 2:].split(maxsplit=20)[19]) + + @classmethod + def cleanup(cls): + cls.objects.filter( + pk__in=[ + worker.pk + for worker in cls.objects.filter(hostname=gethostname()) + if worker.starttime != cls.get_starttime(worker.pid) + ] + ).delete() + return cls.objects.exists() + + +class Websocket(Model): + worker = ForeignKey(Worker, on_delete=CASCADE) + user = ForeignKey(User, on_delete=CASCADE) + + +class Activity(Model): + websockets = ManyToManyField(Websocket, related_name="activity") + status = CharField(max_length=256) diff --git a/chat/async_db.py b/chat/async_db.py new file mode 100644 index 0000000..b37e613 --- /dev/null +++ b/chat/async_db.py @@ -0,0 +1,55 @@ +from django.db import models + + +async def ainsert(con, instance: models.Model, using="default"): + model = type(instance) + concrete_fields = [ + field + for field in model._meta.concrete_fields + if not field.auto_created + ] + # Include the primary key only when the instance has an explicit value. + pk_is_none = getattr(instance, model._meta.pk.attname) is None + insert_columns = [] + insert_values = [] + returning_columns = [] + field_names = [] + for field in concrete_fields: + if not field.primary_key or not pk_is_none: + insert_columns.append(field.column) + insert_values.append(getattr(instance, field.attname)) + returning_columns.append(field.column) + field_names.append(field.attname) + + columns_sql = ", ".join(f'"{column}"' for column in insert_columns) + placeholders_sql = ", ".join(f"${index + 1}" for index in range(len(insert_values))) + returning_sql = ", ".join( + f'"{column}"' for column in returning_columns + ) + row = await con.fetchrow( + ( + f'INSERT INTO "{model._meta.db_table}" ({columns_sql}) ' + f'VALUES ({placeholders_sql}) RETURNING {returning_sql}' + ), + *insert_values, + ) + if row is None: + raise RuntimeError("INSERT did not return a row") + loaded_instance = model.from_db( + using, + field_names, + [row[field.column] for field in concrete_fields], + ) + instance.__dict__.update(loaded_instance.__dict__) + return instance + + +async def adelete(con, instance: models.Model): + model = type(instance) + if instance._state.adding: + raise ValueError("Cannot delete an unsaved instance") + await con.execute( + f'DELETE FROM "{model._meta.db_table}" WHERE "{model._meta.pk.column}"=$1', + instance.pk, + ) + return None diff --git a/chat/lifespan.py b/chat/lifespan.py index 9d409c7..4e1a6e9 100644 --- a/chat/lifespan.py +++ b/chat/lifespan.py @@ -1,3 +1,6 @@ +import os +from socket import gethostname + from asyncio import CancelledError, Queue, sleep from asyncpg import create_pool from contextlib import AsyncExitStack @@ -5,10 +8,12 @@ from typing import TypedDict from django.conf import settings -from .activity import ActivityManager +from .async_db import adelete, ainsert from .utils import CancellingTaskGroup from .websocket import StopSend +from activity.models import Worker, Websocket + class LifespanResponse(TypedDict): type: str | None @@ -19,7 +24,7 @@ class LifespanManager: websockets = {} def __init__(self): - self.activity_manager = ActivityManager() + self.worker_instance = None @staticmethod async def send_pending_msg(send, pending_msg): @@ -30,47 +35,56 @@ class LifespanManager: 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 with create_pool( + host=creds["HOST"], + port=creds["PORT"], + user=creds["USER"], + password=creds["PASSWORD"], + database=creds["NAME"], + ) as pool: + try: + async with AsyncExitStack() as exit_stack: + while (await receive())["type"] != "lifespan.startup": + pass + pending_msg["type"] = "lifespan.startup.complete" + async with pool.acquire() as con: + pid = os.getpid() + self.worker_instance = await ainsert( + con, + Worker( + hostname=gethostname(), + pid=pid, + starttime=Worker.get_starttime(pid), + ), + ) + task_group = await exit_stack.enter_async_context( + CancellingTaskGroup() + ) + task_group.create_task(self.consume_websocket_queue(pool)) + task_group.create_task(self.consume_channel(pool, "chat_channel")) + 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) + if isinstance(self.worker_instance, Worker): + async with pool.acquire() as con: + self.worker_instance = await adelete(con, self.worker_instance) async def consume_websocket_queue(self, pool): # websockets -> pg rows (+ triggers) and activity @@ -81,12 +95,17 @@ class LifespanManager: websocket = self.websockets[item["websocket"]] match item["type"]: case "register": - websocket.websocket_id = ( - await self.activity_manager.register_websocket() + websocket.websocket_instance = ( + await ainsert( + con, + Websocket( + worker=self.worker_instance, user=None + ), + ) ) case "unregister": - await self.activity_manager.unregister_websocket( - websocket.websocket_id, websocket.user_id + websocket.websocket_instance = await adelete( + con, websocket.websocket_instance ) websocket.send_queue.put(StopSend) except CancelledError: diff --git a/chat/settings.py b/chat/settings.py index d72d144..9200742 100644 --- a/chat/settings.py +++ b/chat/settings.py @@ -22,7 +22,9 @@ DEBUG = os.getenv("DEBUG", "1").lower() in ("1", "t", "true") ALLOWED_HOSTS = os.getenv("ALLOWED_HOSTS", "localhost").split(",") INSTALLED_APPS = [ + "activity", "chat", + "user", "django.contrib.admin", "django.contrib.auth", "django.contrib.contenttypes", @@ -86,6 +88,7 @@ AUTH_PASSWORD_VALIDATORS = [ }, {"NAME":"user.password_validation.ComplexityValidator"}, ] +AUTH_USER_MODEL = "user.User" # Internationalization # https://docs.djangoproject.com/en/6.1/topics/i18n/ diff --git a/chat/websocket.py b/chat/websocket.py index 2568912..c6318eb 100644 --- a/chat/websocket.py +++ b/chat/websocket.py @@ -10,7 +10,7 @@ StopSend = object() class Websocket: # keep it dumb, we'd like to run many of these concurrently def __init__(self): - self.websocket_id = -1 + self.websocket_instance = None self.user = AnonymousUser self.send_queue = Queue() @@ -38,30 +38,30 @@ class Websocket: 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 + await wait_for(send_task, 10) 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), - } + try: + 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: + break + finally: + with suppress(Exception): + await wait_for( + shield(send({"type": "websocket.close", "code": 1013})), 10 ) - self.send_queue.task_done() - if item is StopSend: - return diff --git a/gunicorn.conf.py b/gunicorn.conf.py index 3a3f292..060fea5 100644 --- a/gunicorn.conf.py +++ b/gunicorn.conf.py @@ -1,9 +1,9 @@ -from chat.activity import ActivityManager +from activity.models import Worker def on_starting(_): - ActivityManager.setup() + assert not Worker.cleanup() def on_exit(_): - ActivityManager.cleanup() + Worker.cleanup() diff --git a/setup_venv.sh b/setup_venv.sh index f2dcd6e..cc3a7ff 100644 --- a/setup_venv.sh +++ b/setup_venv.sh @@ -12,10 +12,12 @@ CFLAGS="$( python -c 'import sysconfig; print(sysconfig.get_config_var("CFLAGS"))' ) -O2" set +a + +. .venv/bin/activate 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 +PIP_RETRIES=2 PIP_TIMEOUT=2 "${PYTHON%%*/}" -m pip install "${pip_args[@]}" \ + asyncpg django django-pgtrigger gunicorn pip 'psycopg[c]' uvicorn-worker diff --git a/user/__init__.py b/user/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/user/apps.py b/user/apps.py new file mode 100644 index 0000000..35048d4 --- /dev/null +++ b/user/apps.py @@ -0,0 +1,5 @@ +from django.apps import AppConfig + + +class UserConfig(AppConfig): + name = 'user' diff --git a/user/migrations/0001_initial.py b/user/migrations/0001_initial.py new file mode 100644 index 0000000..c138ea5 --- /dev/null +++ b/user/migrations/0001_initial.py @@ -0,0 +1,44 @@ +# Generated by Django 6.1.1 on 2026-09-27 10:49 + +import django.contrib.auth.models +import django.contrib.auth.validators +import django.utils.timezone +from django.db import migrations, models + + +class Migration(migrations.Migration): + + initial = True + + dependencies = [ + ('auth', '0012_alter_user_first_name_max_length'), + ] + + operations = [ + migrations.CreateModel( + name='User', + fields=[ + ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), + ('password', models.CharField(max_length=128, verbose_name='password')), + ('last_login', models.DateTimeField(blank=True, null=True, verbose_name='last login')), + ('is_superuser', models.BooleanField(default=False, help_text='Designates that this user has all permissions without explicitly assigning them.', verbose_name='superuser status')), + ('username', models.CharField(error_messages={'unique': 'A user with that username already exists.'}, help_text='Required. 150 characters or fewer. Letters, digits and @/./+/-/_ only.', max_length=150, unique=True, validators=[django.contrib.auth.validators.UnicodeUsernameValidator()], verbose_name='username')), + ('first_name', models.CharField(blank=True, max_length=150, verbose_name='first name')), + ('last_name', models.CharField(blank=True, max_length=150, verbose_name='last name')), + ('email', models.EmailField(blank=True, max_length=254, verbose_name='email address')), + ('is_staff', models.BooleanField(default=False, help_text='Designates whether the user can log into this admin site.', verbose_name='staff status')), + ('is_active', models.BooleanField(default=True, help_text='Designates whether this user should be treated as active. Unselect this instead of deleting accounts.', verbose_name='active')), + ('date_joined', models.DateTimeField(default=django.utils.timezone.now, verbose_name='date joined')), + ('groups', models.ManyToManyField(blank=True, help_text='The groups this user belongs to. A user will get all permissions granted to each of their groups.', related_name='user_set', related_query_name='user', to='auth.group', verbose_name='groups')), + ('user_permissions', models.ManyToManyField(blank=True, help_text='Specific permissions for this user.', related_name='user_set', related_query_name='user', to='auth.permission', verbose_name='user permissions')), + ], + options={ + 'verbose_name': 'user', + 'verbose_name_plural': 'users', + 'abstract': False, + }, + managers=[ + ('objects', django.contrib.auth.models.UserManager()), + ], + ), + ] diff --git a/user/migrations/__init__.py b/user/migrations/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/user/models.py b/user/models.py new file mode 100644 index 0000000..3d30525 --- /dev/null +++ b/user/models.py @@ -0,0 +1,5 @@ +from django.contrib.auth.models import AbstractUser + + +class User(AbstractUser): + pass -- 2.55.0