Spring Cloud - Getting Retry Working In RestTemplate?

Question

I have been migrating an existing application over to Spring Cloud\'s service discovery, Ribbon load balancing, and circuit breakers. The application already makes extensive

Answer 1

To answer my own question:

Before I get into the details, a cautionary tale:

Eureka's self preservation mode sent me down a rabbit hole while testing the fail-over on my local machine. I recommend turning self preservation mode off while doing your testing. Because I was dropping nodes at a regular rate and then restarting (with a different instance ID using a random value), I tripped Eureka's self preservation mode. I ended up with many instances in Eureka that pointed to the same machine, same port. The fail-over was actually working but the next node that was chosen happened to be another dead instance. Very confusing at first!

I was able to get fail-over working with a modified version of RibbonClientHttpRequestFactory. Because RibbonAutoConfiguration creates a load balanced RestTemplate with this factory, rather then injecting this rest template, I create a new one with my modified version of the request factory:

protected RestTemplate restTemplate;

@Autowired
public void customizeRestTemplate(SpringClientFactory springClientFactory, LoadBalancerClient loadBalancerClient) {
    restTemplate = new RestTemplate();

    // Use a modified version of the http request factory that leverages the load balacing in netflix's RestClient.
    RibbonRetryHttpRequestFactory lFactory = new RibbonRetryHttpRequestFactory(springClientFactory, loadBalancerClient);
    restTemplate.setRequestFactory(lFactory);
}

The modified Request Factory is just a copy of RibbonClientHttpRequestFactory with two minor changes:

1) In createRequest, I removed the code that was selecting a server from the load balancer because the RestClient will do that for us. 2) In the inner class, RibbonHttpRequest, I changed executeInternal to call "executeWithLoadBalancer".

The full class:

@SuppressWarnings("deprecation")
public class RibbonRetryHttpRequestFactory implements ClientHttpRequestFactory {

    private final SpringClientFactory clientFactory;
    private LoadBalancerClient loadBalancer;

    public RibbonRetryHttpRequestFactory(SpringClientFactory clientFactory, LoadBalancerClient loadBalancer) {
        this.clientFactory = clientFactory;
        this.loadBalancer = loadBalancer;
    }

    @Override
    public ClientHttpRequest createRequest(URI originalUri, HttpMethod httpMethod) throws IOException {
        String serviceId = originalUri.getHost();
        IClientConfig clientConfig = clientFactory.getClientConfig(serviceId);

        RestClient client = clientFactory.getClient(serviceId, RestClient.class);
        HttpRequest.Verb verb = HttpRequest.Verb.valueOf(httpMethod.name());
        return new RibbonHttpRequest(originalUri, verb, client, clientConfig);
    }

    public class RibbonHttpRequest extends AbstractClientHttpRequest {

        private HttpRequest.Builder builder;
        private URI uri;
        private HttpRequest.Verb verb;
        private RestClient client;
        private IClientConfig config;
        private ByteArrayOutputStream outputStream = null;

        public RibbonHttpRequest(URI uri, HttpRequest.Verb verb, RestClient client, IClientConfig config) {
            this.uri = uri;
            this.verb = verb;
            this.client = client;
            this.config = config;
            this.builder = HttpRequest.newBuilder().uri(uri).verb(verb);
        }

        @Override
        public HttpMethod getMethod() {
            return HttpMethod.valueOf(verb.name());
        }

        @Override
        public URI getURI() {
            return uri;
        }

        @Override
        protected OutputStream getBodyInternal(HttpHeaders headers) throws IOException {
            if (outputStream == null) {
                outputStream = new ByteArrayOutputStream();
            }
            return outputStream;
        }

        @Override
        protected ClientHttpResponse executeInternal(HttpHeaders headers) throws IOException {
            try {
                addHeaders(headers);
                if (outputStream != null) {
                    outputStream.close();
                    builder.entity(outputStream.toByteArray());
                }
                HttpRequest request = builder.build();
                HttpResponse response = client.executeWithLoadBalancer(request, config);
                return new RibbonHttpResponse(response);
            }
            catch (Exception e) {
                throw new IOException(e);
            }

            //TODO: fix stats, now that execute is not called
            // use execute here so stats are collected
            /*
            return loadBalancer.execute(this.config.getClientName(), new LoadBalancerRequest<ClientHttpResponse>() {
                @Override
                public ClientHttpResponse apply(ServiceInstance instance) throws Exception {}
            });
            */
        }

        private void addHeaders(HttpHeaders headers) {
            for (String name : headers.keySet()) {
                // apache http RequestContent pukes if there is a body and
                // the dynamic headers are already present
                if (!isDynamic(name) || outputStream == null) {
                    List<String> values = headers.get(name);
                    for (String value : values) {
                        builder.header(name, value);
                    }
                }
            }
        }

        private boolean isDynamic(String name) {
            return name.equals("Content-Length") || name.equals("Transfer-Encoding");
        }
    }

    public class RibbonHttpResponse extends AbstractClientHttpResponse {

        private HttpResponse response;
        private HttpHeaders httpHeaders;

        public RibbonHttpResponse(HttpResponse response) {
            this.response = response;
            this.httpHeaders = new HttpHeaders();
            List<Map.Entry<String, String>> headers = response.getHttpHeaders().getAllHeaders();
            for (Map.Entry<String, String> header : headers) {
                this.httpHeaders.add(header.getKey(), header.getValue());
            }
        }

        @Override
        public InputStream getBody() throws IOException {
            return response.getInputStream();
        }

        @Override
        public HttpHeaders getHeaders() {
            return this.httpHeaders;
        }

        @Override
        public int getRawStatusCode() throws IOException {
            return response.getStatus();
        }

        @Override
        public String getStatusText() throws IOException {
            return HttpStatus.valueOf(response.getStatus()).name();
        }

        @Override
        public void close() {
            response.close();
        }
    }
}