Examples with StreamResponse com.linkedin.r2.message.stream.StreamResponse used on opensource projects

Example 76 with StreamResponse

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

the class AbstractAsyncR2StreamServlet method service.

@Override
public void service(final HttpServletRequest req, final HttpServletResponse resp) throws ServletException, IOException {
    final AsyncContext ctx = req.startAsync(req, resp);
    ctx.setTimeout(_timeout);
    final WrappedAsyncContext wrappedCtx = new WrappedAsyncContext(ctx);
    final AsyncEventIOHandler ioHandler = new AsyncEventIOHandler(req.getInputStream(), resp.getOutputStream(), wrappedCtx, MAX_BUFFERED_CHUNKS);
    final RequestContext requestContext = ServletHelper.readRequestContext(req);
    final StreamRequest streamRequest;
    try {
        streamRequest = ServletHelper.readFromServletRequest(req, ioHandler);
    } catch (URISyntaxException e) {
        ServletHelper.writeToServletError(resp, RestStatus.BAD_REQUEST, e.toString());
        wrappedCtx.complete();
        return;
    }
    final AtomicBoolean startedResponding = new AtomicBoolean(false);
    ctx.addListener(new AsyncListener() {

        @Override
        public void onTimeout(AsyncEvent event) throws IOException {
            LOG.error("Server timeout for request: " + formatURI(req.getRequestURI()));
            if (startedResponding.compareAndSet(false, true)) {
                LOG.info("Returning server timeout response");
                ServletHelper.writeToServletError(resp, RestStatus.INTERNAL_SERVER_ERROR, "Server timeout");
            } else {
                req.setAttribute(ASYNC_IOEXCEPTION, new ServletException("Server timeout"));
            }
            ioHandler.exitLoop();
            wrappedCtx.complete();
        }

        @Override
        public void onStartAsync(AsyncEvent event) throws IOException {
        // Nothing to do here
        }

        @Override
        public void onError(AsyncEvent event) throws IOException {
            LOG.error("Server error for request: " + formatURI(req.getRequestURI()));
            if (startedResponding.compareAndSet(false, true)) {
                LOG.info("Returning server error response");
                ServletHelper.writeToServletError(resp, RestStatus.INTERNAL_SERVER_ERROR, "Server error");
            } else {
                req.setAttribute(ASYNC_IOEXCEPTION, new ServletException("Server error"));
            }
            ioHandler.exitLoop();
            wrappedCtx.complete();
        }

        @Override
        public void onComplete(AsyncEvent event) throws IOException {
            Object exception = req.getAttribute(ASYNC_IOEXCEPTION);
            if (exception != null) {
                throw new IOException((Throwable) exception);
            }
        }
    });
    final TransportCallback<StreamResponse> callback = new TransportCallback<StreamResponse>() {

        @Override
        public void onResponse(final TransportResponse<StreamResponse> response) {
            if (startedResponding.compareAndSet(false, true)) {
                ctx.start(new Runnable() {

                    @Override
                    public void run() {
                        try {
                            StreamResponse streamResponse = ServletHelper.writeResponseHeadersToServletResponse(response, resp);
                            streamResponse.getEntityStream().setReader(ioHandler);
                            ioHandler.loop();
                        } catch (Exception e) {
                            req.setAttribute(ASYNC_IOEXCEPTION, e);
                            wrappedCtx.complete();
                        }
                    }
                });
            } else {
                LOG.error("Dropped a response; this is mostly like because that AsyncContext timeout or error had already happened");
            }
        }
    };
    // we have to use a new thread and let this thread return to pool. otherwise the timeout won't start
    ctx.start(new Runnable() {

        @Override
        public void run() {
            try {
                getDispatcher().handleRequest(streamRequest, requestContext, callback);
                ioHandler.loop();
            } catch (Exception e) {
                req.setAttribute(ASYNC_IOEXCEPTION, e);
                wrappedCtx.complete();
            }
        }
    });
}

Example 77 with StreamResponse

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

the class AbstractR2StreamServlet method service.

