Skip to content

Commit

Permalink
fix(node): properly resume
Browse files Browse the repository at this point in the history
  • Loading branch information
@ooliver1
ooliver1 committed Feb 4, 2023
1 parent e563889 commit bb20ad6
Show file tree
Hide file tree
Showing 5 changed files with 150 additions and 117 deletions.
45 changes: 23 additions & 22 deletions mafic/node.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -495,13 +495,6 @@ async def connect(self) -> None:
await self._connect_to_websocket(headers=headers, session=session)
_log.info("Connected to lavalink.", extra={"label": self._label})

_log.debug(
"Creating task to send configuration to resume with key %s",
self._resume_key,
extra={"label": self._label},
)
create_task(self.configure_resuming())

_log.info(
"Creating task for websocket listener...", extra={"label": self._label}
)
Expand Down Expand Up @@ -606,16 +599,24 @@ async def _handle_msg(self, data: IncomingMessage) -> None:
resumed = data["resumed"]
session_id = data["sessionId"]

_log.debug(
"Received session ID %s", session_id, extra={"label": self._label}
)
self._session_id = session_id

if resumed:
_log.info(
"Successfully resumed connection with lavalink.",
extra={"label": self._label},
)
else:
_log.debug(
"Sending configuration to resume with key %s",
self._resume_key,
extra={"label": self._label},
)
await self.configure_resuming()

