예제 #1
0
파일: dbot.py 프로젝트: leenephi/avrae
    def __init__(self, prefix, description=None, testing=False, **options):
        super(Avrae, self).__init__(prefix, help_command=help_command, description=description, **options)
        self.testing = testing
        self.state = "init"
        self.credentials = Credentials()
        if config.TESTING:
            self.mclient = motor.motor_asyncio.AsyncIOMotorClient(self.credentials.test_mongo_url)
        else:
            self.mclient = motor.motor_asyncio.AsyncIOMotorClient(config.MONGO_URL)
        self.mdb = self.mclient[config.MONGODB_DB_NAME]
        self.rdb = self.loop.run_until_complete(self.setup_rdb())
        self.prefixes = dict()
        self.muted = set()
        self.cluster_id = 0

        # sentry
        if config.SENTRY_DSN is not None:
            release = None
            if config.GIT_COMMIT_SHA:
                release = f"avrae-bot@{config.GIT_COMMIT_SHA}"
            sentry_sdk.init(dsn=config.SENTRY_DSN, environment=config.ENVIRONMENT.title(), release=release)

        # ddb entitlements
        if config.TESTING and config.DDB_AUTH_SERVICE_URL is None:
            self.ddb = BeyondClientBase()
        else:
            self.ddb = BeyondClient(self.loop)

        # launchdarkly
        self.ldclient = AsyncLaunchDarklyClient(self.loop, sdk_key=config.LAUNCHDARKLY_SDK_KEY)
예제 #2
0
    def __init__(self, prefix, description=None, testing=False, **options):
        test_guilds = None if not testing else config.COMMAND_TEST_GUILD_IDS

        super().__init__(prefix,
                         help_command=help_command,
                         description=description,
                         test_guilds=test_guilds,
                         sync_commands=False,
                         sync_commands_debug=testing,
                         **options)
        self.testing = testing
        self.state = "init"

        # dbs
        self.mclient = motor.motor_asyncio.AsyncIOMotorClient(config.MONGO_URL)
        self.mdb = self.mclient[config.MONGODB_DB_NAME]
        self.rdb = self.loop.run_until_complete(self.setup_rdb())

        # misc caches
        self.prefixes = dict()
        self.muted = set()
        self.cluster_id = 0

        # launch concurrency
        self.launch_max_concurrency = 1

        # sentry
        if config.SENTRY_DSN is not None:
            release = None
            if config.GIT_COMMIT_SHA:
                release = f"avrae-bot@{config.GIT_COMMIT_SHA}"
            sentry_sdk.init(dsn=config.SENTRY_DSN,
                            environment=config.ENVIRONMENT.title(),
                            release=release)

        # ddb entitlements
        if config.TESTING and config.DDB_AUTH_SERVICE_URL is None:
            self.ddb = BeyondClientBase()
        else:
            self.ddb = BeyondClient(self.loop)

        # launchdarkly
        self.ldclient = AsyncLaunchDarklyClient(
            self.loop, sdk_key=config.LAUNCHDARKLY_SDK_KEY)

        # ddb game log
        self.glclient = GameLogClient(self)
        self.glclient.init()
