io.split.engine.sse.client.SSEClient Maven / Gradle / Ivy
The newest version!
package io.split.engine.sse.client;
import com.google.common.base.Strings;
import io.split.client.RequestDecorator;
import io.split.client.utils.ApacheRequestDecorator;
import io.split.telemetry.domain.StreamingEvent;
import io.split.telemetry.domain.enums.StreamEventsEnum;
import io.split.telemetry.storage.TelemetryRuntimeProducer;
import org.apache.hc.client5.http.classic.methods.HttpGet;
import org.apache.hc.client5.http.impl.classic.CloseableHttpClient;
import org.apache.hc.client5.http.impl.classic.CloseableHttpResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.BufferedReader;
import java.io.EOFException;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.net.SocketException;
import java.net.URI;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.Function;
import static com.google.common.base.Preconditions.checkNotNull;
import static io.split.client.utils.SplitExecutorFactory.buildExecutorService;
public class SSEClient {
public enum StatusMessage {
CONNECTED,
RETRYABLE_ERROR,
NONRETRYABLE_ERROR,
INITIALIZATION_IN_PROGRESS,
FORCED_STOP,
FIRST_EVENT
}
private enum ConnectionState {
OPEN,
CLOSED
}
private final static String SOCKET_CLOSED_MESSAGE = "Socket closed";
private final static String KEEP_ALIVE_PAYLOAD = ":keepalive\n";
private final static long CONNECT_TIMEOUT = 30000;
private static final Lock lock = new ReentrantLock();
private static final Logger _log = LoggerFactory.getLogger(SSEClient.class);
private final ExecutorService _connectionExecutor;
private final CloseableHttpClient _client;
private final Function _eventCallback;
private final Function _statusCallback;
private final AtomicReference _state = new AtomicReference<>(ConnectionState.CLOSED);
private final AtomicReference _ongoingResponse = new AtomicReference<>();
private final AtomicReference _ongoingRequest = new AtomicReference<>();
private AtomicBoolean _forcedStop;
private final RequestDecorator _requestDecorator;
private final TelemetryRuntimeProducer _telemetryRuntimeProducer;
public SSEClient(Function eventCallback,
Function statusCallback,
CloseableHttpClient client,
TelemetryRuntimeProducer telemetryRuntimeProducer,
ThreadFactory threadFactory,
RequestDecorator requestDecorator) {
_eventCallback = eventCallback;
_statusCallback = statusCallback;
_client = client;
_forcedStop = new AtomicBoolean();
_telemetryRuntimeProducer = checkNotNull(telemetryRuntimeProducer);
_connectionExecutor = buildExecutorService(threadFactory, "SPLIT-SSEConnection-%d");
_requestDecorator = requestDecorator;
}
public boolean open(URI uri) {
try {
lock.lock();
if (isOpen()) {
_log.info("SSEClient already open.");
return false;
}
_statusCallback.apply(StatusMessage.INITIALIZATION_IN_PROGRESS);
CountDownLatch signal = new CountDownLatch(1);
_connectionExecutor.submit(() -> connectAndLoop(uri, signal));
try {
if (!signal.await(CONNECT_TIMEOUT, TimeUnit.SECONDS)) {
return false;
}
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
if (e.getMessage() == null) {
_log.info("The thread was interrupted while opening SSEClient");
return false;
}
_log.info(e.getMessage());
return false;
}
return isOpen();
} finally {
lock.unlock();
}
}
public boolean isOpen() {
return (ConnectionState.OPEN.equals(_state.get()));
}
public void close() {
try {
lock.lock();
_forcedStop.set(true);
if (_state.compareAndSet(ConnectionState.OPEN, ConnectionState.CLOSED)) {
if (_ongoingResponse.get() != null) {
try {
_ongoingRequest.get().abort();
_ongoingResponse.get().close();
} catch (IOException e) {
_log.debug(String.format("SSEClient close forced: %s", e.getMessage()));
}
}
}
} finally {
lock.unlock();
}
}
private void connectAndLoop(URI uri, CountDownLatch signal) {
checkNotNull(uri);
checkNotNull(signal);
try {
if (!establishConnection(uri, signal)) {
_statusCallback.apply(StatusMessage.RETRYABLE_ERROR);
return;
}
final InputStream stream = _ongoingResponse.get().getEntity().getContent();
final BufferedReader reader = new BufferedReader(new InputStreamReader(stream));
while (isOpen() && !Thread.currentThread().isInterrupted()) {
try {
handleMessage(readMessageAsString(reader));
} catch (SocketException exc) {
_log.debug(exc.getMessage());
if (SOCKET_CLOSED_MESSAGE.equals(exc.getMessage())) { // Connection closed by us
_statusCallback.apply(StatusMessage.FORCED_STOP);
_telemetryRuntimeProducer.recordStreamingEvents(
new StreamingEvent(StreamEventsEnum.SSE_CONNECTION_ERROR.getType(),
StreamEventsEnum.SseConnectionErrorValues.REQUESTED_CONNECTION_ERROR.getValue(),
System.currentTimeMillis()));
return;
}
// Connection closed by server
_statusCallback.apply(StatusMessage.RETRYABLE_ERROR);
_telemetryRuntimeProducer
.recordStreamingEvents(new StreamingEvent(StreamEventsEnum.SSE_CONNECTION_ERROR.getType(),
StreamEventsEnum.SseConnectionErrorValues.NON_REQUESTED_CONNECTION_ERROR.getValue(),
System.currentTimeMillis()));
return;
} catch (IOException exc) { // Other type of connection error
if (!_forcedStop.get()) {
_log.debug(String.format("SSE connection ended abruptly: %s. Retying", exc.getMessage()));
_telemetryRuntimeProducer.recordStreamingEvents(
new StreamingEvent(StreamEventsEnum.SSE_CONNECTION_ERROR.getType(),
StreamEventsEnum.SseConnectionErrorValues.REQUESTED_CONNECTION_ERROR.getValue(),
System.currentTimeMillis()));
_statusCallback.apply(StatusMessage.RETRYABLE_ERROR);
return;
}
_telemetryRuntimeProducer
.recordStreamingEvents(new StreamingEvent(StreamEventsEnum.SSE_CONNECTION_ERROR.getType(),
StreamEventsEnum.SseConnectionErrorValues.NON_REQUESTED_CONNECTION_ERROR.getValue(),
System.currentTimeMillis()));
}
}
} catch (Exception e) { // Any other error non related to the connection disables streaming altogether
_telemetryRuntimeProducer
.recordStreamingEvents(new StreamingEvent(StreamEventsEnum.SSE_CONNECTION_ERROR.getType(),
StreamEventsEnum.SseConnectionErrorValues.NON_REQUESTED_CONNECTION_ERROR.getValue(),
System.currentTimeMillis()));
_log.warn(e.getMessage(), e);
_statusCallback.apply(StatusMessage.NONRETRYABLE_ERROR);
} finally {
try {
_ongoingResponse.get().close();
} catch (IOException e) {
_log.debug(e.getMessage());
}
_state.set(ConnectionState.CLOSED);
_log.debug("SSEClient finished.");
_forcedStop.set(false);
}
}
private boolean establishConnection(URI uri, CountDownLatch signal) {
HttpGet request = new HttpGet(uri);
request = (HttpGet) ApacheRequestDecorator.decorate(request, _requestDecorator);
_ongoingRequest.set(request);
try {
_ongoingResponse.set(_client.execute(_ongoingRequest.get()));
if (_ongoingResponse.get().getCode() != 200) {
_log.error(String.format("Establishing connection, code error: %s. The url is %s",
_ongoingResponse.get().getCode(), uri.toURL()));
return false;
}
_state.set(ConnectionState.OPEN);
_statusCallback.apply(StatusMessage.CONNECTED);
} catch (IOException exc) {
_log.error(String.format("Error establishConnection to %s", uri), exc);
return false;
} finally {
signal.countDown();
}
return true;
}
static private String readMessageAsString(BufferedReader reader) throws IOException {
StringBuilder lines = new StringBuilder();
while (true) {
String line = reader.readLine();
if (null == line) { // EOF (Remote host closed the connection gracefully and the stream ended)
throw new EOFException("connection closed by remote host");
}
if (line.isEmpty()) { // An empty line is a message separator
return lines.toString();
}
lines.append(line).append("\n");
}
}
private void handleMessage(String message) {
if (Strings.isNullOrEmpty(message) || KEEP_ALIVE_PAYLOAD.equals(message)) {
_log.debug("Keep Alive event");
return;
}
RawEvent e = RawEvent.fromString(message);
_eventCallback.apply(e);
}
}