Merge pull request #107 from rslota/add_ping

Add ping support
joedevivo · Mar 17, 2017 · 52b3853 · 52b3853
2 parents 759cc6a + 20f0096
commit 52b3853
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 6 deletions.
diff --git a/src/h2_client.erl b/src/h2_client.erl
@@ -22,6 +22,7 @@
          start_ssl_upgrade_link/4,
          stop/1,
          send_request/3,
+         send_ping/1,
          sync_request/3,
          get_response/2
         ]).
@@ -144,6 +145,9 @@ send_request(CliPid, Headers, Body) ->
             {ok, StreamId}
     end.
 
+send_ping(CliPid) ->
+    h2_connection:send_ping(CliPid).
+
 -spec get_response(pid(), stream_id()) ->
                           {ok, {hpack:header(), iodata()}}
                            | not_ready

diff --git a/src/h2_connection.erl b/src/h2_connection.erl
@@ -19,6 +19,7 @@
          send_body/3,
          send_body/4,
          send_request/3,
+         send_ping/1,
          is_push/1,
          new_stream/1,
          new_stream/2,
@@ -90,7 +91,8 @@
           stream_callback_mod = application:get_env(chatterbox, stream_callback_mod, chatterbox_static_stream) :: module(),
           buffer = empty :: empty | {binary, binary()} | {frame, h2_frame:header(), binary()},
           continuation = undefined :: undefined | #continuation_state{},
-          flow_control = auto :: auto | manual
+          flow_control = auto :: auto | manual,
+          pings = #{} :: #{binary() => {pid(), non_neg_integer()}}
 }).
 
 -type connection() :: #connection{}.
@@ -243,6 +245,11 @@ send_request(Pid, Headers, Body) ->
     gen_fsm:sync_send_all_state_event(Pid, {send_request, self(), Headers, Body}, infinity),
     ok.
 
+-spec send_ping(pid()) -> ok.
+send_ping(Pid) ->
+    gen_fsm:sync_send_all_state_event(Pid, {send_ping, self()}, infinity),
+    ok.
+
 -spec get_peer(pid()) ->
     {ok, {inet:ip_address(), inet:port_number()}} | {error, term()}.
 get_peer(Pid) ->
@@ -748,13 +755,21 @@ route_frame({H, Ping},
     Ack = h2_frame_ping:ack(Ping),
     socksend(Conn, h2_frame:to_binary(Ack)),
     {next_state, connected, Conn};
-route_frame({H, _Payload},
-            #connection{}=Conn)
+route_frame({H, Payload},
+            #connection{pings = Pings}=Conn)
     when H#frame_header.type == ?PING,
          ?IS_FLAG((H#frame_header.flags), ?FLAG_ACK) ->
-    lager:debug("[~p] Received PING ACK",
-               [Conn#connection.type]),
-    {next_state, connected, Conn};
+    case maps:get(h2_frame_ping:to_binary(Payload), Pings, undefined) of
+        undefined ->
+            lager:debug("[~p] Received unknown PING ACK",
+                        [Conn#connection.type]);
+        {NotifyPid, _} ->
+            lager:debug("[~p] Received PING ACK",
+                        [Conn#connection.type]),
+            NotifyPid ! {'PONG', self()}
+    end,
+    NextPings = maps:remove(Payload, Pings),
+    {next_state, connected, Conn#connection{pings = NextPings}};
 route_frame({H=#frame_header{stream_id=0}, _Payload},
             #connection{}=Conn)
     when H#frame_header.type == ?GOAWAY ->
@@ -1154,6 +1169,22 @@ handle_sync_event({send_request, NotifyPid, Headers, Body}, _F,
         {error, Code} ->
             {reply, {error, Code}, StateName, Conn}
     end;
+handle_sync_event({send_ping, NotifyPid}, _F,
+        StateName,
+        #connection{pings = Pings} = Conn) ->
+    PingValue = crypto:strong_rand_bytes(8),
+    Frame = h2_frame_ping:new(PingValue),
+    Headers = #frame_header{stream_id = 0, flags = 16#0},
+    Binary = h2_frame:to_binary({Headers, Frame}),
+
+    case socksend(Conn, Binary) of
+        ok ->
+            NextPings = maps:put(PingValue, {NotifyPid, erlang:monotonic_time(milli_seconds)}, Pings),
+            NextConn = Conn#connection{pings = NextPings},
+            {reply, ok, StateName, NextConn};
+        {error, _Reason} = Err ->
+            {reply, Err, StateName, Conn}
+    end;
 handle_sync_event(_E, _F, StateName,
                   #connection{}=Conn) ->
     {next_state, StateName, Conn}.