예제 #3
0
파일: dbot.py 프로젝트: leenephi/avrae
class Avrae(commands.AutoShardedBot):
    def __init__(self, prefix, description=None, testing=False, **options):
        super(Avrae, self).__init__(prefix, help_command=help_command, description=description, **options)
        self.testing = testing
        self.state = "init"
        self.credentials = Credentials()
        if config.TESTING:
            self.mclient = motor.motor_asyncio.AsyncIOMotorClient(self.credentials.test_mongo_url)
        else:
            self.mclient = motor.motor_asyncio.AsyncIOMotorClient(config.MONGO_URL)
        self.mdb = self.mclient[config.MONGODB_DB_NAME]
        self.rdb = self.loop.run_until_complete(self.setup_rdb())
        self.prefixes = dict()
        self.muted = set()
        self.cluster_id = 0

        # sentry
        if config.SENTRY_DSN is not None:
            release = None
            if config.GIT_COMMIT_SHA:
                release = f"avrae-bot@{config.GIT_COMMIT_SHA}"
            sentry_sdk.init(dsn=config.SENTRY_DSN, environment=config.ENVIRONMENT.title(), release=release)

        # ddb entitlements
        if config.TESTING and config.DDB_AUTH_SERVICE_URL is None:
            self.ddb = BeyondClientBase()
        else:
            self.ddb = BeyondClient(self.loop)

        # launchdarkly
        self.ldclient = AsyncLaunchDarklyClient(self.loop, sdk_key=config.LAUNCHDARKLY_SDK_KEY)

    async def setup_rdb(self):
        if config.TESTING:
            redis_url = self.credentials.test_redis_url
        else:
            redis_url = config.REDIS_URL
        return RedisIO(await aioredis.create_redis_pool(redis_url, db=config.REDIS_DB_NUM))

    async def get_server_prefix(self, msg):
        return (await get_prefix(self, msg))[-1]

    async def launch_shards(self):
        # set up my shard_ids
        async with clustering.coordination_lock(self.rdb):
            await clustering.coordinate_shards(self)
            if self.shard_ids is not None:
                log.info(f"Launching {len(self.shard_ids)} shards! ({set(self.shard_ids)})")
            await super(Avrae, self).launch_shards()
            log.info(f"Launched {len(self.shards)} shards!")

        if self.is_cluster_0:
            await self.rdb.incr('build_num')

    async def close(self):
        await super().close()
        await self.ddb.close()
        self.ldclient.close()

    @property
    def is_cluster_0(self):
        if self.cluster_id is None:  # we're not running in clustered mode anyway
            return True
        return self.cluster_id == 0

    @staticmethod
    def log_exception(exception=None, context: commands.Context = None):
        if config.SENTRY_DSN is None:
            return

        with sentry_sdk.push_scope() as scope:
            if context:
                # noinspection PyDunderSlots,PyUnresolvedReferences
                # for some reason pycharm doesn't pick up the attribute setter here
                scope.user = {"id": context.author.id, "username": str(context.author)}
                scope.set_tag("message.content", context.message.content)
                scope.set_tag("is_private_message", context.guild is None)
                scope.set_tag("channel.id", context.channel.id)
                scope.set_tag("channel.name", str(context.channel))
                if context.guild is not None:
                    scope.set_tag("guild.id", context.guild.id)
                    scope.set_tag("guild.name", str(context.guild))
            sentry_sdk.capture_exception(exception)