_log.debug(
"Received session ID %s", session_id, extra={"label": self._label}
)
self._session_id = session_id
self._ready.set()
else:
# Of course pyright considers this to be `Never`, so this is to keep types.
Expand Down Expand Up @@ -672,7 +673,7 @@ def voice_update(

return self.__request(
"PATCH",
f"/sessions/{self._session_id}/players/{guild_id}",
f"sessions/{self._session_id}/players/{guild_id}",
{
"voice": {
"sessionId": session_id,
Expand All @@ -693,7 +694,7 @@ def configure_resuming(self) -> Coro[None]:

return self.__request(
"PATCH",
f"/sessions/{self._session_id}",
f"sessions/{self._session_id}",
{
"resumingKey": self._resume_key,
"timeout": 60,
Expand All @@ -712,7 +713,7 @@ def destroy(self, guild_id: int) -> Coro[None]:
_log.debug("Sending request to destroy player", extra={"label": self._label})

return self.__request(
"DELETE", f"/sessions/{self._session_id}/players/{guild_id}"
"DELETE", f"sessions/{self._session_id}/players/{guild_id}"
)

def update(
Expand Down Expand Up @@ -783,7 +784,7 @@ def update(

return self.__request(
"PATCH",
f"/sessions/{self._session_id}/players/{guild_id}",
f"sessions/{self._session_id}/players/{guild_id}",
data,
query,
)
Expand Down Expand Up @@ -878,7 +879,7 @@ async def fetch_tracks(
query = f"{search_type}:{query}"

data: TrackLoadingResult = await self.__request(
"GET", "/loadtracks", params={"identifier": query}
"GET", "loadtracks", params={"identifier": query}
)

if data["loadType"] == "NO_MATCHES":
Expand Down Expand Up @@ -913,7 +914,7 @@ async def decode_track(self, track: str) -> Track:
"""

info: TrackInfo = await self.__request(
"GET", "/decodetrack", params={"encodedTrack": track}
"GET", "decodetrack", params={"encodedTrack": track}
)

return Track.from_data(track=track, info=info)
Expand All @@ -937,7 +938,7 @@ async def decode_tracks(self, tracks: list[str]) -> list[Track]:
"""

track_data: list[TrackWithInfo] = await self.__request(
"POST", "/decodetracks", json=tracks
"POST", "decodetracks", json=tracks
)

return [Track.from_data_with_info(track) for track in track_data]
Expand All @@ -951,7 +952,7 @@ async def fetch_plugins(self) -> list[Plugin]:
The plugins from the node.
"""

plugins: list[PluginData] = await self.__request("GET", "/plugins")
plugins: list[PluginData] = await self.__request("GET", "plugins")

return [Plugin(plugin) for plugin in plugins]

Expand All @@ -965,7 +966,7 @@ async def fetch_route_planner_status(self) -> RoutePlannerStatus | None:
"""

data: RoutePlannerStatusPayload = await self.__request(
"GET", "/routeplanner/status"
"GET", "routeplanner/status"
)

if data["class"] == "RotatingIpRoutePlanner":
Expand Down Expand Up @@ -997,10 +998,10 @@ async def unmark_failed_address(self, address: str) -> None:
"""

await self.__request(
"POST", "/routeplanner/free/address", json={"address": address}
"POST", "routeplanner/free/address", json={"address": address}
)

async def unmark_all_addresses(self) -> None:
"""Unmark all failed addresses so they can be used again."""

await self.__request("POST", "/routeplanner/free/all")
await self.__request("POST", "routeplanner/free/all")
192 changes: 102 additions & 90 deletions test_bot/bot/__main__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,104 +26,116 @@ async def on_ready(self):
if self.ready_ran:
return

# Account for docker still starting up.
await sleep(5)
# Excessively test pool balancing.
await self.pool.create_node(
host="127.0.0.1",
port=6962,
label="US-noshard",
password="haha",
regions=[Group.WEST, Region.OCEANIA, Region.EAST_ASIA],
)
await self.pool.create_node(
host="127.0.0.1",
port=6963,
label="EU-noshard",
password="haha",
regions=[
Group.CENTRAL,
Region.WEST_ASIA,
Region.NORTH_ASIA,
Region.SOUTH_ASIA,
],
)
await self.pool.create_node(
host="127.0.0.1",
port=6964,
label="US-shard0",
password="haha",
regions=[Group.WEST, Region.OCEANIA, Region.EAST_ASIA],
shard_ids=[0],
)
await self.pool.create_node(
host="127.0.0.1",
port=6965,
label="US-shard1",
password="haha",
regions=[Group.WEST, Region.OCEANIA, Region.EAST_ASIA],
shard_ids=[1],
)
await self.pool.create_node(
host="127.0.0.1",
port=6966,
label="EU-shard0-1",
password="haha",
regions=[
Group.CENTRAL,
Region.WEST_ASIA,
Region.NORTH_ASIA,
Region.SOUTH_ASIA,
],
shard_ids=[0],
)
await self.pool.create_node(
host="127.0.0.1",
port=6967,
label="EU-shard0-2",
password="haha",
regions=[
Group.CENTRAL,
Region.WEST_ASIA,
Region.NORTH_ASIA,
Region.SOUTH_ASIA,
],
shard_ids=[0],
)
await self.pool.create_node(
host="127.0.0.1",
port=6968,
label="EU-shard1-1",
password="haha",
regions=[
Group.CENTRAL,
Region.WEST_ASIA,
Region.NORTH_ASIA,
Region.SOUTH_ASIA,
],
shard_ids=[1],
)
await self.pool.create_node(
host="127.0.0.1",
port=6969,
label="EU-shard1-2",
password="haha",
regions=[
Group.CENTRAL,
Region.WEST_ASIA,
Region.NORTH_ASIA,
Region.SOUTH_ASIA,
],
shard_ids=[1],
)
if getenv("TEST_BALANCING"):
# Account for docker still starting up.
await sleep(5)
await self.pool.create_node(
host="127.0.0.1",
port=6962,
label="US-noshard",
password="haha",
regions=[Group.WEST, Region.OCEANIA, Region.EAST_ASIA],
)
await self.pool.create_node(
host="127.0.0.1",
port=6963,
label="EU-noshard",
password="haha",
regions=[
Group.CENTRAL,
Region.WEST_ASIA,
Region.NORTH_ASIA,
Region.SOUTH_ASIA,
],
)
await self.pool.create_node(
host="127.0.0.1",
port=6964,
label="US-shard0",
password="haha",
regions=[Group.WEST, Region.OCEANIA, Region.EAST_ASIA],
shard_ids=[0],
)
await self.pool.create_node(
host="127.0.0.1",
port=6965,
label="US-shard1",
password="haha",
regions=[Group.WEST, Region.OCEANIA, Region.EAST_ASIA],
shard_ids=[1],
)
await self.pool.create_node(
host="127.0.0.1",
port=6966,
label="EU-shard0-1",
password="haha",
regions=[
Group.CENTRAL,
Region.WEST_ASIA,
Region.NORTH_ASIA,
Region.SOUTH_ASIA,
],
shard_ids=[0],
)
await self.pool.create_node(
host="127.0.0.1",
port=6967,
label="EU-shard0-2",
password="haha",
regions=[
Group.CENTRAL,
Region.WEST_ASIA,
Region.NORTH_ASIA,
Region.SOUTH_ASIA,
],
shard_ids=[0],
)
await self.pool.create_node(
host="127.0.0.1",
port=6968,
label="EU-shard1-1",
password="haha",
regions=[
Group.CENTRAL,
Region.WEST_ASIA,
Region.NORTH_ASIA,
Region.SOUTH_ASIA,
],
shard_ids=[1],
)
await self.pool.create_node(
host="127.0.0.1",
port=6969,
label="EU-shard1-2",
password="haha",
regions=[
Group.CENTRAL,
Region.WEST_ASIA,
Region.NORTH_ASIA,
Region.SOUTH_ASIA,
],
shard_ids=[1],
)
else:
await self.pool.create_node(
host="127.0.0.1",
port=6969,
label="MAIN",
password="haha",
)

self.ready_ran = True


intents = Intents.none()
intents.guilds = True
intents.voice_states = True
bot = TestBot(intents=intents, shard_ids=[0, 1], shard_count=2)

if getenv("TEST_BALANCING"):
bot = TestBot(intents=intents, shard_ids=[0, 1], shard_count=2)
else:
bot = TestBot(intents=intents)


class MyPlayer(Player[TestBot]):
Expand Down
22 changes: 22 additions & 0 deletions test_bot/docker-compose.multi-nodeyml
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
version: "3"

services:
bot:
env_file: .env
depends_on:
- lavalink
build: .
volumes:
- ./logs:/bot/logs
- ../mafic:/bot/mafic
environment:
TEST_BALANCING: 1
network_mode: host
lavalink:
build: lavalink
volumes:
- ./logs/lava:/opt/Lavalink/logs
deploy:
replicas: 8
ports:
- "6962-6969:6969"
4 changes: 1 addition & 3 deletions test_bot/docker-compose.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,5 @@ services:
build: lavalink
volumes:
- ./logs/lava:/opt/Lavalink/logs
deploy:
replicas: 8
ports:
- "6962-6969:6969"
- "6969:6969"
4 changes: 2 additions & 2 deletions test_bot/lavalink/entrypoint.sh
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@

#!/bin/sh

id -u lavalink &>/dev/null || groupadd -g 322 lavalink
id -u lavalink &>/dev/null || useradd -u 322 -g 322 lavalink
_=$(id -u lavalink &>/dev/null || groupadd -g 322 lavalink)
_=$(id -u lavalink &>/dev/null || useradd -u 322 -g 322 lavalink)
chown -R 322:322 ./logs
exec runuser -u lavalink -- "$@"

0 comments on commit bb20ad6

Please sign in to comment.