Browse Source
This allows people to run their >2,500 guild bot in a single process without the headaches of IPC/RPC or much difficulty.pull/447/head
Rapptz
8 years ago
9 changed files with 341 additions and 31 deletions
@ -0,0 +1,174 @@ |
|||
# -*- coding: utf-8 -*- |
|||
|
|||
""" |
|||
The MIT License (MIT) |
|||
|
|||
Copyright (c) 2015-2016 Rapptz |
|||
|
|||
Permission is hereby granted, free of charge, to any person obtaining a |
|||
copy of this software and associated documentation files (the "Software"), |
|||
to deal in the Software without restriction, including without limitation |
|||
the rights to use, copy, modify, merge, publish, distribute, sublicense, |
|||
and/or sell copies of the Software, and to permit persons to whom the |
|||
Software is furnished to do so, subject to the following conditions: |
|||
|
|||
The above copyright notice and this permission notice shall be included in |
|||
all copies or substantial portions of the Software. |
|||
|
|||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS |
|||
OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, |
|||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE |
|||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER |
|||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING |
|||
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER |
|||
DEALINGS IN THE SOFTWARE. |
|||
""" |
|||
|
|||
from .state import AutoShardedConnectionState |
|||
from .client import Client |
|||
from .gateway import * |
|||
from .errors import ConnectionClosed |
|||
from . import compat |
|||
|
|||
import asyncio |
|||
import logging |
|||
|
|||
log = logging.getLogger(__name__) |
|||
|
|||
class Shard: |
|||
def __init__(self, ws, client): |
|||
self.ws = ws |
|||
self._client = client |
|||
self.loop = self._client.loop |
|||
self._current = asyncio.Future(loop=self.loop) |
|||
self._current.set_result(None) # we just need an already done future |
|||
|
|||
@property |
|||
def id(self): |
|||
return self.ws.shard_id |
|||
|
|||
@asyncio.coroutine |
|||
def poll(self): |
|||
try: |
|||
yield from self.ws.poll_event() |
|||
except (ReconnectWebSocket, ResumeWebSocket) as e: |
|||
resume = type(e) is ResumeWebSocket |
|||
log.info('Got ' + type(e).__name__) |
|||
self.ws = yield from DiscordWebSocket.from_client(self._client, resume=resume, |
|||
shard_id=self.id, |
|||
session=self.ws.session_id, |
|||
sequence=self.ws.sequence) |
|||
except ConnectionClosed as e: |
|||
yield from self._client.close() |
|||
if e.code != 1000: |
|||
raise |
|||
|
|||
def get_future(self): |
|||
if self._current.done(): |
|||
self._current = compat.create_task(self.poll(), loop=self.loop) |
|||
|
|||
return self._current |
|||
|
|||
class AutoShardedClient(Client): |
|||
"""A client similar to :class:`Client` except it handles the complications |
|||
of sharding for the user into a more manageable and transparent single |
|||
process bot. |
|||
|
|||
When using this client, you will be able to use it as-if it was a regular |
|||
:class:`Client` with a single shard when implementation wise internally it |
|||
is split up into multiple shards. This allows you to not have to deal with |
|||
IPC or other complicated infrastructure. |
|||
|
|||
It is recommended to use this client only if you have surpassed at least |
|||
1000 guilds. |
|||
|
|||
If no :attr:`shard_count` is provided, then the library will use the |
|||
Bot Gateway endpoint call to figure out how many shards to use. |
|||
""" |
|||
def __init__(self, *args, loop=None, **kwargs): |
|||
kwargs.pop('shard_id', None) |
|||
super().__init__(*args, loop=loop, **kwargs) |
|||
|
|||
self.connection = AutoShardedConnectionState(dispatch=self.dispatch, chunker=self.request_offline_members, |
|||
syncer=self._syncer, http=self.http, loop=self.loop, **kwargs) |
|||
|
|||
# instead of a single websocket, we have multiple |
|||
# the index is the shard_id |
|||
self.shards = [] |
|||
|
|||
@asyncio.coroutine |
|||
def request_offline_members(self, guild, *, shard_id=None): |
|||
"""|coro| |
|||
|
|||
Requests previously offline members from the guild to be filled up |
|||
into the :attr:`Guild.members` cache. This function is usually not |
|||
called. |
|||
|
|||
When the client logs on and connects to the websocket, Discord does |
|||
not provide the library with offline members if the number of members |
|||
in the guild is larger than 250. You can check if a guild is large |
|||
if :attr:`Guild.large` is ``True``. |
|||
|
|||
Parameters |
|||
----------- |
|||
guild: :class:`Guild` or list |
|||
The guild to request offline members for. If this parameter is a |
|||
list then it is interpreted as a list of guilds to request offline |
|||
members for. |
|||
""" |
|||
|
|||
try: |
|||
guild_id = guild.id |
|||
shard_id = shard_id or guild.shard_id |
|||
except AttributeError: |
|||
guild_id = [s.id for s in guild] |
|||
|
|||
payload = { |
|||
'op': 8, |
|||
'd': { |
|||
'guild_id': guild_id, |
|||
'query': '', |
|||
'limit': 0 |
|||
} |
|||
} |
|||
|
|||
ws = self.shards[shard_id].ws |
|||
yield from ws.send_as_json(payload) |
|||
|
|||
@asyncio.coroutine |
|||
def connect(self): |
|||
"""|coro| |
|||
|
|||
Creates a websocket connection and lets the websocket listen |
|||
to messages from discord. |
|||
|
|||
Raises |
|||
------- |
|||
GatewayNotFound |
|||
If the gateway to connect to discord is not found. Usually if this |
|||
is thrown then there is a discord API outage. |
|||
ConnectionClosed |
|||
The websocket connection has been terminated. |
|||
""" |
|||
ret = yield from DiscordWebSocket.from_sharded_client(self) |
|||
self.shards = [Shard(ws, self) for ws in ret] |
|||
|
|||
while not self.is_closed: |
|||
pollers = [shard.get_future() for shard in self.shards] |
|||
yield from asyncio.wait(pollers, loop=self.loop, return_when=asyncio.FIRST_COMPLETED) |
|||
|
|||
@asyncio.coroutine |
|||
def close(self): |
|||
"""|coro| |
|||
|
|||
Closes the connection to discord. |
|||
""" |
|||
if self.is_closed: |
|||
return |
|||
|
|||
for shard in self.shards: |
|||
yield from shard.ws.close() |
|||
|
|||
yield from self.http.close() |
|||
self._closed.set() |
|||
self._is_ready.clear() |
|||
Loading…
Reference in new issue