예제 #4
0
class Avrae(commands.AutoShardedBot):
    def __init__(self, prefix, description=None, testing=False, **options):
        test_guilds = None if not testing else config.COMMAND_TEST_GUILD_IDS

        super().__init__(prefix,
                         help_command=help_command,
                         description=description,
                         test_guilds=test_guilds,
                         sync_commands=False,
                         sync_commands_debug=testing,
                         **options)
        self.testing = testing
        self.state = "init"

        # dbs
        self.mclient = motor.motor_asyncio.AsyncIOMotorClient(config.MONGO_URL)
        self.mdb = self.mclient[config.MONGODB_DB_NAME]
        self.rdb = self.loop.run_until_complete(self.setup_rdb())

        # misc caches
        self.prefixes = dict()
        self.muted = set()
        self.cluster_id = 0

        # launch concurrency
        self.launch_max_concurrency = 1

        # sentry
        if config.SENTRY_DSN is not None:
            release = None
            if config.GIT_COMMIT_SHA:
                release = f"avrae-bot@{config.GIT_COMMIT_SHA}"
            sentry_sdk.init(dsn=config.SENTRY_DSN,
                            environment=config.ENVIRONMENT.title(),
                            release=release)

        # ddb entitlements
        if config.TESTING and config.DDB_AUTH_SERVICE_URL is None:
            self.ddb = BeyondClientBase()
        else:
            self.ddb = BeyondClient(self.loop)

        # launchdarkly
        self.ldclient = AsyncLaunchDarklyClient(
            self.loop, sdk_key=config.LAUNCHDARKLY_SDK_KEY)

        # ddb game log
        self.glclient = GameLogClient(self)
        self.glclient.init()

    async def setup_rdb(self):
        return RedisIO(await
                       aioredis.create_redis_pool(config.REDIS_URL,
                                                  db=config.REDIS_DB_NUM))

    async def get_guild_prefix(self, guild: discord.Guild) -> str:
        guild_id = str(guild.id)
        if guild_id in self.prefixes:
            return self.prefixes.get(guild_id, config.DEFAULT_PREFIX)
        # load from db and cache
        gp_obj = await self.mdb.prefixes.find_one({"guild_id": guild_id})
        if gp_obj is None:
            gp = config.DEFAULT_PREFIX
        else:
            gp = gp_obj.get("prefix", config.DEFAULT_PREFIX)
        self.prefixes[guild_id] = gp
        return gp

    @property
    def is_cluster_0(self):
        if self.cluster_id is None:  # we're not running in clustered mode anyway
            return True
        return self.cluster_id == 0

    @staticmethod
    def log_exception(exception=None, ctx: context.AvraeContext = None):
        if config.SENTRY_DSN is None:
            return

        with sentry_sdk.push_scope() as scope:
            if ctx:
                # noinspection PyDunderSlots,PyUnresolvedReferences
                # for some reason pycharm doesn't pick up the attribute setter here
                scope.user = {"id": ctx.author.id, "username": str(ctx.author)}
                scope.set_tag("message.content", ctx.message.content)
                scope.set_tag("is_private_message", ctx.guild is None)
                scope.set_tag("channel.id", ctx.channel.id)
                scope.set_tag("channel.name", str(ctx.channel))
                if ctx.guild is not None:
                    scope.set_tag("guild.id", ctx.guild.id)
                    scope.set_tag("guild.name", str(ctx.guild))
            sentry_sdk.capture_exception(exception)

    async def launch_shards(self):
        # set up my shard_ids
        async with clustering.coordination_lock(self.rdb):
            await clustering.coordinate_shards(self)
            log.info(f"I am cluster {self.cluster_id}.")
            if self.shard_ids is not None:
                log.info(
                    f"Launching {len(self.shard_ids)} shards! ({self.shard_ids})"
                )

        # if we are cluster 0, we are responsible for handling application command sync
        if self.is_cluster_0:
            self._sync_commands = True

        # release lock and launch
        await super().launch_shards()
        log.info(f"Launched {len(self.shards)} shards!")

        if self.is_cluster_0:
            await self.rdb.incr('build_num')

    async def before_identify_hook(self, shard_id, *, initial=False):
        bucket_id = shard_id % self.launch_max_concurrency
        # dummy call to initialize monitoring - see note on returning 0.0 at
        # https://psutil.readthedocs.io/en/latest/index.html#psutil.cpu_percent
        psutil.cpu_percent()

        async def pre_lock_check(first=False):
            """
            Before attempting a lock, CPU utilization should be <75% to prevent a huge spike in CPU when connecting
            up to 16 shards at once. This also allows multiple clusters to startup concurrently.
            """
            if not first:
                await asyncio.sleep(0.2)
            wait_start = time.monotonic()
            while psutil.cpu_percent() > 75:
                t = random.uniform(5, 15)
                log.info(
                    f"[C{self.cluster_id}] CPU usage is high, waiting {t:.2f}s!"
                )
                await asyncio.sleep(t)
                if time.monotonic(
                ) - wait_start > 300:  # liveness: wait no more than 5 minutes
                    break

        # wait until the bucket is available and try to acquire the lock
        await clustering.wait_bucket_available(shard_id,
                                               bucket_id,
                                               self.rdb,
                                               pre_lock_hook=pre_lock_check)

    async def get_context(self, *args, **kwargs) -> context.AvraeContext:
        return await super().get_context(*args,
                                         cls=context.AvraeContext,
                                         **kwargs)

    async def close(self):
        # note: when closing the bot 2 errors are emitted:
        #
        # ERROR:asyncio: An open stream object is being garbage collected; call "stream.close()" explicitly.
        # ERROR:asyncio: An open stream object is being garbage collected; call "stream.close()" explicitly.
        #
        # These are caused by aioredis streams being GC'ed when discord.py cancels the tasks that create them
        # (because of course d.py decides it wants to cancel *all* tasks on its loop...)
        await super().close()
        await self.ddb.close()
        await self.rdb.close()
        await self.glclient.close()
        self.mclient.close()
        self.ldclient.close()