/*
* Copyright 2016-present Facebook, Inc.
*
* Licensed under the Apache License, Version 2.0 (the "License"); you may
* not use this file except in compliance with the License. You may obtain
* a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* License for the specific language governing permissions and limitations
* under the License.
*/
package com.facebook.buck.slb;
import com.google.common.annotations.VisibleForTesting;
import java.io.IOException;
import java.io.InputStream;
import java.net.URI;
import okhttp3.Call;
import okhttp3.Response;
public class LoadBalancedHttpResponse extends OkHttpResponseWrapper {
private final HttpLoadBalancer loadBalancer;
private final URI server;
private boolean hasConnectionResultBeenReported;
private static final boolean FIX_HTTP_BOTTLENECK =
"true".equals(System.getProperty("buck.fix_http_bottleneck"));
public static LoadBalancedHttpResponse createLoadBalancedResponse(
URI server, HttpLoadBalancer loadBalancer, Call call) throws IOException {
try {
return new LoadBalancedHttpResponse(server, loadBalancer, call.execute());
} catch (IOException e) {
if (FIX_HTTP_BOTTLENECK) {
loadBalancer.reportRequestException(server);
}
throw e;
}
}
@VisibleForTesting
LoadBalancedHttpResponse(URI server, HttpLoadBalancer loadBalancer, Response response) {
super(response);
this.loadBalancer = loadBalancer;
this.server = server;
this.hasConnectionResultBeenReported = false;
}
@Override
public long contentLength() throws IOException {
try {
return super.contentLength();
} catch (IOException exception) {
reportConnectionResultIfFirst(false);
throw exception;
}
}
@Override
public InputStream getBody() {
return new LoadBalancedInputStream(getResponse().body().byteStream());
}
@Override
public void close() throws IOException {
super.close();
// We can only be sure a connection was successful after all data was read and the
// connection successfully closed.
reportConnectionResultIfFirst(true);
}
private void reportConnectionResultIfFirst(boolean successful) {
if (hasConnectionResultBeenReported) {
return;
}
hasConnectionResultBeenReported = true;
if (successful) {
loadBalancer.reportRequestSuccess(server);
} else {
loadBalancer.reportRequestException(server);
}
}
private class LoadBalancedInputStream extends InputStream {
private final InputStream rawStream;
public LoadBalancedInputStream(InputStream stream) {
this.rawStream = stream;
}
@Override
public void close() throws IOException {
rawStream.close();
}
@Override
public int read() throws IOException {
try {
return rawStream.read();
} catch (IOException e) {
reportConnectionResultIfFirst(false);
throw e;
}
}
@Override
public int read(byte[] b) throws IOException {
try {
return rawStream.read(b);
} catch (IOException e) {
reportConnectionResultIfFirst(false);
throw e;
}
}
@Override
public int read(byte[] b, int off, int len) throws IOException {
try {
return rawStream.read(b, off, len);
} catch (IOException e) {
reportConnectionResultIfFirst(false);
throw e;
}
}
@Override
public long skip(long n) throws IOException {
try {
return rawStream.skip(n);
} catch (IOException e) {
reportConnectionResultIfFirst(false);
throw e;
}
}
@Override
public int available() throws IOException {
try {
return rawStream.available();
} catch (IOException e) {
reportConnectionResultIfFirst(false);
throw e;
}
}
@Override
public synchronized void mark(int readlimit) {
rawStream.mark(readlimit);
}
@Override
public synchronized void reset() throws IOException {
try {
rawStream.reset();
} catch (IOException e) {
reportConnectionResultIfFirst(false);
throw e;
}
}
@Override
public boolean markSupported() {
return rawStream.markSupported();
}
}
}