Skip to content

Commit

Permalink
Add Shard Id helpers
Browse files Browse the repository at this point in the history
Add helpers to retrieve the shard Id for guilds, and count how many
guilds are handled by a Shard.

Helpers to retrieve the shard Id of a guild have been added as:

- `Guild::shard_id`
- `GuildId::shard_id`

These are in two forms: one working with the cache feature, and one
without. The function that works with the cache will automatically
retrieve the total number of shards from the Cache, while the uncached
version requires passing in the total number of shards used.

With the cache enabled, this might look like:

```rust
guild.shard_id();
// which calls:
guild_id.shard_id();
```

Without the cache enabled, this looks like:

```rust
let shard_count = 7;

guild.shard_id(shard_count);
// which calls:
guild_id.shard_id(shard_count);
```

These two variants on `Guild` and `GuildId` are helper sugar methods
over the new function `utils::shard_id`, which accepts a `guild_id` and
a `shard_count`:

```rust
use serenity::utils;

assert_eq!(utils::shard_id(81384788765712384, 17), 7);
```

You would use `utils::shard_id` when you have the total number of shards
due to `{Guild,GuildId}::shard_id` unlocking the cache to retrieve the
total number of shards. This avoids some amount of work
  • Loading branch information
Zeyla Hellyer committed Apr 13, 2017
1 parent c4d0b58 commit 1561f9e
Show file tree
Hide file tree
Showing 7 changed files with 217 additions and 1 deletion.
25 changes: 24 additions & 1 deletion src/client/gateway/shard.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ use ::internal::prelude::*;
use ::internal::ws_impl::{ReceiverExt, SenderExt};
use ::model::event::{Event, GatewayEvent, ReadyEvent};
use ::model::{Game, GuildId, OnlineStatus};
use ::utils;

#[cfg(feature="cache")]
use ::client::CACHE;
Expand Down Expand Up @@ -482,7 +483,7 @@ impl Shard {
Ok(())
}

/// Requests that one or multiple [`Guild`]s be synced.
/// Requests that one or multiple [`Guild`]s be chunked.
///
/// This will ask Discord to start sending member chunks for large guilds
/// (250 members+). If a guild is over 250 members, then a full member list
Expand All @@ -507,6 +508,28 @@ impl Shard {
let _ = self.keepalive_channel.send(GatewayStatus::SendMessage(msg));
}

/// Calculates the number of guilds that the shard is responsible for.
///
/// If sharding is not being used (i.e. 1 shard), then the total number of
/// guilds in the [`Cache`] will be used.
///
/// **Note**: Requires the `cache` feature be enabled.
///
/// [`Cache`]: ../../ext/cache/struct.Cache.html
#[cfg(feature="cache")]
pub fn guilds_handled(&self) -> u16 {
let cache = CACHE.read().unwrap();

if let Some((shard_id, shard_count)) = self.shard_info.map(|s| (s[0], s[1])) {
cache.guilds
.keys()
.filter(|guild_id| utils::shard_id(guild_id.0, shard_count) == shard_id)
.count() as u16
} else {
cache.guilds.len() as u16
}
}

