129 lines
4.8 KiB
Python
129 lines
4.8 KiB
Python
import asyncio
|
|
import logging
|
|
from typing import Optional
|
|
|
|
from meta import conf, appname
|
|
from meta.logger import log_context, log_action_stack, logging_context, log_app, log_wrap
|
|
from meta.ipc import AppClient
|
|
from meta.app import appname_from_shard
|
|
from meta.sharding import shard_count
|
|
|
|
from data import Database
|
|
|
|
from .events import command_event_handler, guild_event_handler, voice_event_handler
|
|
from .snapshot import shard_snapshot, ShardSnapshot
|
|
from .data import AnalyticsData
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
for name in conf.config.options('LOGGING_LEVELS', no_defaults=True):
|
|
logging.getLogger(name).setLevel(conf.logging_levels[name])
|
|
|
|
|
|
class AnalyticsServer:
|
|
# TODO: Move these to the config
|
|
# How often to request snapshots
|
|
snap_period = 120
|
|
# How soon after a snapshot failure (e.g. not all shards online) to retry
|
|
snap_retry_period = 10
|
|
|
|
def __init__(self) -> None:
|
|
self.db = Database(conf.data['args'])
|
|
self.data = self.db.load_registry(AnalyticsData())
|
|
|
|
self.event_handlers = [
|
|
command_event_handler,
|
|
guild_event_handler,
|
|
voice_event_handler
|
|
]
|
|
|
|
self.talk = AppClient(
|
|
conf.analytics['appname'],
|
|
appname,
|
|
{'host': conf.analytics['server_host'], 'port': int(conf.analytics['server_port'])},
|
|
{'host': conf.appipc['server_host'], 'port': int(conf.appipc['server_port'])}
|
|
)
|
|
self.talk_shard_snapshot = self.talk.register_route()(shard_snapshot)
|
|
|
|
self._snap_task: Optional[asyncio.Task] = None
|
|
|
|
async def attach_event_handlers(self):
|
|
for handler in self.event_handlers:
|
|
await handler.attach(self.talk)
|
|
|
|
@log_wrap(action='Snap')
|
|
async def take_snapshot(self):
|
|
# Check if all the shards are registered on shard_talk
|
|
expected_peers = [appname_from_shard(i) for i in range(0, shard_count)]
|
|
if missing := [peer for peer in expected_peers if peer not in self.talk.peers]:
|
|
# We are missing peer(s)!
|
|
logger.warning(
|
|
f"Analytics could not take snapshot because peers are missing: {', '.join(missing)}"
|
|
)
|
|
return False
|
|
|
|
# Everyone is here, ask for shard snapshots
|
|
results = await self.talk_shard_snapshot().broadcast()
|
|
|
|
# Make sure everyone sent results and there were no exceptions (e.g. concurrency)
|
|
if not all(result is not None and not isinstance(result, Exception) for result in results.values()):
|
|
# This should essentially never happen
|
|
# Either some of the shards could not make a snapshot (e.g. Discord client issues)
|
|
# or they disconnected in the process.
|
|
logger.warning(
|
|
f"Analytics could not take snapshot because some peers failed! Partial snapshot: {results}"
|
|
)
|
|
return False
|
|
|
|
# Now we have a dictionary of shard snapshots, aggregate, pull in remaining data, and store.
|
|
# TODO Possibly move this out into snapshots.py?
|
|
aggregate = {field: 0 for field in ShardSnapshot._fields}
|
|
for result in results.values():
|
|
for field, num in result._asdict().items():
|
|
aggregate[field] += num
|
|
|
|
row = await self.data.Snapshots.create(
|
|
appname=appname,
|
|
guild_count=aggregate['guild_count'],
|
|
member_count=aggregate['member_count'],
|
|
user_count=aggregate['user_count'],
|
|
in_voice=aggregate['voice_count'],
|
|
)
|
|
logger.info(f"Created snapshot: {row.data!r}")
|
|
return True
|
|
|
|
@log_wrap(action='SnapLoop')
|
|
async def snapshot_loop(self):
|
|
while True:
|
|
try:
|
|
result = await self.take_snapshot()
|
|
if result:
|
|
await asyncio.sleep(self.snap_period)
|
|
else:
|
|
logger.info("Snapshot failed, retrying after %d seconds", self.snap_retry_period)
|
|
await asyncio.sleep(self.snap_retry_period)
|
|
except asyncio.CancelledError:
|
|
logger.info("Snapshot loop cancelled, closing.")
|
|
return
|
|
except Exception:
|
|
logger.exception(
|
|
"Unhandled exception during snapshot loop. Ignoring and continuing cautiously."
|
|
)
|
|
await asyncio.sleep(self.snap_retry_period)
|
|
|
|
async def run(self):
|
|
log_action_stack.set(['Analytics'])
|
|
log_app.set(conf.analytics['appname'])
|
|
|
|
async with await self.db.connect():
|
|
await self.talk.connect()
|
|
await self.attach_event_handlers()
|
|
self._snap_task = asyncio.create_task(self.snapshot_loop())
|
|
await asyncio.gather(*(handler._consumer_task for handler in self.event_handlers))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
server = AnalyticsServer()
|
|
asyncio.run(server.run())
|