Skip to main content

twilight_gateway/
lib.rs

1#![cfg_attr(docsrs, feature(doc_cfg))]
2#![doc = include_str!("../README.md")]
3#![warn(
4    clippy::missing_const_for_fn,
5    clippy::missing_docs_in_private_items,
6    clippy::pedantic,
7    missing_docs,
8    unsafe_code
9)]
10#![allow(
11    clippy::module_name_repetitions,
12    clippy::must_use_candidate,
13    clippy::unnecessary_wraps
14)]
15
16pub mod error;
17
18mod channel;
19mod command;
20#[cfg(any(feature = "zlib", feature = "zstd"))]
21mod compression;
22mod config;
23mod event;
24mod json;
25mod latency;
26mod message;
27mod ratelimiter;
28mod session;
29mod shard;
30mod stream;
31
32pub use self::{
33    channel::MessageSender,
34    command::Command,
35    config::{Config, ConfigBuilder},
36    event::EventTypeFlags,
37    json::parse,
38    latency::Latency,
39    message::Message,
40    ratelimiter::CommandRatelimiter,
41    session::Session,
42    shard::{Shard, ShardState},
43    stream::StreamExt,
44};
45pub use twilight_model::gateway::{CloseFrame, Intents, ShardId};
46
47#[doc(no_inline)]
48pub use twilight_gateway_queue as queue;
49#[doc(no_inline)]
50pub use twilight_model::gateway::event::{Event, EventType};
51
52#[cfg(feature = "twilight-http")]
53use self::error::{StartRecommendedError, StartRecommendedErrorType};
54#[cfg(feature = "twilight-http")]
55use twilight_http::Client;
56
57/// Discord Gateway API version used by this crate.
58pub const API_VERSION: u8 = 10;
59
60/// Creates an iterator of a single bucket's worth of shard identifiers.
61///
62/// Each bucket holds a consecutive range of identifiers and all identifiers
63/// share the same `total` value.
64///
65/// # Strategy
66///
67/// Shards may be bucketed per-thread for a thread-per-core architecture and/or
68/// per-machine for horizontal scaling.
69///
70/// # Examples
71///
72/// Create 1 bucket with the recommended shard count:
73///
74/// ```no_run
75/// use std::env;
76/// use twilight_http::Client;
77///
78/// # #[tokio::main(flavor = "current_thread")]
79/// # async fn main() -> anyhow::Result<()> {
80/// let http = Client::new(env::var("TOKEN")?);
81/// let info = http.gateway().authed().await?.model().await?;
82///
83/// let shards = twilight_gateway::bucket(0, 1, info.shards);
84/// assert_eq!(shards.len(), info.shards as usize);
85/// # anyhow::Ok(())
86/// # }
87/// ```
88///
89/// Create 2 buckets with 25 identifiers:
90///
91/// ```
92/// let bucket_1 = twilight_gateway::bucket(0, 2, 25);
93/// let bucket_2 = twilight_gateway::bucket(1, 2, 25);
94///
95/// assert_eq!(bucket_1.len(), 13);
96/// assert_eq!(bucket_2.len(), 12);
97/// ```
98///
99/// # Panics
100///
101/// Panics if the bucket id is greater than or equal to the total number of
102/// buckets.
103#[track_caller]
104pub fn bucket(
105    bucket_id: u16,
106    buckets: u16,
107    shards: u32,
108) -> impl DoubleEndedIterator<Item = ShardId> + ExactSizeIterator {
109    let bucket_id = u32::from(bucket_id);
110    let buckets = u32::from(buckets);
111    assert!(bucket_id < buckets, "bucket_id must be less than buckets");
112
113    let (q, r) = (shards / buckets, shards % buckets);
114
115    let len = q + u32::from(bucket_id < r);
116    let start = bucket_id * q + r.min(bucket_id);
117
118    (start..start + len).map(move |id| ShardId::new(id, shards))
119}
120
121/// Create a single bucket's worth of shards.
122///
123/// Passing a primary config is required. Further customization of this config
124/// may be performed in the callback.
125///
126/// Internally calls [`create_iterator`] with `(bucket_id..total).step_by(concurrency)`.
127///
128/// # Panics
129///
130/// Panics if `bucket_id >= total`, `bucket_id >= concurrency`, or `concurrency >= total`.
131///
132/// Panics if loading TLS certificates fails.
133#[deprecated = "creates non-consecutive shards; use `bucket` instead"]
134#[track_caller]
135pub fn create_bucket<F, Q>(
136    bucket_id: u16,
137    concurrency: u16,
138    total: u32,
139    config: Config<Q>,
140    per_shard_config: F,
141) -> impl ExactSizeIterator<Item = Shard<Q>>
142where
143    F: FnMut(ShardId, ConfigBuilder<Q>) -> Config<Q>,
144    Q: Clone,
145{
146    assert!(
147        u32::from(bucket_id) < total,
148        "bucket id must be less than the total"
149    );
150    assert!(
151        bucket_id < concurrency,
152        "bucket id must be less than concurrency"
153    );
154    assert!(
155        (u32::from(concurrency)) < total,
156        "concurrency must be less than the total"
157    );
158
159    #[allow(deprecated)]
160    create_iterator(
161        (u32::from(bucket_id)..total).step_by(concurrency.into()),
162        total,
163        config,
164        per_shard_config,
165    )
166}
167
168/// Create a iterator of shards.
169///
170/// Passing a primary config is required. Further customization of this config
171/// may be performed in the callback.
172///
173/// # Examples
174///
175/// Start 10 out of 10 shards and count them:
176///
177/// ```no_run
178/// use std::{collections::HashMap, env, sync::Arc};
179/// use twilight_gateway::{Config, Intents};
180///
181/// let token = env::var("DISCORD_TOKEN")?;
182///
183/// let config = Config::new(token.clone(), Intents::GUILDS);
184/// let shards = twilight_gateway::create_iterator(0..10, 10, config, |_, builder| builder.build());
185///
186/// assert_eq!(shards.len(), 10);
187/// # Ok::<(), Box<dyn std::error::Error>>(())
188/// ```
189///
190/// # Panics
191///
192/// Panics if `range` contains values larger than `total`.
193///
194/// Panics if loading TLS certificates fails.
195#[deprecated = "use `bucket` instead"]
196#[track_caller]
197pub fn create_iterator<F, Q>(
198    numbers: impl ExactSizeIterator<Item = u32>,
199    total: u32,
200    config: Config<Q>,
201    mut per_shard_config: F,
202) -> impl ExactSizeIterator<Item = Shard<Q>>
203where
204    F: FnMut(ShardId, ConfigBuilder<Q>) -> Config<Q>,
205    Q: Clone,
206{
207    numbers.map(move |index| {
208        let id = ShardId::new(index, total);
209        let config = per_shard_config(id, ConfigBuilder::from(config.clone()));
210
211        Shard::with_config(id, config)
212    })
213}
214
215/// Create a range of shards from Discord's recommendation.
216///
217/// Passing a primary config is required. Further customization of this config
218/// may be performed in the callback.
219///
220/// Internally calls [`create_iterator`] with the values from [`GetGatewayAuthed`].
221///
222/// # Errors
223///
224/// Returns a [`StartRecommendedErrorType::Deserializing`] error type if the
225/// response body failed to deserialize.
226///
227/// Returns a [`StartRecommendedErrorType::Request`] error type if the request
228/// failed to complete.
229///
230/// # Panics
231///
232/// Panics if loading TLS certificates fails.
233///
234/// [`GetGatewayAuthed`]: twilight_http::request::GetGatewayAuthed
235#[cfg(feature = "twilight-http")]
236pub async fn create_recommended<F, Q>(
237    client: &Client,
238    config: Config<Q>,
239    per_shard_config: F,
240) -> Result<impl ExactSizeIterator<Item = Shard<Q>> + use<F, Q>, StartRecommendedError>
241where
242    F: FnMut(ShardId, ConfigBuilder<Q>) -> Config<Q>,
243    Q: Clone,
244{
245    let request = client.gateway().authed();
246    let response = request.await.map_err(|source| StartRecommendedError {
247        kind: StartRecommendedErrorType::Request,
248        source: Some(Box::new(source)),
249    })?;
250    let info = response
251        .model()
252        .await
253        .map_err(|source| StartRecommendedError {
254            kind: StartRecommendedErrorType::Deserializing,
255            source: Some(Box::new(source)),
256        })?;
257
258    #[allow(deprecated)]
259    Ok(create_iterator(
260        0..info.shards,
261        info.shards,
262        config,
263        per_shard_config,
264    ))
265}