3939import static org .junit .Assert .assertEquals ;
4040import static org .junit .Assert .assertFalse ;
4141import static org .junit .Assert .assertNotNull ;
42+ import static org .junit .Assert .assertNull ;
4243import static org .junit .Assert .assertTrue ;
4344import static org .junit .Assert .fail ;
4445
@@ -181,7 +182,7 @@ address, NioSocketChannel.class, channelOptions, group, newNegotiator(),
181182 public void overrideDefaultUserAgent () throws Exception {
182183 startServer ();
183184 NettyClientTransport transport = newTransport (newNegotiator (),
184- DEFAULT_MAX_MESSAGE_SIZE , GrpcUtil .DEFAULT_MAX_HEADER_LIST_SIZE , "testUserAgent" );
185+ DEFAULT_MAX_MESSAGE_SIZE , GrpcUtil .DEFAULT_MAX_HEADER_LIST_SIZE , "testUserAgent" , true );
185186 transport .start (clientTransportListener );
186187
187188 new Rpc (transport , new Metadata ()).halfClose ().waitForResponse ();
@@ -198,7 +199,7 @@ public void maxMessageSizeShouldBeEnforced() throws Throwable {
198199 startServer ();
199200 // Allow the response payloads of up to 1 byte.
200201 NettyClientTransport transport = newTransport (newNegotiator (),
201- 1 , GrpcUtil .DEFAULT_MAX_HEADER_LIST_SIZE , null );
202+ 1 , GrpcUtil .DEFAULT_MAX_HEADER_LIST_SIZE , null , true );
202203 transport .start (clientTransportListener );
203204
204205 try {
@@ -280,7 +281,7 @@ public void maxHeaderListSizeShouldBeEnforcedOnClient() throws Exception {
280281 startServer ();
281282
282283 NettyClientTransport transport =
283- newTransport (newNegotiator (), DEFAULT_MAX_MESSAGE_SIZE , 1 , null );
284+ newTransport (newNegotiator (), DEFAULT_MAX_MESSAGE_SIZE , 1 , null , true );
284285 transport .start (clientTransportListener );
285286
286287 try {
@@ -346,6 +347,30 @@ public void clientStreamGetsAttributes() throws Exception {
346347 assertEquals (address , rpc .stream .getAttributes ().get (Grpc .TRANSPORT_ATTR_REMOTE_ADDR ));
347348 }
348349
350+ @ Test
351+ public void keepAliveEnabled () throws Exception {
352+ startServer ();
353+ NettyClientTransport transport = newTransport (newNegotiator (), DEFAULT_MAX_MESSAGE_SIZE ,
354+ GrpcUtil .DEFAULT_MAX_HEADER_LIST_SIZE , null /* user agent */ , true /* keep alive */ );
355+ transport .start (clientTransportListener );
356+ Rpc rpc = new Rpc (transport ).halfClose ();
357+ rpc .waitForResponse ();
358+
359+ assertNotNull (transport .keepAliveManager ());
360+ }
361+
362+ @ Test
363+ public void keepAliveDisabled () throws Exception {
364+ startServer ();
365+ NettyClientTransport transport = newTransport (newNegotiator (), DEFAULT_MAX_MESSAGE_SIZE ,
366+ GrpcUtil .DEFAULT_MAX_HEADER_LIST_SIZE , null /* user agent */ , false /* keep alive */ );
367+ transport .start (clientTransportListener );
368+ Rpc rpc = new Rpc (transport ).halfClose ();
369+ rpc .waitForResponse ();
370+
371+ assertNull (transport .keepAliveManager ());
372+ }
373+
349374 private Throwable getRootCause (Throwable t ) {
350375 if (t .getCause () == null ) {
351376 return t ;
@@ -361,15 +386,18 @@ private ProtocolNegotiator newNegotiator() throws IOException {
361386 }
362387
363388 private NettyClientTransport newTransport (ProtocolNegotiator negotiator ) {
364- return newTransport (negotiator ,
365- DEFAULT_MAX_MESSAGE_SIZE , GrpcUtil . DEFAULT_MAX_HEADER_LIST_SIZE , null /* user agent */ );
389+ return newTransport (negotiator , DEFAULT_MAX_MESSAGE_SIZE , GrpcUtil . DEFAULT_MAX_HEADER_LIST_SIZE ,
390+ null /* user agent */ , true /* keep alive */ );
366391 }
367392
368- private NettyClientTransport newTransport (
369- ProtocolNegotiator negotiator , int maxMsgSize , int maxHeaderListSize , String userAgent ) {
393+ private NettyClientTransport newTransport (ProtocolNegotiator negotiator , int maxMsgSize ,
394+ int maxHeaderListSize , String userAgent , boolean enableKeepAlive ) {
370395 NettyClientTransport transport = new NettyClientTransport (
371396 address , NioSocketChannel .class , new HashMap <ChannelOption <?>, Object >(), group , negotiator ,
372397 DEFAULT_WINDOW_SIZE , maxMsgSize , maxHeaderListSize , authority , userAgent );
398+ if (enableKeepAlive ) {
399+ transport .enableKeepAlive (true , 1000 , 1000 );
400+ }
373401 transports .add (transport );
374402 return transport ;
375403 }
0 commit comments