@@ -468,40 +468,50 @@ public void testDispatchRequestAddsAndFreesBytesOnError() {
468468 }
469469
470470 public void testDispatchRequestAddsAndFreesBytesOnlyOnceOnError () {
471- int contentLength = BREAKER_LIMIT .bytesAsInt ();
472- String content = randomAlphaOfLength ((int ) Math .round (contentLength / inFlightRequestsBreaker .getOverhead ()));
473- // we will produce an error in the rest handler and one more when sending the error response
474- RestRequest request = testRestRequest ("/error" , content , XContentType .JSON );
475- ExceptionThrowingChannel channel = new ExceptionThrowingChannel (request , true );
476-
477- restController .dispatchRequest (request , channel , client .threadPool ().getThreadContext ());
478-
479- assertEquals (0 , inFlightRequestsBreaker .getTrippedCount ());
480- assertEquals (0 , inFlightRequestsBreaker .getUsed ());
471+ try {
472+ RestController .PERMIT_DOUBLE_RESPONSE = true ;
473+ int contentLength = BREAKER_LIMIT .bytesAsInt ();
474+ String content = randomAlphaOfLength ((int ) Math .round (contentLength / inFlightRequestsBreaker .getOverhead ()));
475+ // we will produce an error in the rest handler and one more when sending the error response
476+ RestRequest request = testRestRequest ("/error" , content , XContentType .JSON );
477+ ExceptionThrowingChannel channel = new ExceptionThrowingChannel (request , true );
478+
479+ restController .dispatchRequest (request , channel , client .threadPool ().getThreadContext ());
480+
481+ assertEquals (0 , inFlightRequestsBreaker .getTrippedCount ());
482+ assertEquals (0 , inFlightRequestsBreaker .getUsed ());
483+ } finally {
484+ RestController .PERMIT_DOUBLE_RESPONSE = false ;
485+ }
481486 }
482487
483488 public void testDispatchRequestAddsAndFreesBytesOnlyOnceOnErrorDuringSend () {
484- int contentLength = Math .toIntExact (BREAKER_LIMIT .getBytes ());
485- String content = randomAlphaOfLength ((int ) Math .round (contentLength / inFlightRequestsBreaker .getOverhead ()));
486- // use a real recycler that tracks leaks and create some content bytes in the test handler to check for leaks
487- final BytesRefRecycler recycler = new BytesRefRecycler (new MockPageCacheRecycler (Settings .EMPTY ));
488- restController .registerHandler (
489- new Route (GET , "/foo" ),
490- (request , c , client ) -> new RestToXContentListener <>(c ).onResponse ((b , p ) -> b .startObject ().endObject ())
491- );
492- // we will produce an error in the rest handler and one more when sending the error response
493- RestRequest request = testRestRequest ("/foo" , content , XContentType .JSON );
494- ExceptionThrowingChannel channel = new ExceptionThrowingChannel (request , true ) {
495- @ Override
496- protected BytesStream newBytesOutput () {
497- return new RecyclerBytesStreamOutput (recycler );
498- }
499- };
489+ try {
490+ RestController .PERMIT_DOUBLE_RESPONSE = true ;
491+ int contentLength = Math .toIntExact (BREAKER_LIMIT .getBytes ());
492+ String content = randomAlphaOfLength ((int ) Math .round (contentLength / inFlightRequestsBreaker .getOverhead ()));
493+ // use a real recycler that tracks leaks and create some content bytes in the test handler to check for leaks
494+ final BytesRefRecycler recycler = new BytesRefRecycler (new MockPageCacheRecycler (Settings .EMPTY ));
495+ restController .registerHandler (
496+ new Route (GET , "/foo" ),
497+ (request , c , client ) -> new RestToXContentListener <>(c ).onResponse ((b , p ) -> b .startObject ().endObject ())
498+ );
499+ // we will produce an error in the rest handler and one more when sending the error response
500+ RestRequest request = testRestRequest ("/foo" , content , XContentType .JSON );
501+ ExceptionThrowingChannel channel = new ExceptionThrowingChannel (request , true ) {
502+ @ Override
503+ protected BytesStream newBytesOutput () {
504+ return new RecyclerBytesStreamOutput (recycler );
505+ }
506+ };
500507
501- restController .dispatchRequest (request , channel , client .threadPool ().getThreadContext ());
508+ restController .dispatchRequest (request , channel , client .threadPool ().getThreadContext ());
502509
503- assertEquals (0 , inFlightRequestsBreaker .getTrippedCount ());
504- assertEquals (0 , inFlightRequestsBreaker .getUsed ());
510+ assertEquals (0 , inFlightRequestsBreaker .getTrippedCount ());
511+ assertEquals (0 , inFlightRequestsBreaker .getUsed ());
512+ } finally {
513+ RestController .PERMIT_DOUBLE_RESPONSE = false ;
514+ }
505515 }
506516
507517 public void testDispatchRequestLimitsBytes () {
0 commit comments