Search in sources :

Example 6 with Writer

use of com.linkedin.r2.message.stream.entitystream.Writer in project rest.li by linkedin.

the class TestStreamingCompression method testDecompress.

private void testDecompress(StreamingCompressor compressor, byte[] uncompressed, byte[] compressed) throws CompressionException, ExecutionException, InterruptedException {
    ByteWriter writer = new ByteWriter(compressed);
    EntityStream compressedStream = EntityStreams.newEntityStream(writer);
    EntityStream uncompressedStream = compressor.inflate(compressedStream);
    FutureCallback<byte[]> callback = new FutureCallback<byte[]>();
    uncompressedStream.setReader(new ByteReader(callback));
    byte[] result = callback.get();
    Assert.assertEquals(result, uncompressed);
}
Also used : EntityStream(com.linkedin.r2.message.stream.entitystream.EntityStream) FutureCallback(com.linkedin.common.callback.FutureCallback)

Example 7 with Writer

use of com.linkedin.r2.message.stream.entitystream.Writer in project rest.li by linkedin.

the class TestStreamingCompression method testCompressThenDecompress.

private void testCompressThenDecompress(StreamingCompressor compressor, byte[] origin) throws CompressionException, ExecutionException, InterruptedException {
    ByteWriter writer = new ByteWriter(origin);
    EntityStream uncompressedStream = EntityStreams.newEntityStream(writer);
    EntityStream compressedStream = compressor.deflate(uncompressedStream);
    EntityStream decompressedStream = compressor.inflate(compressedStream);
    FutureCallback<byte[]> callback = new FutureCallback<byte[]>();
    decompressedStream.setReader(new ByteReader(callback));
    byte[] result = callback.get();
    Assert.assertEquals(result, origin);
}
Also used : EntityStream(com.linkedin.r2.message.stream.entitystream.EntityStream) FutureCallback(com.linkedin.common.callback.FutureCallback)

Example 8 with Writer

use of com.linkedin.r2.message.stream.entitystream.Writer in project rest.li by linkedin.

the class TestStreamRequest method testBackPressure.

@Test
public void testBackPressure() throws Exception {
    for (Client client : clients()) {
        final long totalBytes = SMALL_BYTES_NUM;
        TimedBytesWriter writer = new TimedBytesWriter(totalBytes, BYTE);
        EntityStream entityStream = EntityStreams.newEntityStream(writer);
        StreamRequestBuilder builder = new StreamRequestBuilder(Bootstrap.createHttpURI(PORT, RATE_LIMITED_URI));
        StreamRequest request = builder.setMethod("POST").build(entityStream);
        final AtomicInteger status = new AtomicInteger(-1);
        final CountDownLatch latch = new CountDownLatch(1);
        Callback<StreamResponse> callback = expectSuccessCallback(latch, status);
        client.streamRequest(request, callback);
        latch.await(60000, TimeUnit.MILLISECONDS);
        Assert.assertEquals(status.get(), RestStatus.OK);
        TimedBytesReader reader = _rateLimitedRequestHandler.getReader();
        Assert.assertNotNull(reader);
        Assert.assertEquals(totalBytes, reader.getTotalBytes());
        Assert.assertTrue(reader.allBytesCorrect());
        long clientSendTimespan = writer.getStopTime() - writer.getStartTime();
        long serverReceiveTimespan = reader.getStopTime() - reader.getStartTime();
        Assert.assertTrue(serverReceiveTimespan > 1000);
        double diff = Math.abs(serverReceiveTimespan - clientSendTimespan);
        double diffRatio = diff / clientSendTimespan;
        // make it generous to reduce the chance occasional test failures
        Assert.assertTrue(diffRatio < 0.2);
    }
}
Also used : EntityStream(com.linkedin.r2.message.stream.entitystream.EntityStream) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) StreamResponse(com.linkedin.r2.message.stream.StreamResponse) Client(com.linkedin.r2.transport.common.Client) CountDownLatch(java.util.concurrent.CountDownLatch) StreamRequestBuilder(com.linkedin.r2.message.stream.StreamRequestBuilder) StreamRequest(com.linkedin.r2.message.stream.StreamRequest) Test(org.testng.annotations.Test)

Example 9 with Writer

use of com.linkedin.r2.message.stream.entitystream.Writer in project rest.li by linkedin.

the class TestEntityStream method testRaceBetweenDoneAndCancel.

