聊聊spring cloud的AsyncLoadBalancerAutoConfiguration

/**
 * Auto configuration for Ribbon (client side load balancing).
 *
 * @author Rob Worsnop
 */
@Configuration
@ConditionalOnBean(LoadBalancerClient.class)
@ConditionalOnClass(AsyncRestTemplate.class)
public class AsyncLoadBalancerAutoConfiguration {

    @Configuration
    static class AsyncRestTemplateCustomizerConfig {
        @LoadBalanced
        @Autowired(required = false)
        private List<AsyncRestTemplate> restTemplates = Collections.emptyList();

        @Bean
        public SmartInitializingSingleton loadBalancedAsyncRestTemplateInitializer(
                final List<AsyncRestTemplateCustomizer> customizers) {
            return new SmartInitializingSingleton() {
                @Override
                public void afterSingletonsInstantiated() {
                    for (AsyncRestTemplate restTemplate : AsyncRestTemplateCustomizerConfig.this.restTemplates) {
                        for (AsyncRestTemplateCustomizer customizer : customizers) {
                            customizer.customize(restTemplate);
                        }
                    }
                }
            };
        }
    }

    @Configuration
    static class LoadBalancerInterceptorConfig {
        @Bean
        public AsyncLoadBalancerInterceptor asyncLoadBalancerInterceptor(LoadBalancerClient loadBalancerClient) {
            return new AsyncLoadBalancerInterceptor(loadBalancerClient);
        }

        @Bean
        public AsyncRestTemplateCustomizer asyncRestTemplateCustomizer(
                final AsyncLoadBalancerInterceptor loadBalancerInterceptor) {
            return new AsyncRestTemplateCustomizer() {
                @Override
                public void customize(AsyncRestTemplate restTemplate) {
                    List<AsyncClientHttpRequestInterceptor> list = new ArrayList<>(
                            restTemplate.getInterceptors());
                    list.add(loadBalancerInterceptor);
                    restTemplate.setInterceptors(list);
                }
            };
        }
    }
}

public interface AsyncRestTemplateCustomizer {
    void customize(AsyncRestTemplate restTemplate);
}

public class AsyncLoadBalancerInterceptor implements AsyncClientHttpRequestInterceptor {

    private LoadBalancerClient loadBalancer;

    public AsyncLoadBalancerInterceptor(LoadBalancerClient loadBalancer) {
        this.loadBalancer = loadBalancer;
    }

    @Override
    public ListenableFuture<ClientHttpResponse> intercept(final HttpRequest request, final byte[] body,
            final AsyncClientHttpRequestExecution execution) throws IOException {
        final URI originalUri = request.getURI();
        String serviceName = originalUri.getHost();
        return this.loadBalancer.execute(serviceName,
                new LoadBalancerRequest<ListenableFuture<ClientHttpResponse>>() {
                    @Override
                    public ListenableFuture<ClientHttpResponse> apply(final ServiceInstance instance)
                            throws Exception {
                        HttpRequest serviceRequest = new ServiceRequestWrapper(request,
                                instance, loadBalancer);
                        return execution.executeAsync(serviceRequest, body);
                    }

                });
    }
}

/**
 * Abstract base for {@link AsyncClientHttpRequest} that makes sure that headers and body
 * are not written multiple times.
 *
 * @author Arjen Poutsma
 * @since 4.0
 * @deprecated as of Spring 5.0, in favor of {@link org.springframework.http.client.reactive.AbstractClientHttpRequest}
 */
@Deprecated
abstract class AbstractAsyncClientHttpRequest implements AsyncClientHttpRequest {

    private final HttpHeaders headers = new HttpHeaders();

    private boolean executed = false;

    @Override
    public final HttpHeaders getHeaders() {
        return (this.executed ? HttpHeaders.readOnlyHttpHeaders(this.headers) : this.headers);
    }

    @Override
    public final OutputStream getBody() throws IOException {
        assertNotExecuted();
        return getBodyInternal(this.headers);
    }

    @Override
    public ListenableFuture<ClientHttpResponse> executeAsync() throws IOException {
        assertNotExecuted();
        ListenableFuture<ClientHttpResponse> result = executeInternal(this.headers);
        this.executed = true;
        return result;
    }

