From 8c31cbc5b2726749074874414357bc0eb131f78a Mon Sep 17 00:00:00 2001 From: Interitio Date: Tue, 15 Sep 2026 22:14:44 +1000 Subject: [PATCH] feat: Add batching end-of-focus notification --- hyperfocus/twitch/component.py | 162 ++++++++++++++++++++++++++++++++- 1 file changed, 160 insertions(+), 2 deletions(-) diff --git a/hyperfocus/twitch/component.py b/hyperfocus/twitch/component.py index 041350b..59b34f9 100644 --- a/hyperfocus/twitch/component.py +++ b/hyperfocus/twitch/component.py @@ -2,10 +2,11 @@ import asyncio from collections import defaultdict from string import punctuation import datetime as dt -from datetime import datetime, timedelta +from datetime import date, datetime, timedelta from math import ceil import time import random +from typing import Optional import twitchio from twitchio import Scopes @@ -46,7 +47,85 @@ responses = """ """.splitlines() +class FocusNotifier: + def __init__(self, channel: twitchio.PartialUser, send_as: str): + # Users which have their focus end in the next 10s + self.prenotify_queue: dict[str, datetime] = {} + + # Users which have already had their focus end, waiting for prenotified + self.notify_queue: list[str] = [] + + # Task which moves prenotify to notify, and either consumes and notifies + # or sleeps and reschedules + self.notify_task: Optional[asyncio.Task] = None + + # Channel which we can send to, to notify + self.channel: twitchio.PartialUser = channel + self.send_as = send_as + + # TODO: Locking? Ensure we never lose a notify task? + + def focusendingsoon(self, username: str, focusends: datetime): + logger.debug(f"Prenotify for {username=} at {focusends=}") + self.prenotify_queue[username] = focusends + self.schedule_notify() + + def schedule_notify(self): + # Idempotently schedule notification + if ( + self.notify_task is None + or self.notify_task.cancelled() + or self.notify_task.done() + ): + self.notify_task = asyncio.create_task(self._notify()) + + async def _notify(self): + # TODO: Logging + # The notify task + # Move all expired prenotify to notify queue + logger.debug( + f"Scheduling notify task: {self.prenotify_queue=} {self.notify_queue=}" + ) + now = utc_now() + for username, ending in list(self.prenotify_queue.items()): + if ending <= now: + self.prenotify_queue.pop(username) + self.notify_queue.append(username) + + # If there are any prenotifies left + if self.prenotify_queue: + # Wait until troubleshooting prenotify + 1 and rerun + sleep_until = max(self.prenotify_queue.values()) + dt.timedelta( + milliseconds=500 + ) + await asyncio.sleep((sleep_until - now).total_seconds()) + + # This is sync, no possible race conditions + self.notify_task = None + self.schedule_notify() + else: + # If there aren't + # Set the notify_task to None to indicate we are done waiting + # Then any new prenotifies will schedule a new one + self.notify_task = None + # Consume notify queue and actually notify + if self.notify_queue: + names = tuple(self.notify_queue) + self.notify_queue.clear() + await self._actually_notify(*names) + + async def _actually_notify(self, *usernames): + # Send message to channel with given names + namestr = ", ".join(usernames) + message = f"🚨 {namestr} your HYPERFOCUS lockdown is complete. Nico approves. kawaii16Hype ✨" + await self.channel.send_message(message, self.send_as) + + class FocusComponent(cmds.Component): + # If user's focus sessions end within ~10 seconds of each other, notify at the same time + _notify_buffer = 10 + _last_seen_threshold = 4 * 3600 + def __init__(self, bot: Bot): self.bot = bot self.data = bot.dbconn.load_registry(HyperfocusData()) @@ -56,12 +135,82 @@ class FocusComponent(cmds.Component): self._last_seen: dict[int, dict[int, datetime]] = defaultdict(dict) self._last_deleted: dict[int, datetime] = {} + + # Map of userid -> (focus end, focus preending task) + self._focus_preending: dict[int, asyncio.Task] = {} + self._channel_notifiers: dict[int, FocusNotifier] = {} + self._cid_channel_cache: dict[int, twitchio.PartialUser] = {} + self.hyperfocus_lock = asyncio.Lock() # ----- API ----- async def component_load(self): await self.data.init() + def schedule_focusend_notify( + self, focuser: Hyperfocuser, initial_cid: int | None = None + ): + # Cancel current focus preend if it exists + self.cancel_focusend_notify(focuser.profileid) + + # Schedule a new one + task = asyncio.create_task(self._preend_notify_task(focuser, initial_cid)) + self._focus_preending[focuser.profileid] = task + + def cancel_focusend_notify(self, profileid): + maybetask = self._focus_preending.pop(profileid, None) + if maybetask and not maybetask.done() and not maybetask.cancelled(): + maybetask.cancel() + + async def _preend_notify_task( + self, focuser: Hyperfocuser, initial_cid: int | None = None + ): + logger.debug(f"Scheduling prenotify for {focuser=}") + prenotify_at = focuser.ends_at - dt.timedelta(seconds=self._notify_buffer) + now = utc_now() + if prenotify_at > now: + await asyncio.sleep((prenotify_at - now).total_seconds()) + + now = utc_now() + + # In case the focuser was modified by another client + profileid = focuser.profileid + profile = await self.bot.profiles.profiles.get_profile(profileid) + maybefocuser = await self.get_hyperfocus(profileid) + + # Note this is all sync past here, so no danger of race condition with cancelling the notify + self._focus_preending.pop(profileid) + if not maybefocuser: + # Focus ended prematurely, so we do as well + return + if maybefocuser.ends_at - dt.timedelta(seconds=self._notify_buffer) >= now: + self.schedule_focusend_notify(maybefocuser, initial_cid=initial_cid) + else: + # Trigger the prenotify + # (it will go through anyway) + logger.debug(f"Prenotifying for {maybefocuser=}") + targets = set() + if initial_cid is not None: + targets.add(initial_cid) + + last_seen_cids = self._last_seen[profileid].items() + for cid, last_seen in last_seen_cids: + if (now - last_seen).total_seconds() < self._last_seen_threshold: + targets.add(cid) + logger.debug(f"Prenotifying for {maybefocuser=} in {targets=}") + for cid in targets: + notifier = self._channel_notifiers.get(cid) + if notifier is None: + channel = self._cid_channel_cache.get(cid) + if channel is not None: + notifier = FocusNotifier(channel, send_as=self.bot.bot_id) + self._channel_notifiers[cid] = notifier + if notifier is not None: + notifier.focusendingsoon( + username=profile.nickname or "Unknown", + focusends=maybefocuser.ends_at, + ) + async def dispatch_focuser_update( self, profileid: int, @@ -85,14 +234,18 @@ class FocusComponent(cmds.Component): now = utc_now() if focuser is not None: + self.schedule_focusend_notify(focuser, initial_cid=communityid) + # If we are sending active, send to all last seens newer than a certain date, and delete any old ones last_seen_cids = self._last_seen[profileid].items() for cid, last_seen in last_seen_cids: - if (now - last_seen).total_seconds() < 4 * 3600: + if (now - last_seen).total_seconds() < self._last_seen_threshold: targets.add(cid) for cid in targets: await self.channel.send_hyperfocus_patch(cid, focuser) else: + self.cancel_focusend_notify(profileid) + # If we are deleting, send to *all* last seens, and then delete any old ones # targets.update(self._last_seen[profileid].keys()) # for cid in targets: @@ -153,6 +306,9 @@ class FocusComponent(cmds.Component): now = utc_now() self._last_seen[profile.profileid][comm.communityid] = now + self._cid_channel_cache[comm.communityid] = ( + payload.broadcaster + ) # TODO: Better last seen channel info hyperfocused = await self.get_hyperfocus(profile.profileid) @@ -196,6 +352,7 @@ class FocusComponent(cmds.Component): if hyperfocused: # Send an update to the channel await self.channel.send_hyperfocus_patch(comm.communityid, hyperfocused) + self.schedule_focusend_notify(hyperfocused, initial_cid=comm.communityid) # ------ Commands ----- @cmds.command( @@ -256,6 +413,7 @@ class FocusComponent(cmds.Component): async with self.hyperfocus_lock: row = await Hyperfocuser.fetch(profile.profileid) if row: + self.cancel_focusend_notify(profile.profileid) await row.delete() await self.channel.send_hyperfocus_del(profile.profileid)