@Override
protected void service(final HttpServletRequest req, final HttpServletResponse resp) throws ServletException, IOException {
    final SyncIOHandler ioHandler = new SyncIOHandler(req.getInputStream(), resp.getOutputStream(), 2, _ioHandlerTimeout);
    RequestContext requestContext = ServletHelper.readRequestContext(req);
    StreamRequest streamRequest;
    try {
        streamRequest = ServletHelper.readFromServletRequest(req, ioHandler);
    } catch (URISyntaxException e) {
        ServletHelper.writeToServletError(resp, RestStatus.BAD_REQUEST, e.toString());
        return;
    }
    TransportCallback<StreamResponse> callback = new TransportCallback<StreamResponse>() {

        @Override
        public void onResponse(TransportResponse<StreamResponse> response) {
            StreamResponse streamResponse = ServletHelper.writeResponseHeadersToServletResponse(response, resp);
            streamResponse.getEntityStream().setReader(ioHandler);
        }
    };
    getDispatcher().handleRequest(streamRequest, requestContext, callback);
    ioHandler.loop();
}

Example 78 with StreamResponse

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

the class TestMIMEIntegrationReaderDrain method executeRequestWithDrainStrategy.

///////////////////////////////////////////////////////////////////////////////////////
private MimeMultipart executeRequestWithDrainStrategy(final int chunkSize, final List<MimeBodyPart> bodyPartList, final String drainStrategy, final String serverHeaderPrefix) throws Exception {
    MimeMultipart multiPartMimeBody = new MimeMultipart();
    //Add your body parts
    for (final MimeBodyPart bodyPart : bodyPartList) {
        multiPartMimeBody.addBodyPart(bodyPart);
    }
    final ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
    multiPartMimeBody.writeTo(byteArrayOutputStream);
    final ByteString requestPayload = ByteString.copy(byteArrayOutputStream.toByteArray());
    final VariableByteStringWriter variableByteStringWriter = new VariableByteStringWriter(requestPayload, chunkSize);
    final EntityStream entityStream = EntityStreams.newEntityStream(variableByteStringWriter);
    final StreamRequestBuilder builder = new StreamRequestBuilder(Bootstrap.createHttpURI(PORT, SERVER_URI));
    StreamRequest request = builder.setMethod("POST").setHeader(HEADER_CONTENT_TYPE, multiPartMimeBody.getContentType()).setHeader(DRAIN_HEADER, drainStrategy).build(entityStream);
    final AtomicInteger status = new AtomicInteger(-1);
    final CountDownLatch latch = new CountDownLatch(1);
    final Map<String, String> responseHeaders = new HashMap<String, String>();
    Callback<StreamResponse> callback = expectSuccessCallback(latch, status, responseHeaders);
    _client.streamRequest(request, callback);
    latch.await(TEST_TIMEOUT, TimeUnit.MILLISECONDS);
    Assert.assertEquals(status.get(), RestStatus.OK);
    Assert.assertEquals(responseHeaders.get(DRAIN_HEADER), serverHeaderPrefix + drainStrategy);
    return multiPartMimeBody;
}

Example 79 with StreamResponse

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

the class TestMIMEIntegrationReaderWriter method executeRequestAndAssert.

