chorus/src/limit.rs

use crate::api::limits::{Limit, LimitType, Limits};

use reqwest::{Client, RequestBuilder, Response};
use std::collections::{HashMap, VecDeque};

// Note: There seem to be some overlapping request limiters. We need to make sure that sending a
// request checks for all the request limiters that apply, and blocks if any of the limiters are 0

#[allow(dead_code)]
pub struct LimitedRequester {
    http: Client,
    requests: VecDeque<RequestBuilder>,
    last_reset_epoch: i64,
    limits_rate: HashMap<LimitType, Limit>,
}

impl LimitedRequester {
    /// Create a new `LimitedRequester`. `LimitedRequester`s use a `VecDeque` to store requests and
    /// send them to the server using a `Client`. It keeps track of the remaining requests that can
    /// be send within the `Limit` of an external API Ratelimiter, and looks at the returned request
    /// headers to see if it can find Ratelimit info to update itself.
    #[allow(dead_code)]
    pub async fn new(api_url: String) -> Self {
        LimitedRequester {
            http: Client::new(),
            requests: VecDeque::new(),
            last_reset_epoch: chrono::Utc::now().timestamp(),
            limits_rate: Limits::check_limits(api_url).await,
        }
    }

    fn add_to_queue(request: RequestBuilder, queue: &mut VecDeque<RequestBuilder>) {
        queue.push_back(request);
    }

    fn update_limit_entry(entry: &mut Limit, reset: u64, remaining: u64, limit: u64) {
        if reset != entry.reset {
            entry.reset = reset;
            entry.remaining = limit;
            entry.limit = limit;
        } else {
            entry.remaining = remaining;
            entry.limit = limit;
        }
    }

    fn decrement_limit_entry(entry: &mut Limit) {
        entry.remaining -= 1;
    }

    fn update_limits(&mut self, response: Response, limit_type: LimitType) {
        // TODO: Make this work
        let remaining = match response.headers().get("X-RateLimit-Remaining") {
            Some(remaining) => remaining.to_str().unwrap().parse::<u64>().unwrap(),
            None => return, //false,
        };
        let limit = match response.headers().get("X-RateLimit-Limit") {
            Some(limit) => limit.to_str().unwrap().parse::<u64>().unwrap(),
            None => return, //false,
        };
        let reset = match response.headers().get("X-RateLimit-Reset") {
            Some(reset) => reset.to_str().unwrap().parse::<u64>().unwrap(),
            None => return, //false,
        };

        let mut limits_copy = self.limits_rate.clone();
        let status = response.status();
        let status_str = status.as_str();

        if status_str.chars().next().unwrap() == '4' {
            limits_copy.get_mut(&LimitType::Error).unwrap().remaining -= 1;
        }

        limits_copy.get_mut(&LimitType::Global).unwrap().remaining -= 1;
        limits_copy.get_mut(&LimitType::Ip).unwrap().remaining -= 1;

        match limit_type {
            // Error, Global and Ip get handled seperately.
            LimitType::Error => {}
            LimitType::Global => {}
            LimitType::Ip => {}
            LimitType::AuthLogin => {
                let entry = limits_copy.get_mut(&LimitType::AuthLogin).unwrap();
                if reset != entry.reset {
                    LimitedRequester::update_limit_entry(entry, reset, limit, limit);
                } else {
                    LimitedRequester::update_limit_entry(entry, reset, remaining, limit);
                }
            }
            LimitType::AbsoluteRegister => {
                let entry = limits_copy.get_mut(&LimitType::AbsoluteRegister).unwrap();
                if reset != entry.reset {
                    LimitedRequester::update_limit_entry(entry, reset, limit, limit);
                } else {
                    LimitedRequester::update_limit_entry(entry, reset, remaining, limit);
                }
                // AbsoluteRegister and AuthRegister both need to be updated, if a Register event
                // happens.
                limits_copy
                    .get_mut(&LimitType::AuthRegister)
                    .unwrap()
                    .remaining -= 1;
            }
            LimitType::AuthRegister => {
                let entry = limits_copy.get_mut(&LimitType::AuthRegister).unwrap();
                if reset != entry.reset {
                    LimitedRequester::update_limit_entry(entry, reset, limit, limit);
                } else {
                    LimitedRequester::update_limit_entry(entry, reset, remaining, limit);
                }
                // AbsoluteRegister and AuthRegister both need to be updated, if a Register event
                // happens.
                limits_copy
                    .get_mut(&LimitType::AbsoluteRegister)
                    .unwrap()
                    .remaining -= 1;
            }
            LimitType::AbsoluteMessage => {
                let entry = limits_copy.get_mut(&LimitType::AbsoluteMessage).unwrap();
                if reset != entry.reset {
                    LimitedRequester::update_limit_entry(entry, reset, limit, limit);
                } else {
                    LimitedRequester::update_limit_entry(entry, reset, remaining, limit);
                }
            }
            LimitType::Channel => {
                let entry = limits_copy.get_mut(&LimitType::Channel).unwrap();
                if reset != entry.reset {
                    LimitedRequester::update_limit_entry(entry, reset, limit, limit);
                } else {
                    LimitedRequester::update_limit_entry(entry, reset, remaining, limit);
                }
            }
            LimitType::Guild => {
                let entry = limits_copy.get_mut(&LimitType::Guild).unwrap();
                if reset != entry.reset {
                    LimitedRequester::update_limit_entry(entry, reset, limit, limit);
                } else {
                    LimitedRequester::update_limit_entry(entry, reset, remaining, limit);
                }
            }
            LimitType::Webhook => {
                let entry = limits_copy.get_mut(&LimitType::Webhook).unwrap();
                if reset != entry.reset {
                    LimitedRequester::update_limit_entry(entry, reset, limit, limit);
                } else {
                    LimitedRequester::update_limit_entry(entry, reset, remaining, limit);
                }
            }
        }
    }
}