private void testRaceBetweenDoneAndCancel(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.done();
            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));
    // in any case, reader shouldn't fail
    Assert.assertEquals(reader.errorTimes(), 0);
    // if done wins the race
    if (reader.doneTimes() > 0) {
        Assert.assertEquals(reader.doneTimes(), 1);
        Assert.assertEquals(observer.doneTimes(), 1);
        Assert.assertEquals(observer.errorTimes(), 0);
        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);
    }
}
Also used : EntityStream(com.linkedin.r2.message.stream.entitystream.EntityStream) CountDownLatch(java.util.concurrent.CountDownLatch)

Example 10 with Writer

use of com.linkedin.r2.message.stream.entitystream.Writer in project rest.li by linkedin.

the class TestEntityStream method testReaderInitErrorThrowInStreaming.

@Test
public void testReaderInitErrorThrowInStreaming() throws Exception {
    final CountDownLatch latch = new CountDownLatch(1);
    final TestWriter writer = new TestWriter() {

        @Override
        public void onWritePossible() {
            super.onWritePossible();
            write();
        }
    };
    TestObserver observer = new TestObserver();
    final EntityStream es = EntityStreams.newEntityStream(writer);
    es.addObserver(observer);
    final ControlReader reader = new ControlReader() {

        @Override
        public void onInit(ReadHandle rh) {
            rh.request(1);
        }

        @Override
        public void onDataAvailable(ByteString data) {
            super.onDataAvailable(data);
            throw new RuntimeException();
        }
    };
    ExecutorService executor = Executors.newFixedThreadPool(1);
    executor.submit(new Runnable() {

        @Override
        public void run() {
            es.setReader(reader);
            latch.countDown();
        }
    });
    Assert.assertTrue(latch.await(1000, TimeUnit.MILLISECONDS));
    Assert.assertEquals(reader.errorTimes(), 1);
    Assert.assertEquals(reader.getChunkCount(), 1);
    Assert.assertEquals(observer.errorTimes(), 1);
    Assert.assertEquals(reader.getChunkCount(), 1);
    Assert.assertEquals(writer.abortedTimes(), 1);
    Assert.assertEquals(writer.getWritePossibleCount(), 1);
}
Also used : EntityStream(com.linkedin.r2.message.stream.entitystream.EntityStream) ReadHandle(com.linkedin.r2.message.stream.entitystream.ReadHandle) ByteString(com.linkedin.data.ByteString) ScheduledExecutorService(java.util.concurrent.ScheduledExecutorService) ExecutorService(java.util.concurrent.ExecutorService) CountDownLatch(java.util.concurrent.CountDownLatch) Test(org.testng.annotations.Test)

Aggregations

Test (org.testng.annotations.Test)29 EntityStream (com.linkedin.r2.message.stream.entitystream.EntityStream)22 CountDownLatch (java.util.concurrent.CountDownLatch)17 StreamRequest (com.linkedin.r2.message.stream.StreamRequest)16 StreamResponse (com.linkedin.r2.message.stream.StreamResponse)16 StreamRequestBuilder (com.linkedin.r2.message.stream.StreamRequestBuilder)12 ByteString (com.linkedin.data.ByteString)11 FutureCallback (com.linkedin.common.callback.FutureCallback)10 MultiPartMIMEWriter (com.linkedin.multipart.MultiPartMIMEWriter)8 RequestContext (com.linkedin.r2.message.RequestContext)6 ScheduledExecutorService (java.util.concurrent.ScheduledExecutorService)6 ReadHandle (com.linkedin.r2.message.stream.entitystream.ReadHandle)5 None (com.linkedin.common.util.None)4 SinglePartMIMEFullReaderCallback (com.linkedin.multipart.utils.MIMETestUtils.SinglePartMIMEFullReaderCallback)4 RestRequestBuilder (com.linkedin.r2.message.rest.RestRequestBuilder)4 ByteStringWriter (com.linkedin.r2.message.stream.entitystream.ByteStringWriter)4 WriteHandle (com.linkedin.r2.message.stream.entitystream.WriteHandle)4 TransportDispatcherBuilder (com.linkedin.r2.transport.common.bridge.server.TransportDispatcherBuilder)4 RestLiAttachmentReader (com.linkedin.restli.common.attachments.RestLiAttachmentReader)4 URI (java.net.URI)4