    /**
     * Asserts that this request has not been {@linkplain #executeAsync() executed} yet.
     * @throws IllegalStateException if this request has been executed
     */
    protected void assertNotExecuted() {
        Assert.state(!this.executed, "ClientHttpRequest already executed");
    }

    /**
     * Abstract template method that returns the body.
     * @param headers the HTTP headers
     * @return the body output stream
     */
    protected abstract OutputStream getBodyInternal(HttpHeaders headers) throws IOException;

    /**
     * Abstract template method that writes the given headers and content to the HTTP request.
     * @param headers the HTTP headers
     * @return the response object for the executed request
     */
    protected abstract ListenableFuture<ClientHttpResponse> executeInternal(HttpHeaders headers)
            throws IOException;

}

/**
 * {@link org.springframework.http.client.ClientHttpRequest} implementation
 * that uses standard Java facilities to execute streaming requests. Created
 * via the {@link org.springframework.http.client.SimpleClientHttpRequestFactory}.
 *
 * @author Arjen Poutsma
 * @since 3.0
 * @see org.springframework.http.client.SimpleClientHttpRequestFactory#createRequest
 * @see org.springframework.http.client.support.AsyncHttpAccessor
 * @see org.springframework.web.client.AsyncRestTemplate
 * @deprecated as of Spring 5.0, with no direct replacement
 */
@Deprecated
final class SimpleStreamingAsyncClientHttpRequest extends AbstractAsyncClientHttpRequest {

    private final HttpURLConnection connection;

    private final int chunkSize;

    @Nullable
    private OutputStream body;

    private final boolean outputStreaming;

    private final AsyncListenableTaskExecutor taskExecutor;

    SimpleStreamingAsyncClientHttpRequest(HttpURLConnection connection, int chunkSize,
            boolean outputStreaming, AsyncListenableTaskExecutor taskExecutor) {

        this.connection = connection;
        this.chunkSize = chunkSize;
        this.outputStreaming = outputStreaming;
        this.taskExecutor = taskExecutor;
    }

    @Override
    public String getMethodValue() {
        return this.connection.getRequestMethod();
    }

    @Override
    public URI getURI() {
        try {
            return this.connection.getURL().toURI();
        }
        catch (URISyntaxException ex) {
            throw new IllegalStateException(
                    "Could not get HttpURLConnection URI: " + ex.getMessage(), ex);
        }
    }

    @Override
    protected OutputStream getBodyInternal(HttpHeaders headers) throws IOException {
        if (this.body == null) {
            if (this.outputStreaming) {
                long contentLength = headers.getContentLength();
                if (contentLength >= 0) {
                    this.connection.setFixedLengthStreamingMode(contentLength);
                }
                else {
                    this.connection.setChunkedStreamingMode(this.chunkSize);
                }
            }
            SimpleBufferingClientHttpRequest.addHeaders(this.connection, headers);
            this.connection.connect();
            this.body = this.connection.getOutputStream();
        }
        return StreamUtils.nonClosing(this.body);
    }

    @Override
    protected ListenableFuture<ClientHttpResponse> executeInternal(final HttpHeaders headers) throws IOException {
        return this.taskExecutor.submitListenable(new Callable<ClientHttpResponse>() {
            @Override
            public ClientHttpResponse call() throws Exception {
                try {
                    if (body != null) {
                        body.close();
                    }
                    else {
                        SimpleBufferingClientHttpRequest.addHeaders(connection, headers);
                        connection.connect();
                        // Immediately trigger the request in a no-output scenario as well
                        connection.getResponseCode();
                    }
                }
                catch (IOException ex) {
                    // ignore
                }
                return new SimpleClientHttpResponse(connection);
            }
        });

    }

}

/**
 * {@link ClientHttpRequest} implementation based on Netty 4.
 *
 * <p>Created via the {@link Netty4ClientHttpRequestFactory}.
 *
 * @author Arjen Poutsma
 * @author Rossen Stoyanchev
 * @author Brian Clozel
 * @since 4.1.2
 * @deprecated as of Spring 5.0, in favor of
 * {@link org.springframework.http.client.reactive.ReactorClientHttpConnector}
 */
@Deprecated
class Netty4ClientHttpRequest extends AbstractAsyncClientHttpRequest implements ClientHttpRequest {