///////////////////////////////////////////////////////////////////////////////////////
//This method will execute a StreamRequest using a payload comprised of the provided MultiPartMIMEWriter. The request
//makes its way to a server which will read all the parts one by one and store them in a list. Subsequently,
//an assertion is performed to make sure that each part was read and saved properly on the server.
private void executeRequestAndAssert(final MultiPartMIMEWriter requestWriter, final List<MIMEDataPart> expectedParts) throws Exception {
    final StreamRequest streamRequest = MultiPartMIMEStreamRequestFactory.generateMultiPartMIMEStreamRequest(Bootstrap.createHttpURI(PORT, SERVER_URI), "mixed", requestWriter, Collections.<String, String>emptyMap());
    final AtomicInteger status = new AtomicInteger(-1);
    final CountDownLatch latch = new CountDownLatch(1);
    Callback<StreamResponse> callback = expectSuccessCallback(latch, status, new HashMap<String, String>());
    _client.streamRequest(streamRequest, callback);
    latch.await(TEST_TIMEOUT, TimeUnit.MILLISECONDS);
    Assert.assertEquals(status.get(), RestStatus.OK);
    List<SinglePartMIMEReaderCallbackImpl> singlePartMIMEReaderCallbacks = _mimeServerRequestHandler.getTestMultiPartMIMEReaderCallback().getSinglePartMIMEReaderCallbacks();
    _mimeServerRequestHandler.getTestMultiPartMIMEReaderCallback().getSinglePartMIMEReaderCallbacks();
    Assert.assertEquals(singlePartMIMEReaderCallbacks.size(), expectedParts.size());
    for (int i = 0; i < singlePartMIMEReaderCallbacks.size(); i++) {
        //Actual
        final SinglePartMIMEReaderCallbackImpl currentCallback = singlePartMIMEReaderCallbacks.get(i);
        //Expected
        final MIMEDataPart currentExpectedPart = expectedParts.get(i);
        Assert.assertEquals(currentCallback.getHeaders(), currentExpectedPart.getPartHeaders());
        Assert.assertEquals(currentCallback.getFinishedData(), currentExpectedPart.getPartData());
    }
}

Example 80 with StreamResponse

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

the class TestMIMEChainingReader method testMimeReaderDataSource.

//Verifies that a multi part mime reader can be used as a data source to the writer.
//To make the test easier to write, we simply chain back to the client in the form of simulating a response.
@Test(dataProvider = "chunkSizes")
public void testMimeReaderDataSource(final int chunkSize) throws Exception {
    final List<MultiPartMIMEDataSourceWriter> dataSources = generateInputStreamDataSources(chunkSize, _scheduledExecutorService);
    final MultiPartMIMEWriter writer = new MultiPartMIMEWriter.Builder().appendDataSources(dataSources).build();
    final StreamRequest streamRequest = mock(StreamRequest.class);
    when(streamRequest.getEntityStream()).thenReturn(writer.getEntityStream());
    final String contentTypeHeader = "multipart/mixed; boundary=" + writer.getBoundary();
    when(streamRequest.getHeader(MultiPartMIMEUtils.CONTENT_TYPE_HEADER)).thenReturn(contentTypeHeader);
    //Client side preparation to read the part back on the callback
    //Note the chunks size will carry over since the client is controlling how much data he gets back
    //based on the chunk size he writes.
    final CountDownLatch latch = new CountDownLatch(1);
    ClientMultiPartMIMEReaderReceiverCallback _clientReceiver = new ClientMultiPartMIMEReaderReceiverCallback(latch);
    Callback<StreamResponse> callback = generateSuccessChainCallback(_clientReceiver);
    //Server side start
    MultiPartMIMEReader reader = MultiPartMIMEReader.createAndAcquireStream(streamRequest);
    ServerMultiPartMIMEChainReaderWriterCallback _serverSender = new ServerMultiPartMIMEChainReaderWriterCallback(callback, reader);
    reader.registerReaderCallback(_serverSender);
    latch.await(_testTimeout, TimeUnit.MILLISECONDS);
    //Verify client. No need to verify the server.
    List<MIMETestUtils.SinglePartMIMEFullReaderCallback> singlePartMIMEReaderCallbacks = _clientReceiver.getSinglePartMIMEReaderCallbacks();
    Assert.assertEquals(singlePartMIMEReaderCallbacks.size(), 4);
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(0).getFinishedData(), BODY_A.getPartData());
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(0).getHeaders(), BODY_A.getPartHeaders());
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(1).getFinishedData(), BODY_B.getPartData());
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(1).getHeaders(), BODY_B.getPartHeaders());
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(2).getFinishedData(), BODY_C.getPartData());
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(2).getHeaders(), BODY_C.getPartHeaders());
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(3).getFinishedData(), BODY_D.getPartData());
    Assert.assertEquals(singlePartMIMEReaderCallbacks.get(3).getHeaders(), BODY_D.getPartHeaders());
}