use of com.linkedin.r2.message.stream.entitystream.Observer in project rest.li by linkedin.
the class LoggingFilter method onStreamResponse.
@Override
public void onStreamResponse(StreamResponse res, RequestContext requestContext, Map<String, String> wireAttrs, NextFilter<StreamRequest, StreamResponse> nextFilter) {
EntityStream entityStream = res.getEntityStream();
entityStream.addObserver(new Observer() {
private long startTime;
private long bytesNum = 0;
@Override
public void onDataAvailable(ByteString data) {
if (bytesNum == 0) {
startTime = System.nanoTime();
}
bytesNum += data.length();
}
@Override
public void onDone() {
long stopTime = System.nanoTime();
_log.info("Status: success. Total bytes streamed: " + bytesNum + ". Total stream time: " + (stopTime - startTime) + " nano seconds.");
}
@Override
public void onError(Throwable e) {
long stopTime = System.nanoTime();
_log.error("Status: failed. Total bytes streamed: " + bytesNum + ". Total stream time before failure: " + (stopTime - startTime) + " nano seconds.");
}
});
}
use of com.linkedin.r2.message.stream.entitystream.Observer in project rest.li by linkedin.
the class HttpDispatcher method handleRequest.
/**
* handle a {@link com.linkedin.r2.message.stream.StreamRequest} using the given request context.
* @see com.linkedin.r2.transport.common.bridge.server.TransportDispatcher#handleStreamRequest
*
* @param req the request to be handled.
* @param context the request context.
* @param callback the callback to be invoked with the response or error.
*/
public void handleRequest(StreamRequest req, RequestContext context, final TransportCallback<StreamResponse> callback) {
markOnRequestTimings(context);
final Map<String, String> headers = new HashMap<>(req.getHeaders());
final Map<String, String> wireAttrs = WireAttributeHelper.removeWireAttributes(headers);
final BaseConnector connector = new BaseConnector();
try {
MessageType.Type msgType = MessageType.getMessageType(wireAttrs, MessageType.Type.REST);
switch(msgType) {
default:
case REST:
req.getEntityStream().setReader(connector);
StreamRequest newReq = req.builder().build(EntityStreams.newEntityStream(connector));
// decorate the call back so that if response is error or response finishes streaming,
// we cancel the request stream
TransportCallback<StreamResponse> decorateCallback = new TransportCallback<StreamResponse>() {
@Override
public void onResponse(TransportResponse<StreamResponse> response) {
// no need to check StreamException because that's handled by HttpBridge.httpToStreamCallback
if (response.hasError()) {
connector.cancel();
} else {
Observer observer = new Observer() {
@Override
public void onDataAvailable(ByteString data) {
// do nothing
}
@Override
public void onDone() {
connector.cancel();
}
@Override
public void onError(Throwable e) {
connector.cancel();
}
};
response.getResponse().getEntityStream().addObserver(observer);
}
callback.onResponse(response);
}
};
_dispatcher.handleStreamRequest(HttpBridge.toStreamRequest(newReq, headers), wireAttrs, context, HttpBridge.httpToStreamCallback(decorateCallback));
}
} catch (Exception e) {
connector.cancel();
callback.onResponse(TransportResponseImpl.<StreamResponse>error(e, Collections.<String, String>emptyMap()));
}
}
use of com.linkedin.r2.message.stream.entitystream.Observer in project rest.li by linkedin.
the class TestEntityStream method testCancelWhenNotWritePossible.
@Test
public void testCancelWhenNotWritePossible() throws Exception {
TestWriter writer = new TestWriter();
TestObserver observer = new TestObserver();
final ControlReader reader = new ControlReader();
EntityStream es = EntityStreams.newEntityStream(writer);
es.addObserver(observer);
es.setReader(reader);
reader.read(10);
ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
final CountDownLatch latch = new CountDownLatch(1);
scheduler.schedule(new Runnable() {
@Override
public void run() {
reader.cancel();
latch.countDown();
}
}, 50, TimeUnit.MILLISECONDS);
while (writer.remaining() > 0) {
writer.write();
}
Assert.assertTrue(latch.await(1000, TimeUnit.MILLISECONDS));
Assert.assertEquals(writer.abortedTimes(), 1);
Assert.assertEquals(observer.getChunkCount(), 10);
Assert.assertEquals(observer.doneTimes(), 0);
Assert.assertEquals(observer.errorTimes(), 1);
Assert.assertEquals(observer.getLastEvent(), "onError");
Assert.assertEquals(reader.getChunkCount(), 10);
Assert.assertEquals(reader.doneTimes(), 0);
Assert.assertEquals(reader.errorTimes(), 0);
scheduler.shutdown();
}
use of com.linkedin.r2.message.stream.entitystream.Observer in project rest.li by linkedin.
the class TestEntityStream method testWriterInitError.
@Test
public void testWriterInitError() throws Exception {
final CountDownLatch latch = new CountDownLatch(1);
final TestWriter writer = new TestWriter() {
@Override
public void onInit(WriteHandle wh) {
throw new RuntimeException();
}
};
TestObserver observer = new TestObserver();
final EntityStream es = EntityStreams.newEntityStream(writer);
es.addObserver(observer);
final ControlReader reader = new ControlReader();
ExecutorService executor = Executors.newFixedThreadPool(1);
executor.submit(new Runnable() {
@Override
public void run() {
es.setReader(reader);
reader.read(5);
latch.countDown();
}
});
Assert.assertTrue(latch.await(1000, TimeUnit.MILLISECONDS));
Assert.assertEquals(reader.errorTimes(), 1);
Assert.assertEquals(reader.getChunkCount(), 0);
Assert.assertEquals(observer.errorTimes(), 1);
Assert.assertEquals(reader.getChunkCount(), 0);
Assert.assertEquals(writer.abortedTimes(), 1);
Assert.assertEquals(writer.getWritePossibleCount(), 0);
}
use of com.linkedin.r2.message.stream.entitystream.Observer in project rest.li by linkedin.
the class TestEntityStream method testRaceBetweenErrorAndCancel.
private void testRaceBetweenErrorAndCancel(ExecutorService executor) throws Exception {
final CountDownLatch startLatch = new CountDownLatch(1);
final CountDownLatch finishLatch = new CountDownLatch(2);
final CountDownLatch prepareLatch = new CountDownLatch(2);
final TestWriter writer = new TestWriter();
TestObserver observer = new TestObserver();
final ControlReader reader = new ControlReader();
EntityStream es = EntityStreams.newEntityStream(writer);
es.addObserver(observer);
es.setReader(reader);
reader.read(1000);
executor.submit(new Callable<Object>() {
@Override
public Object call() throws Exception {
while (writer.remaining() > 100) {
writer.write();
}
prepareLatch.countDown();
startLatch.await();
writer.error(new RuntimeException("writer has problem"));
finishLatch.countDown();
return null;
}
});
executor.submit(new Callable<Object>() {
@Override
public Object call() throws Exception {
prepareLatch.countDown();
startLatch.await();
reader.cancel();
finishLatch.countDown();
return null;
}
});
prepareLatch.await();
startLatch.countDown();
Assert.assertTrue(finishLatch.await(1000, TimeUnit.MILLISECONDS));
// if error wins the race
if (reader.errorTimes() > 0) {
Assert.assertEquals(reader.doneTimes(), 0);
Assert.assertEquals(observer.doneTimes(), 0);
Assert.assertEquals(observer.errorTimes(), 1);
Assert.assertEquals(writer.abortedTimes(), 0);
} else // if cancel wins the race
{
Assert.assertEquals(observer.doneTimes(), 0);
Assert.assertEquals(observer.errorTimes(), 1);
Assert.assertEquals(writer.abortedTimes(), 1);
Assert.assertEquals(reader.doneTimes(), 0);
Assert.assertEquals(reader.errorTimes(), 0);
}
}
Aggregations