Search in sources :

Example 11 with Observer

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.");
        }
    });
}
Also used : EntityStream(com.linkedin.r2.message.stream.entitystream.EntityStream) ByteString(com.linkedin.data.ByteString) Observer(com.linkedin.r2.message.stream.entitystream.Observer)

Example 12 with Observer

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()));
    }
}
Also used : TransportCallback(com.linkedin.r2.transport.common.bridge.common.TransportCallback) HashMap(java.util.HashMap) BaseConnector(com.linkedin.r2.message.stream.entitystream.BaseConnector) ByteString(com.linkedin.data.ByteString) StreamResponse(com.linkedin.r2.message.stream.StreamResponse) ByteString(com.linkedin.data.ByteString) TransportResponse(com.linkedin.r2.transport.common.bridge.common.TransportResponse) StreamRequest(com.linkedin.r2.message.stream.StreamRequest) Observer(com.linkedin.r2.message.stream.entitystream.Observer) MessageType(com.linkedin.r2.transport.common.MessageType)

Example 13 with Observer

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();
}
Also used : EntityStream(com.linkedin.r2.message.stream.entitystream.EntityStream) ScheduledExecutorService(java.util.concurrent.ScheduledExecutorService) CountDownLatch(java.util.concurrent.CountDownLatch) Test(org.testng.annotations.Test)

Example 14 with Observer

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);
}
Also used : EntityStream(com.linkedin.r2.message.stream.entitystream.EntityStream) WriteHandle(com.linkedin.r2.message.stream.entitystream.WriteHandle) ScheduledExecutorService(java.util.concurrent.ScheduledExecutorService) ExecutorService(java.util.concurrent.ExecutorService) CountDownLatch(java.util.concurrent.CountDownLatch) Test(org.testng.annotations.Test)

Example 15 with Observer

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);
    }
}
Also used : EntityStream(com.linkedin.r2.message.stream.entitystream.EntityStream) CountDownLatch(java.util.concurrent.CountDownLatch)

Aggregations

EntityStream (com.linkedin.r2.message.stream.entitystream.EntityStream)15 Test (org.testng.annotations.Test)10 CountDownLatch (java.util.concurrent.CountDownLatch)9 ByteString (com.linkedin.data.ByteString)8 Observer (com.linkedin.r2.message.stream.entitystream.Observer)7 ScheduledExecutorService (java.util.concurrent.ScheduledExecutorService)6 ExecutorService (java.util.concurrent.ExecutorService)4 ReadHandle (com.linkedin.r2.message.stream.entitystream.ReadHandle)3 StreamRequest (com.linkedin.r2.message.stream.StreamRequest)2 StreamResponse (com.linkedin.r2.message.stream.StreamResponse)2 WriteHandle (com.linkedin.r2.message.stream.entitystream.WriteHandle)2 Callback (com.linkedin.common.callback.Callback)1 MultiPartMIMEFullReaderCallback (com.linkedin.multipart.utils.MIMETestUtils.MultiPartMIMEFullReaderCallback)1 SinglePartMIMEFullReaderCallback (com.linkedin.multipart.utils.MIMETestUtils.SinglePartMIMEFullReaderCallback)1 RequestContext (com.linkedin.r2.message.RequestContext)1 StreamRequestBuilder (com.linkedin.r2.message.stream.StreamRequestBuilder)1 BaseConnector (com.linkedin.r2.message.stream.entitystream.BaseConnector)1 ByteStringWriter (com.linkedin.r2.message.stream.entitystream.ByteStringWriter)1 FullEntityObserver (com.linkedin.r2.message.stream.entitystream.FullEntityObserver)1 MessageType (com.linkedin.r2.transport.common.MessageType)1