How to refactor the code when the response should always be closed properly?

In one of our project, we need to use httpClient to get some data from backend services. We found the legacy code doesn't close the response properly.

The code is like:

HttpResponse response = httpClient.execute(request);
int statusCode = response.getStatusLine().getStatusCode();
if (statusCode == 200) {
    return EntityUtils.toString(response.getEntity());
} else {
    throw new InvalidResponseException();
}

When statusCode is 200, EntityUtils.toString will consume the content of the response and then close it properly. But in other cases, the response is not closed and we will have http connection leak (After some time, the httpClient pool is run up that we can't get new thread)

There are many code like this over the codebase, so I want to use the load design pattern to make it simple.

I defined a HttpClientWrapper, like:

class HttpClientWrapper {
    private HttpClient client;

    public HttpClientWrapper(HttpClient client) {
        this.client = client;
    }

    public <T> T execute(HttpRequestBase request, WithResponse<T> handler) {
        HttpResponse response = null;
        try {
            response = client.execute(request);
            return handler.withResponse(response);
        } catch (Exception e) {
            throw new RuntimeException(e);
        } finally {
            if (response != null) {
                EntityUtils.consumeQuietly(response.getEntity());
            }
        }
    }
}

interface WithResponse<T> {
    T withResponse(HttpResponse response) throws Exception;
}

I consumed the response in the wrapper in finally, so the response will always be closed properly. And I can change the existing code using it happily:

return new HttpClientWrapper(httpClient).execute(request, new WithResponse<String>() {
    String withResponse(HttpResponse response) throws Exception {
        int statusCode = response.getStatusLine().getStatusCode();
        if (statusCode == 200) {
            return EntityUtils.toString(response.getEntity());
        } else {
            throw new InvalidResponseException();
        }
    }
});

I don't need to worry the leak any more.

But suddenly I found such a block of code:

Stopwatch stopwatch = new Stopwatch();

try {
    stopwatch.start();
    HttpResponse response = httpClient.execute( request );
    stopwatch.stop();

    MDC.put( "backendTime", String.valueOf( stopwatch.elapsed( TimeUnit.MILLISECONDS ) ) );

    return EntityUtils.toString(response.getEntity());

} catch( IOException e ) {
    throw new RuntimeException( e );
}

It need to check how long the httpClient used to get the response! I can't use the HttpClientWrapper here, because I can't find a way to measure the part of the whole process per current design.

I have 2 options now:

1. 
   

   Don't use the HttpClientWrapper for this code, and we need to close the response manually. (But the way of getting response is is not consistent anymore)

   
   


2. 
   

   Modify the HttpClientWrapper, to make it complex and flex enough for this requirement. (But there is only one place needs it)

   
   

I like neither, is there any other better solution?