@@ -498,40 +498,50 @@ public void testDispatchRequestAddsAndFreesBytesOnError() {
498498    }
499499
500500    public  void  testDispatchRequestAddsAndFreesBytesOnlyOnceOnError () {
501-         int  contentLength  = BREAKER_LIMIT .bytesAsInt ();
502-         String  content  = randomAlphaOfLength ((int ) Math .round (contentLength  / inFlightRequestsBreaker .getOverhead ()));
503-         // we will produce an error in the rest handler and one more when sending the error response 
504-         RestRequest  request  = testRestRequest ("/error" , content , XContentType .JSON );
505-         ExceptionThrowingChannel  channel  = new  ExceptionThrowingChannel (request , randomBoolean ());
506- 
507-         restController .dispatchRequest (request , channel , client .threadPool ().getThreadContext ());
508- 
509-         assertEquals (0 , inFlightRequestsBreaker .getTrippedCount ());
510-         assertEquals (0 , inFlightRequestsBreaker .getUsed ());
501+         try  {
502+             RestController .PERMIT_DOUBLE_RESPONSE  = true ;
503+             int  contentLength  = BREAKER_LIMIT .bytesAsInt ();
504+             String  content  = randomAlphaOfLength ((int ) Math .round (contentLength  / inFlightRequestsBreaker .getOverhead ()));
505+             // we will produce an error in the rest handler and one more when sending the error response 
506+             RestRequest  request  = testRestRequest ("/error" , content , XContentType .JSON );
507+             ExceptionThrowingChannel  channel  = new  ExceptionThrowingChannel (request , randomBoolean ());
508+ 
509+             restController .dispatchRequest (request , channel , client .threadPool ().getThreadContext ());
510+ 
511+             assertEquals (0 , inFlightRequestsBreaker .getTrippedCount ());
512+             assertEquals (0 , inFlightRequestsBreaker .getUsed ());
513+         } finally  {
514+             RestController .PERMIT_DOUBLE_RESPONSE  = false ;
515+         }
511516    }
512517
513518    public  void  testDispatchRequestAddsAndFreesBytesOnlyOnceOnErrorDuringSend () {
514-         int  contentLength  = Math .toIntExact (BREAKER_LIMIT .getBytes ());
515-         String  content  = randomAlphaOfLength ((int ) Math .round (contentLength  / inFlightRequestsBreaker .getOverhead ()));
516-         // use a real recycler that tracks leaks and create some content bytes in the test handler to check for leaks 
517-         final  BytesRefRecycler  recycler  = new  BytesRefRecycler (new  MockPageCacheRecycler (Settings .EMPTY ));
518-         restController .registerHandler (
519-             new  Route (GET , "/foo" ),
520-             (request , c , client ) -> new  RestToXContentListener <>(c ).onResponse ((b , p ) -> b .startObject ().endObject ())
521-         );
522-         // we will produce an error in the rest handler and one more when sending the error response 
523-         RestRequest  request  = testRestRequest ("/foo" , content , XContentType .JSON );
524-         ExceptionThrowingChannel  channel  = new  ExceptionThrowingChannel (request , randomBoolean ()) {
525-             @ Override 
526-             protected  BytesStream  newBytesOutput () {
527-                 return  new  RecyclerBytesStreamOutput (recycler );
528-             }
529-         };
519+         try  {
520+             RestController .PERMIT_DOUBLE_RESPONSE  = true ;
521+             int  contentLength  = Math .toIntExact (BREAKER_LIMIT .getBytes ());
522+             String  content  = randomAlphaOfLength ((int ) Math .round (contentLength  / inFlightRequestsBreaker .getOverhead ()));
523+             // use a real recycler that tracks leaks and create some content bytes in the test handler to check for leaks 
524+             final  BytesRefRecycler  recycler  = new  BytesRefRecycler (new  MockPageCacheRecycler (Settings .EMPTY ));
525+             restController .registerHandler (
526+                 new  Route (GET , "/foo" ),
527+                 (request , c , client ) -> new  RestToXContentListener <>(c ).onResponse ((b , p ) -> b .startObject ().endObject ())
528+             );
529+             // we will produce an error in the rest handler and one more when sending the error response 
530+             RestRequest  request  = testRestRequest ("/foo" , content , XContentType .JSON );
531+             ExceptionThrowingChannel  channel  = new  ExceptionThrowingChannel (request , randomBoolean ()) {
532+                 @ Override 
533+                 protected  BytesStream  newBytesOutput () {
534+                     return  new  RecyclerBytesStreamOutput (recycler );
535+                 }
536+             };
530537
531-         restController .dispatchRequest (request , channel , client .threadPool ().getThreadContext ());
538+              restController .dispatchRequest (request , channel , client .threadPool ().getThreadContext ());
532539
533-         assertEquals (0 , inFlightRequestsBreaker .getTrippedCount ());
534-         assertEquals (0 , inFlightRequestsBreaker .getUsed ());
540+             assertEquals (0 , inFlightRequestsBreaker .getTrippedCount ());
541+             assertEquals (0 , inFlightRequestsBreaker .getUsed ());
542+         } finally  {
543+             RestController .PERMIT_DOUBLE_RESPONSE  = false ;
544+         }
535545    }
536546
537547    public  void  testDispatchRequestLimitsBytes () {
0 commit comments