#[allow(unused_variables)]
fn handle_dispatch(&mut self, event: &Event) {
#[cfg(feature="voice")]
Expand Down
4 changes: 4 additions & 0 deletions src/ext/cache/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,8 @@ pub struct Cache {
/// A map of direct message channels that the current user has open with
/// other users.
pub private_channels: HashMap<ChannelId, Arc<RwLock<PrivateChannel>>>,
/// The total number of shards being used by the bot.
pub shard_count: u64,
/// A list of guilds which are "unavailable". Refer to the documentation for
/// [`Event::GuildUnavailable`] for more information on when this can occur.
///
Expand Down Expand Up @@ -877,6 +879,7 @@ impl Cache {
}

self.presences.extend(ready.presences);
self.shard_count = ready.shard.map_or(1, |s| s[1]);
self.user = ready.user;
}

Expand Down Expand Up @@ -938,6 +941,7 @@ impl Default for Cache {
notes: HashMap::default(),
presences: HashMap::default(),
private_channels: HashMap::default(),
shard_count: 1,
unavailable_guilds: HashSet::default(),
user: CurrentUser {
avatar: None,
Expand Down
43 changes: 43 additions & 0 deletions src/model/guild/guild_id.rs
Original file line number Diff line number Diff line change
Expand Up @@ -420,6 +420,49 @@ impl GuildId {
rest::edit_member(self.0, user_id.into().0, &map)
}

/// Returns the Id of the shard associated with the guild.
///
/// When the cache is enabled this will automatically retrieve the total
/// number of shards.
///
/// **Note**: When the cache is enabled, this function unlocks the cache to
/// retrieve the total number of shards in use. If you already have the
/// total, consider using [`utils::shard_id`].
///
/// [`utils::shard_id`]: ../utils/fn.shard_id.html
#[cfg(feature="cache")]
#[inline]
pub fn shard_id(&self) -> u64 {
::utils::shard_id(self.0, CACHE.read().unwrap().shard_count)
}

/// Returns the Id of the shard associated with the guild.
///
/// When the cache is enabled this will automatically retrieve the total
/// number of shards.
///
/// When the cache is not enabled, the total number of shards being used
/// will need to be passed.
///
/// # Examples
///
/// Retrieve the Id of the shard for a guild with Id `81384788765712384`,
/// using 17 shards:
///
/// ```rust
/// use serenity::model::GuildId;
/// use serenity::utils;
///
/// let guild_id = GuildId(81384788765712384);
///
/// assert_eq!(guild_id.shard_id(17), 7);
/// ```
#[cfg(not(feature="cache"))]
#[inline]
pub fn shard_id(&self, shard_count: u64) -> u64 {
::utils::shard_id(self.0, shard_count)
}

/// Starts an integration sync for the given integration Id.
///
/// Requires the [Manage Guild] permission.
Expand Down
42 changes: 42 additions & 0 deletions src/model/guild/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -887,6 +887,48 @@ impl Guild {
permissions
}

/// Returns the Id of the shard associated with the guild.
///
/// When the cache is enabled this will automatically retrieve the total
/// number of shards.
///
/// **Note**: When the cache is enabled, this function unlocks the cache to
/// retrieve the total number of shards in use. If you already have the
/// total, consider using [`utils::shard_id`].
///
/// [`utils::shard_id`]: ../utils/fn.shard_id.html
#[cfg(feature="cache")]
#[inline]
pub fn shard_id(&self) -> u64 {
self.id.shard_id()
}

/// Returns the Id of the shard associated with the guild.
///
/// When the cache is enabled this will automatically retrieve the total
/// number of shards.
///
/// When the cache is not enabled, the total number of shards being used
/// will need to be passed.
///
/// # Examples
///
/// Retrieve the Id of the shard for a guild with Id `81384788765712384`,
/// using 17 shards:
///
/// ```rust,ignore
/// use serenity::utils;
///
/// // assumes a `guild` has already been bound
///
/// assert_eq!(guild.shard_id(17), 7);
/// ```
#[cfg(not(feature="cache"))]
#[inline]
pub fn shard_id(&self, shard_count: u64) -> u64 {
self.id.shard_id(shard_count)
}

/// Returns the formatted URL of the guild's splash image, if one exists.
pub fn splash_url(&self) -> Option<String> {
self.icon.as_ref().map(|icon|
Expand Down
42 changes: 42 additions & 0 deletions src/model/guild/partial_guild.rs
Original file line number Diff line number Diff line change
Expand Up @@ -404,6 +404,48 @@ impl PartialGuild {
self.id.move_member(user_id, channel_id)
}

/// Returns the Id of the shard associated with the guild.
///
/// When the cache is enabled this will automatically retrieve the total
/// number of shards.
///
/// **Note**: When the cache is enabled, this function unlocks the cache to
/// retrieve the total number of shards in use. If you already have the
/// total, consider using [`utils::shard_id`].
///
/// [`utils::shard_id`]: ../utils/fn.shard_id.html
#[cfg(feature="cache")]
#[inline]
pub fn shard_id(&self) -> u64 {
self.id.shard_id()
}

/// Returns the Id of the shard associated with the guild.
///
/// When the cache is enabled this will automatically retrieve the total
/// number of shards.
///
/// When the cache is not enabled, the total number of shards being used
/// will need to be passed.
///
/// # Examples
///
/// Retrieve the Id of the shard for a guild with Id `81384788765712384`,
/// using 17 shards:
///
/// ```rust,ignore
/// use serenity::utils;
///
/// // assumes a `guild` has already been bound
///
/// assert_eq!(guild.shard_id(17), 7);
/// ```
#[cfg(not(feature="cache"))]
#[inline]
pub fn shard_id(&self, shard_count: u64) -> u64 {
self.id.shard_id(shard_count)
}

/// Returns the formatted URL of the guild's splash image, if one exists.
pub fn splash_url(&self) -> Option<String> {
self.icon.as_ref().map(|icon|
Expand Down
44 changes: 44 additions & 0 deletions src/model/invite.rs
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,50 @@ pub struct InviteGuild {
pub splash_hash: Option<String>,
}

impl InviteGuild {
/// Returns the Id of the shard associated with the guild.
///
/// When the cache is enabled this will automatically retrieve the total
/// number of shards.
///
/// **Note**: When the cache is enabled, this function unlocks the cache to
/// retrieve the total number of shards in use. If you already have the
/// total, consider using [`utils::shard_id`].
///
/// [`utils::shard_id`]: ../utils/fn.shard_id.html
#[cfg(feature="cache")]
#[inline]
pub fn shard_id(&self) -> u64 {
self.id.shard_id()
}

/// Returns the Id of the shard associated with the guild.
///
/// When the cache is enabled this will automatically retrieve the total
/// number of shards.
///
/// When the cache is not enabled, the total number of shards being used
/// will need to be passed.
///
/// # Examples
///
/// Retrieve the Id of the shard for a guild with Id `81384788765712384`,
/// using 17 shards:
///
/// ```rust,ignore
/// use serenity::utils;
///
/// // assumes a `guild` has already been bound
///
/// assert_eq!(guild.shard_id(17), 7);
/// ```
#[cfg(not(feature="cache"))]
#[inline]
pub fn shard_id(&self, shard_count: u64) -> u64 {
self.id.shard_id(shard_count)
}
}

/// Detailed information about an invite.
/// This information can only be retrieved by anyone with the [Manage Guild]
/// permission. Otherwise, a minimal amount of information can be retrieved via
Expand Down
18 changes: 18 additions & 0 deletions src/utils/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -233,3 +233,21 @@ pub fn parse_quotes(s: &str) -> Vec<String> {

args
}

/// Calculates the Id of the shard responsible for a guild, given its Id and
/// total number of shards used.
///
/// # Examples
///
/// Retrieve the Id of the shard for a guild with Id `81384788765712384`, using
/// 17 shards:
///
/// ```rust
/// use serenity::utils;
///
/// assert_eq!(utils::shard_id(81384788765712384, 17), 7);
/// ```
#[inline]
pub fn shard_id(guild_id: u64, shard_count: u64) -> u64 {
(guild_id >> 22) % shard_count
}

0 comments on commit 1561f9e

Please sign in to comment.