From b468883e2eb5fe1f454220b2f6f85463d83e44b1 Mon Sep 17 00:00:00 2001 From: spiral Date: Fri, 24 Sep 2021 13:04:00 -0400 Subject: [PATCH 1/4] gateway: add max_concurrency support --- litecord/blueprints/gateway.py | 10 +++++++++- litecord/gateway/websocket.py | 10 +++++++++- manage/cmd/migration/scripts/14_max_concurrency.sql | 3 +++ 3 files changed, 21 insertions(+), 2 deletions(-) create mode 100644 manage/cmd/migration/scripts/14_max_concurrency.sql diff --git a/litecord/blueprints/gateway.py b/litecord/blueprints/gateway.py index bb90d42..95ce132 100644 --- a/litecord/blueprints/gateway.py +++ b/litecord/blueprints/gateway.py @@ -51,6 +51,14 @@ async def api_gateway_bot(): user_id, ) + max_concurrency = await app.db.fetchval( + """select max_concurrency + from users + where id = $1 + """, + user_id, + ) + shards = max(int(guild_count / 1000), 1) # get _ws.session ratelimit @@ -78,7 +86,7 @@ async def api_gateway_bot(): "total": bucket.requests, "remaining": bucket._tokens, "reset_after": int(reset_after_ts * 1000), - "max_concurrency": 1, + "max_concurrency": max_concurrency, }, } ) diff --git a/litecord/gateway/websocket.py b/litecord/gateway/websocket.py index c306591..3190024 100644 --- a/litecord/gateway/websocket.py +++ b/litecord/gateway/websocket.py @@ -811,7 +811,15 @@ class GatewayWebsocket: except (Unauthorized, Forbidden): raise WebsocketClose(4004, "Authentication failed") - await self._connect_ratelimit(user_id) + max_concurrency = await self.app.db.fetchval( + """select max_concurrency + from users + where id = $1 + """, + user_id, + ) + + await self._connect_ratelimit(f"{str(user_id)}%{shard[0]%max_concurrency}") bot = await self.app.db.fetchval( """ diff --git a/manage/cmd/migration/scripts/14_max_concurrency.sql b/manage/cmd/migration/scripts/14_max_concurrency.sql new file mode 100644 index 0000000..3366bc9 --- /dev/null +++ b/manage/cmd/migration/scripts/14_max_concurrency.sql @@ -0,0 +1,3 @@ +alter table users + add column max_concurrency int not null default 1 + check(bot = true or max_concurrency = 1); \ No newline at end of file From bba48f7d0f97a9dbaf3cefcae6743e76d3536eaa Mon Sep 17 00:00:00 2001 From: spiral Date: Fri, 24 Sep 2021 13:40:03 -0400 Subject: [PATCH 2/4] fix formatting --- litecord/blueprints/gateway.py | 2 +- litecord/gateway/websocket.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litecord/blueprints/gateway.py b/litecord/blueprints/gateway.py index 95ce132..5a4601b 100644 --- a/litecord/blueprints/gateway.py +++ b/litecord/blueprints/gateway.py @@ -56,7 +56,7 @@ async def api_gateway_bot(): from users where id = $1 """, - user_id, + user_id, ) shards = max(int(guild_count / 1000), 1) diff --git a/litecord/gateway/websocket.py b/litecord/gateway/websocket.py index 3190024..c535c1a 100644 --- a/litecord/gateway/websocket.py +++ b/litecord/gateway/websocket.py @@ -816,7 +816,7 @@ class GatewayWebsocket: from users where id = $1 """, - user_id, + user_id, ) await self._connect_ratelimit(f"{str(user_id)}%{shard[0]%max_concurrency}") From 2ad6b29175c27b8d05d713df97837324d9d22c57 Mon Sep 17 00:00:00 2001 From: spiral Date: Wed, 10 Nov 2021 00:29:07 -0500 Subject: [PATCH 3/4] add manage command set_max_concurrency --- manage/cmd/users.py | 47 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/manage/cmd/users.py b/manage/cmd/users.py index efd3cf4..48717a0 100644 --- a/manage/cmd/users.py +++ b/manage/cmd/users.py @@ -94,6 +94,42 @@ async def adduser(ctx, args): print(f'\tdiscrim: {user["discriminator"]}') +async def set_max_concurrency(ctx, args): + """Update the `max_concurrency` for a bot. + This can only be set for bot accounts! + """ + + if int(args.max_concurrency) < 1: + return print("max_concurrency must be >0") + + bot = await ctx.db.fetchval( + """ + select bot + from users + where id = $1 + """, + int(args.user_id), + ) + + if bot == None: + return print("user not found") + + if bot == False: + return print("user must be a bot") + + await ctx.db.execute( + """ + update users + set max_concurrency = $1 + where id = $2 + """, + int(args.max_concurrency), + int(args.user_id), + ) + + print(f"OK: set max_concurrency={args.max_concurrency} for {args.user_id}") + + async def set_flag(ctx, args): """Setting a 'staff' flag gives the user access to the Admin API. Beware of that. @@ -198,6 +234,17 @@ def setup(subparser): setup_test_parser.set_defaults(func=adduser) + set_max_concurrency_parser = subparser.add_parser( + "set_max_concurrency", + help="set `max_concurrency` for a user", + description=set_max_concurrency.__doc__, + ) + set_max_concurrency_parser.add_argument("user_id") + set_max_concurrency_parser.add_argument( + "max_concurrency", help="the `max_concurrency` value to set" + ) + set_max_concurrency_parser.set_defaults(func=set_max_concurrency) + setflag_parser = subparser.add_parser( "setflag", help="set a flag for a user", description=set_flag.__doc__ ) From e3f894330dafdbe7ab1d1e958cf0d0f3da1d1a40 Mon Sep 17 00:00:00 2001 From: spiral Date: Wed, 10 Nov 2021 00:44:54 -0500 Subject: [PATCH 4/4] fix formatting --- manage/cmd/users.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/manage/cmd/users.py b/manage/cmd/users.py index 48717a0..e54ee87 100644 --- a/manage/cmd/users.py +++ b/manage/cmd/users.py @@ -98,7 +98,7 @@ async def set_max_concurrency(ctx, args): """Update the `max_concurrency` for a bot. This can only be set for bot accounts! """ - + if int(args.max_concurrency) < 1: return print("max_concurrency must be >0")