add websocket ping support

examples/websocket.xml.in

@@ -31,10 +31,18 @@
  <websocket type="message">{"uid":"%%_uid%%", "data":"data"}</websocket>
  </request>
 
+ <request>
+ <websocket type="ping" ack="bidi_ack">hi</websocket>
+ </request>
+
  <request>
  <websocket type="message" ack="no_ack">{"key":"value"}</websocket>
  </request>
 
+ <for var="i" from="1" to="3" incr="1">
+ <thinktime value="30"/>
+ </for>
+
  <request>
  <websocket type="close"></websocket>
  </request>

include/ts_websocket.hrl

@@ -44,7 +44,8 @@
 
 -record(websocket_session, {
  status, % status of handshake
- accept % Sec-Websocket-Accept header value
+ accept, % Sec-Websocket-Accept header value
+ ping % ping payload, expected in pong response
  }).
 
 %% opcode of websocket

src/lib/websocket.erl

@@ -28,7 +28,7 @@
 -author('[email protected]').
 
 -export([get_handshake/6, check_handshake/2, encode_binary/1, encode_text/1,
- encode_close/1, encode/2, decode/1]).
+ encode_close/1, encode_ping/1, encode_pong/1, encode/2, decode/1]).
 
 -include("ts_profile.hrl").
 -include("ts_config.hrl").
@@ -113,6 +113,12 @@ encode_close(Reason) ->
  Data = <<StatusCode/binary, Reason/binary>>,
  encode(Data, ?OP_CLOSE).
 
+encode_ping(Data) ->
+ encode(Data, ?OP_PING).
+
+encode_pong(Data) ->
+ encode(Data, ?OP_PONG).
+
 encode(Data, Opcode) ->
  Key = crypto:strong_rand_bytes(4),
  PayloadLen = erlang:size(Data),

src/tsung/ts_websocket.erl

@@ -98,7 +98,11 @@ get_message(#websocket_request{type = message, data = Data, frame = Frame},
 get_message(#websocket_request{type = close},
  #state_rcv{session = WebsocketSession})
  when WebsocketSession#websocket_session.status == connected ->
- {websocket:encode_close(<<"close">>), WebsocketSession}.
+ {websocket:encode_close(<<"close">>), WebsocketSession};
+get_message(#websocket_request{type = ping, data = Data},
+ #state_rcv{session = WebsocketSession})
+ when WebsocketSession#websocket_session.status == connected ->
+ {websocket:encode_ping(list_to_binary(Data)), WebsocketSession#websocket_session{ping = list_to_binary(Data)}}.
 
 %%----------------------------------------------------------------------
 %% Function: parse/2
@@ -137,10 +141,16 @@ parse(Data, State=#state_rcv{acc = [],
 %% normal websocket message
 parse(Data, State=#state_rcv{acc = [], session = WebsocketSession})
  when WebsocketSession#websocket_session.status == connected ->
+ Expected = WebsocketSession#websocket_session.ping,
  case websocket:decode(Data) of
- {?OP_CLOSE, _Reason, _} ->
- ?DebugF("receive close from server: ~p~n", [_Reason]),
- {State#state_rcv{ack_done = true}, [], true};
+ {?OP_PONG, Expected, << >>} ->
+ ?DebugF("received pong: [~p] ~n", [Expected]),
+ ts_mon_cache:add({count, websocket_pong}),
+ {State#state_rcv{ack_done = true, session = WebsocketSession#websocket_session{ping = none}}, [], false};
+ {?OP_PONG, _Unexpected, << >>} ->
+ ts_mon_cache:add({count, error_websocket_pong}),
+ ?DebugF("unexpected pong: [~p] ~n", [_Unexpected]),
+ {State#state_rcv{ack_done = true}, [], false};
  {_Opcode, _Payload, Left} ->
  ?DebugF("receive from server: ~p ~p~n", [_Opcode, _Payload]),
  {State#state_rcv{ack_done = true, acc = Left}, [], false};
@@ -154,8 +164,19 @@ parse(Data, State=#state_rcv{acc = Acc, datasize = DataSize}) ->
  parse(<< Acc/binary, Data/binary >>,
  State#state_rcv{acc = [], datasize = NewSize}).
 
-parse_bidi(Data, State) ->
- ts_plugin:parse_bidi(Data, State).
+parse_bidi(Data, State=#state_rcv{acc=[]}) ->
+ case websocket:decode(Data) of
+ {?OP_CLOSE, _Reason, _} ->
+ ?DebugF("receive close from server: ~p~n", [_Reason]),
+ {websocket:encode_close(<< >>), State, think};
+ {?OP_PING, Payload, Tail} ->
+ ts_mon_cache:add({count, websocket_ping}),
+ {websocket:encode_pong(Payload), State#state_rcv{acc = Tail}, think};
+ {_Opcode, _Payload, Tail} ->
+ {nodata, State#state_rcv{acc = Tail}, think}
+ end;
+parse_bidi(Data, State=#state_rcv{acc=Acc}) ->
+ parse_bidi(<< Acc/binary, Data/binary >>, State#state_rcv{acc = []}).
 
 %%----------------------------------------------------------------------
 %% Function: parse_config/2

tsung-1.0.dtd

@@ -376,8 +376,8 @@ repeat | if | change_type | foreach | set_option | interaction | abort )*>
  >
 <!ELEMENT websocket (#PCDATA | http_header)* >
 <!ATTLIST websocket
- type (connect | message | close) #REQUIRED
- ack (no_ack | parse) #IMPLIED
+ type (connect | message | ping | close) #REQUIRED
+ ack (no_ack | parse | bidi_ack) #IMPLIED
  frame (binary | text) #IMPLIED
  origin CDATA ""
  subprotocols CDATA ""