Add rate limit handling

3 files changed, 86 insertions, 40 deletions

diff --git a/lib/http.ml b/lib/http.ml
index 3e10eb8..d2dff65 100644
--- a/lib/http.ml
+++ b/lib/http.ml
@@ -5,6 +5,9 @@ module Make(T : S.Token) = struct
 
     module Base = struct
         exception Invalid_Method
+        exception Bad_response_headers
+
+        let rl = ref Rl.empty
 
         let base_url = "https://discordapp.com/api/v7"
 
@@ -24,22 +27,35 @@ module Make(T : S.Token) = struct
                 "Content-Type", "application/json";
             ]
 
-        (* TODO Finish processor *)
-        let process_response ((_resp:Response.t), body) =
+        let process_response path ((resp:Response.t), body) =
+            (match Response.headers resp
+            |> Rl.rl_of_header with
+            | Some r -> Mvar.put (Rl.find_exn !rl path) r
+            | None -> raise Bad_response_headers)
+            >>= fun () ->
             body |> Cohttp_async.Body.to_string >>| Yojson.Safe.from_string
 
         let request ?(body=`Null) m path =
-            let uri = process_url path in
-            let headers = process_request_headers () in
-            let body = process_request_body body in
-            (match m with
-            | `DELETE -> Cohttp_async.Client.delete ~headers ~body uri
-            | `GET -> Cohttp_async.Client.get ~headers uri
-            | `PATCH -> Cohttp_async.Client.patch ~headers ~body uri
-            | `POST -> Cohttp_async.Client.post ~headers ~body uri
-            | `PUT -> Cohttp_async.Client.put ~headers ~body uri
-            | _ -> raise Invalid_Method)
-            >>= process_response
+            rl := Rl.update ~f:(function
+                | None -> Mvar.create ()
+                | Some r -> r
+            ) !rl path;
+            let limit = Rl.find_exn !rl path in
+            Mvar.take limit >>= fun limit ->
+            let process () =
+                let uri = process_url path in
+                let headers = process_request_headers () in
+                let body = process_request_body body in
+                (match m with
+                | `DELETE -> Cohttp_async.Client.delete ~headers ~body uri
+                | `GET -> Cohttp_async.Client.get ~headers uri
+                | `PATCH -> Cohttp_async.Client.patch ~headers ~body uri
+                | `POST -> Cohttp_async.Client.post ~headers ~body uri
+                | `PUT -> Cohttp_async.Client.put ~headers ~body uri
+                | _ -> raise Invalid_Method)
+                >>= process_response path
+            in if limit.remaining > 0 then process ()
+            else Clock.at (Core.Time.(Span.of_int_sec limit.reset |> of_span_since_epoch)) >>= process
     end
 
     let get_gateway () =
diff --git a/lib/rl.ml b/lib/rl.ml
new file mode 100644
index 0000000..316abcf
--- /dev/null
+++ b/lib/rl.ml
@@ -0,0 +1,29 @@
+open Core
+open Async
+
+module RouteMap = Map.Make(String)
+
+type rl = {
+    limit: int;
+    remaining: int;
+    reset: int;
+}
+
+type t = ((rl, read_write) Mvar.t) RouteMap.t
+
+let rl_of_header h =
+    let module C = Cohttp.Header in
+    match C.get h "X-RateLimit-Limit", C.get h "X-RateLimit-Remaining", C.get h "X-RateLimit-Reset" with
+    | Some lim, Some rem, Some re ->
+        let limit = Int.of_string lim in
+        let remaining = Int.of_string rem in
+        let reset = Int.of_string re in
+        Some { limit; remaining; reset; }
+    | _ -> None
+
+let update = RouteMap.update
+let empty : t = RouteMap.empty
+let find = RouteMap.find
+let find_exn m s = match find m s with
+    | Some r -> r
+    | None -> raise (Not_found_s (String.sexp_of_t s))
\ No newline at end of file
diff --git a/lib/s.ml b/lib/s.ml
index 95eaee3..9ac86ad 100644
--- a/lib/s.ml
+++ b/lib/s.ml
@@ -21,32 +21,32 @@ module type Dispatch = sig
     | READY of Yojson.Safe.json
     | RESUMED of Yojson.Safe.json
     | INVALID_SESSION of Yojson.Safe.json
-    | CHANNEL_CREATE of Channel.t
-    | CHANNEL_UPDATE of Channel.t
-    | CHANNEL_DELETE of Channel.t
+    | CHANNEL_CREATE of Channel_t.t
+    | CHANNEL_UPDATE of Channel_t.t
+    | CHANNEL_DELETE of Channel_t.t
     | CHANNEL_PINS_UPDATE of Yojson.Safe.json
-    | GUILD_CREATE of Guild.t
-    | GUILD_UPDATE of Guild.t
-    | GUILD_DELETE of Guild.t
-    | GUILD_BAN_ADD of Ban.t
-    | GUILD_BAN_REMOVE of Ban.t
+    | GUILD_CREATE of Guild_t.t
+    | GUILD_UPDATE of Guild_t.t
+    | GUILD_DELETE of Guild_t.t
+    | GUILD_BAN_ADD of Ban_t.t
+    | GUILD_BAN_REMOVE of Ban_t.t
     | GUILD_EMOJIS_UPDATE of Yojson.Safe.json
     | GUILD_INTEGRATIONS_UPDATE of Yojson.Safe.json
-    | GUILD_MEMBER_ADD of Member.t
-    | GUILD_MEMBER_REMOVE of Member.t
-    | GUILD_MEMBER_UPDATE of Member.t
-    | GUILD_MEMBERS_CHUNK of Member.t list
-    | GUILD_ROLE_CREATE of Role.t (* * Guild.t *)
-    | GUILD_ROLE_UPDATE of Role.t (* * Guild.t *)
-    | GUILD_ROLE_DELETE of Role.t (* * Guild.t *)
-    | MESSAGE_CREATE of Message.t
-    | MESSAGE_UPDATE of Message.t
-    | MESSAGE_DELETE of Message.t
-    | MESSAGE_BULK_DELETE of Message.t list
-    | MESSAGE_REACTION_ADD of (* Message.t * *) Reaction.t
-    | MESSAGE_REACTION_REMOVE of (* Message.t * *) Reaction.t
-    | MESSAGE_REACTION_REMOVE_ALL of (* Message.t * *) Reaction.t list
-    | PRESENCE_UPDATE of Presence.t
+    | GUILD_MEMBER_ADD of Member_t.t
+    | GUILD_MEMBER_REMOVE of Member_t.t
+    | GUILD_MEMBER_UPDATE of Member_t.t
+    | GUILD_MEMBERS_CHUNK of Member_t.t list
+    | GUILD_ROLE_CREATE of Role_t.t (* * Guild.t *)
+    | GUILD_ROLE_UPDATE of Role_t.t (* * Guild.t *)
+    | GUILD_ROLE_DELETE of Role_t.t (* * Guild.t *)
+    | MESSAGE_CREATE of Message_t.t
+    | MESSAGE_UPDATE of Message_t.t
+    | MESSAGE_DELETE of Message_t.t
+    | MESSAGE_BULK_DELETE of Message_t.t list
+    | MESSAGE_REACTION_ADD of (* Message.t * *) Reaction_t.t
+    | MESSAGE_REACTION_REMOVE of (* Message.t * *) Reaction_t.t
+    | MESSAGE_REACTION_REMOVE_ALL of (* Message.t * *) Reaction_t.t list
+    | PRESENCE_UPDATE of Presence_t.t
     | TYPING_START of Yojson.Safe.json
     | USER_UPDATE of Yojson.Safe.json
     | VOICE_STATE_UPDATE of Yojson.Safe.json
@@ -55,8 +55,8 @@ module type Dispatch = sig
 
     exception Invalid_event of string
 
-    val event_of_string : contents:Yojson.Safe.json -> string -> dispatch_event
-    val dispatch : ev:string -> Yojson.Safe.json -> unit
+    val event_of_string : contents:string -> string -> dispatch_event
+    val dispatch : ev:string -> string -> unit
 end
 
 module type Http = sig
@@ -72,6 +72,7 @@ module type Http = sig
         val process_request_headers : unit -> Cohttp.Header.t
 
         val process_response :
+            string ->
             Cohttp_async.Response.t * Cohttp_async.Body.t ->
             Yojson.Safe.json Deferred.t
 
@@ -277,7 +278,7 @@ module type Sharder = sig
         val request_guild_members :
             ?query:string ->
             ?limit:int ->
-            guild:Snowflake.t ->
+            guild:Snowflake_t.t ->
             shard ->
             shard Deferred.t
 
@@ -301,7 +302,7 @@ module type Sharder = sig
     val request_guild_members :
         ?query:string ->
         ?limit:int ->
-        guild:Snowflake.t ->
+        guild:Snowflake_t.t ->
         t ->
         Shard.shard list Deferred.t
 end
\ No newline at end of file