diff options
| -rw-r--r-- | lib/client/shardManager.ml | 69 |
1 files changed, 43 insertions, 26 deletions
diff --git a/lib/client/shardManager.ml b/lib/client/shardManager.ml index 16aa531..cf849bc 100644 --- a/lib/client/shardManager.ml +++ b/lib/client/shardManager.ml @@ -42,27 +42,26 @@ module Opcode = struct | HEARTBEAT_ACK -> "HEARTBEAT_ACK" end +module ShardMap = Map.Make( + struct + type t = int + let compare: int -> int -> int = Pervasives.compare + end +) + module Shard = struct type t = { send: (Frame.t -> unit Lwt.t); recv: (unit -> Frame.t Lwt. t); + id: int; hb_interval: int; session_id: string; seq: int; } - let connect uri = - let url = Uri.to_string uri in - let ip = Ipaddr.V4 Ipaddr.V4.any in - Websocket_lwt.with_connection (`TLS (`Hostname url, `IP ip, `Port 443)) uri - >|= fun (recv, send) -> - (* Start heartbeat sequencing *) - { send; recv; hb_interval = 42500; session_id = ""; seq = 0; } - let send payload shard = payload |> shard.send - |> ignore let wrap_payload d op = `Assoc [ @@ -70,8 +69,11 @@ module Shard = struct ("d", d) ] - let identify ?(threshold=250) shard id total token = - let p = wrap_payload(`Assoc [ + let create_frame content = + Frame.create ~content () + + let identify ?(threshold=250) shard total token = + let p = wrap_payload (`Assoc [ ("token", `String token); ("properties", `Assoc [ ("$os", `String Sys.os_type); @@ -79,9 +81,9 @@ module Shard = struct ("$device", `String "animus"); ]); ("large_threshold", `Int threshold); - ("shard", `List [`Int id; `Int total]); + ("shard", `List [`Int shard.id; `Int total]); ]) (Opcode.to_int IDENTIFY) in - let p = Frame.create ~content:(Yojson.Basic.to_string p) () in + let p = create_frame (Yojson.Basic.to_string p) in send p shard let resume shard token = @@ -90,20 +92,26 @@ module Shard = struct ("session_id", `String shard.session_id); ("seq", `Int shard.seq); ]) (Opcode.to_int RESUME) in - let p = Frame.create ~content:(Yojson.Basic.to_string p) () in + let p = create_frame (Yojson.Basic.to_string p) in send p shard let heartbeat shard = - let p = Frame.create () in + let p = wrap_payload (`Int shard.seq) (Opcode.to_int HEARTBEAT) in + let p = create_frame (Yojson.Basic.to_string p) in send p shard -end -module ShardMap = Map.Make( - struct - type t = int - let compare: int -> int -> int = Pervasives.compare - end -) + (* Use options *) + let connect ~options uri id total token = + let url = Uri.to_string uri in + let ip = Ipaddr.V4 Ipaddr.V4.any in + Websocket_lwt.with_connection (`TLS (`Hostname url, `IP ip, `Port 443)) uri + >|= fun (recv, send) -> + let shard = { send; recv; hb_interval = 42500; id; session_id = ""; seq = 0; } in + heartbeat shard >>= (fun _ -> + identify shard total token (* This needs to be handled, I'm just pleasing the compiler *) + ) |> ignore; + shard +end type t = { shards: Shard.t ShardMap.t; @@ -111,16 +119,25 @@ type t = { token: string; } -let create_shard manager = - Shard.connect manager.gateway_url - >|= fun shard -> +let create_shard ?(options=[]) manager = let id = (ShardMap.cardinal manager.shards) + 1 in + Shard.connect ~options manager.gateway_url id (ShardMap.cardinal manager.shards) manager.token + >|= fun shard -> ShardMap.add id shard manager.shards +let update_shard id shard manager = + match ShardMap.exists (fun k _ -> id == k) manager.shards with + | true -> ShardMap.add id shard manager.shards + | false -> manager.shards + +let heartbeat manager id = + let shard = ShardMap.find id manager.shards in + Shard.heartbeat shard + let identify manager id = let total = ShardMap.cardinal manager.shards in let shard = ShardMap.find id manager.shards in - Shard.identify shard id total manager.token + Shard.identify shard total manager.token let resume manager id = let shard = ShardMap.find id manager.shards in |