    private final Bootstrap bootstrap;

    private final URI uri;

    private final HttpMethod method;

    private final ByteBufOutputStream body;

    public Netty4ClientHttpRequest(Bootstrap bootstrap, URI uri, HttpMethod method) {
        this.bootstrap = bootstrap;
        this.uri = uri;
        this.method = method;
        this.body = new ByteBufOutputStream(Unpooled.buffer(1024));
    }

    @Override
    public HttpMethod getMethod() {
        return this.method;
    }

    @Override
    public String getMethodValue() {
        return this.method.name();
    }

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

    @Override
    public ClientHttpResponse execute() throws IOException {
        try {
            return executeAsync().get();
        }
        catch (InterruptedException ex) {
            Thread.currentThread().interrupt();
            throw new IOException("Interrupted during request execution", ex);
        }
        catch (ExecutionException ex) {
            if (ex.getCause() instanceof IOException) {
                throw (IOException) ex.getCause();
            }
            else {
                throw new IOException(ex.getMessage(), ex.getCause());
            }
        }
    }

    @Override
    protected OutputStream getBodyInternal(HttpHeaders headers) throws IOException {
        return this.body;
    }

    @Override
    protected ListenableFuture<ClientHttpResponse> executeInternal(final HttpHeaders headers) throws IOException {
        final SettableListenableFuture<ClientHttpResponse> responseFuture = new SettableListenableFuture<>();

        ChannelFutureListener connectionListener = future -> {
            if (future.isSuccess()) {
                Channel channel = future.channel();
                channel.pipeline().addLast(new RequestExecuteHandler(responseFuture));
                FullHttpRequest nettyRequest = createFullHttpRequest(headers);
                channel.writeAndFlush(nettyRequest);
            }
            else {
                responseFuture.setException(future.cause());
            }
        };

        this.bootstrap.connect(this.uri.getHost(), getPort(this.uri)).addListener(connectionListener);
        return responseFuture;
    }

    private FullHttpRequest createFullHttpRequest(HttpHeaders headers) {
        io.netty.handler.codec.http.HttpMethod nettyMethod =
                io.netty.handler.codec.http.HttpMethod.valueOf(this.method.name());

        String authority = this.uri.getRawAuthority();
        String path = this.uri.toString().substring(this.uri.toString().indexOf(authority) + authority.length());
        FullHttpRequest nettyRequest = new DefaultFullHttpRequest(
                HttpVersion.HTTP_1_1, nettyMethod, path, this.body.buffer());

        nettyRequest.headers().set(HttpHeaders.HOST, this.uri.getHost() + ":" + getPort(uri));
        nettyRequest.headers().set(HttpHeaders.CONNECTION, "close");
        headers.forEach((headerName, headerValues) -> nettyRequest.headers().add(headerName, headerValues));
        if (!nettyRequest.headers().contains(HttpHeaders.CONTENT_LENGTH) && this.body.buffer().readableBytes() > 0) {
            nettyRequest.headers().set(HttpHeaders.CONTENT_LENGTH, this.body.buffer().readableBytes());
        }

        return nettyRequest;
    }

    private static int getPort(URI uri) {
        int port = uri.getPort();
        if (port == -1) {
            if ("http".equalsIgnoreCase(uri.getScheme())) {
                port = 80;
            }
            else if ("https".equalsIgnoreCase(uri.getScheme())) {
                port = 443;
            }
        }
        return port;
    }

    /**
     * A SimpleChannelInboundHandler to update the given SettableListenableFuture.
     */
    private static class RequestExecuteHandler extends SimpleChannelInboundHandler<FullHttpResponse> {

        private final SettableListenableFuture<ClientHttpResponse> responseFuture;

        public RequestExecuteHandler(SettableListenableFuture<ClientHttpResponse> responseFuture) {
            this.responseFuture = responseFuture;
        }

        @Override
        protected void channelRead0(ChannelHandlerContext context, FullHttpResponse response) throws Exception {
            this.responseFuture.set(new Netty4ClientHttpResponse(context, response));
        }

        @Override
        public void exceptionCaught(ChannelHandlerContext context, Throwable cause) throws Exception {
            this.responseFuture.setException(cause);
        }
    }

}