use of com.linkedin.r2.message.stream.entitystream.Writer in project rest.li by linkedin.
the class TestRestLiServer method testMultipartRelatedRequestNoUserAttachments.
@Test
public void testMultipartRelatedRequestNoUserAttachments() throws Exception {
// This test verifies the server's ability to handle a multipart related request that has only one part which is
// the rest.li payload; meaning there are no user defined attachments. Technically the client builders shouldn't do
// this but we allow this to keep the protocol somewhat flexible.
final AsyncStatusCollectionResource statusResource = getMockResource(AsyncStatusCollectionResource.class);
statusResource.streamingAction(EasyMock.<String>anyObject(), EasyMock.<RestLiAttachmentReader>anyObject(), EasyMock.<Callback<Long>>anyObject());
EasyMock.expectLastCall().andAnswer(new IAnswer<Object>() {
@Override
public Object answer() throws Throwable {
// Verify there are no attachments.
final RestLiAttachmentReader attachmentReader = (RestLiAttachmentReader) EasyMock.getCurrentArguments()[1];
Assert.assertNull(attachmentReader);
// Verify the action param.
Assert.assertEquals((String) EasyMock.getCurrentArguments()[0], "someMetadata");
// Now respond back to the request.
@SuppressWarnings("unchecked") Callback<Long> callback = (Callback<Long>) EasyMock.getCurrentArguments()[2];
callback.onSuccess(1234l);
return null;
}
});
replay(statusResource);
// Now we create a multipart/related payload.
final String payload = "{\"metadata\": \"someMetadata\"}";
final ByteStringWriter byteStringWriter = new ByteStringWriter(ByteString.copyString(payload, Charset.defaultCharset()));
final MultiPartMIMEWriter.Builder builder = new MultiPartMIMEWriter.Builder();
final MultiPartMIMEWriter writer = AttachmentUtils.createMultiPartMIMEWriter(byteStringWriter, "application/json", builder);
final StreamRequest streamRequest = MultiPartMIMEStreamRequestFactory.generateMultiPartMIMEStreamRequest(new URI("/asyncstatuses/?action=streamingAction"), "related", writer, Collections.<String, String>emptyMap(), "POST", ImmutableMap.of(RestConstants.HEADER_ACCEPT, RestConstants.HEADER_VALUE_MULTIPART_RELATED), Collections.emptyList());
final Callback<StreamResponse> callback = new Callback<StreamResponse>() {
@Override
public void onSuccess(StreamResponse streamResponse) {
Messages.toRestResponse(streamResponse, new Callback<RestResponse>() {
@Override
public void onError(Throwable e) {
Assert.fail();
}
@Override
public void onSuccess(RestResponse result) {
Assert.assertEquals(result.getStatus(), 200);
try {
ContentType contentType = new ContentType(streamResponse.getHeader(RestConstants.HEADER_CONTENT_TYPE));
Assert.assertEquals(contentType.getBaseType(), RestConstants.HEADER_VALUE_APPLICATION_JSON);
} catch (ParseException parseException) {
Assert.fail();
}
// Verify the response body
Assert.assertEquals(result.getEntity().asAvroString(), "{\"value\":1234}");
EasyMock.verify(statusResource);
EasyMock.reset(statusResource);
}
});
}
@Override
public void onError(Throwable e) {
fail();
}
};
_server.handleRequest(streamRequest, new RequestContext(), callback);
}
use of com.linkedin.r2.message.stream.entitystream.Writer in project rest.li by linkedin.
the class TestRestLiServer method testRequestAttachmentsResponseAttachmentsException.
@Test
public void testRequestAttachmentsResponseAttachmentsException() throws Exception {
// This test verifies the server's behavior in the face of an exception. In this case the resource method
// threw an exception AFTER setting response attachments. Additionally the resource method failed to absorb any
// incoming request attachments. We verify in this test that StreamResponseCallbackAdaptor in RestLiServer
// drains and absorbs all bytes from the incoming request and that the response attachments set by the resource method
// are told to abort.
// Define the server side resource attachments to be sent back.
final RestLiResponseAttachments.Builder responseAttachmentsBuilder = new RestLiResponseAttachments.Builder();
final RestLiTestAttachmentDataSource toBeAbortedDataSource = RestLiTestAttachmentDataSource.createWithRandomPayload("1");
responseAttachmentsBuilder.appendSingleAttachment(toBeAbortedDataSource);
Capture<ResourceContext> resourceContextCapture = EasyMock.newCapture();
final AsyncStatusCollectionResource statusResource = getMockResource(AsyncStatusCollectionResource.class, capture(resourceContextCapture));
statusResource.streamingAction(EasyMock.<String>anyObject(), EasyMock.<RestLiAttachmentReader>anyObject(), EasyMock.<Callback<Long>>anyObject());
EasyMock.expectLastCall().andAnswer(new IAnswer<Object>() {
@Override
public Object answer() throws Throwable {
// Verify there are still attachments to be read.
final RestLiAttachmentReader attachmentReader = (RestLiAttachmentReader) EasyMock.getCurrentArguments()[1];
Assert.assertFalse(attachmentReader.haveAllAttachmentsFinished());
// Verify the action param.
Assert.assertEquals((String) EasyMock.getCurrentArguments()[0], "someMetadata");
// Set the response attachments
resourceContextCapture.getValue().setResponseAttachments(responseAttachmentsBuilder.build());
// Now throw an exception.
@SuppressWarnings("unchecked") Callback<Long> callback = (Callback<Long>) EasyMock.getCurrentArguments()[2];
callback.onError(new RestLiServiceException(HttpStatus.S_409_CONFLICT, "Some conflict"));
return null;
}
});
replay(statusResource);
// Now we create a multipart/related payload.
final String payload = "{\"metadata\": \"someMetadata\"}";
final ByteStringWriter byteStringWriter = new ByteStringWriter(ByteString.copyString(payload, Charset.defaultCharset()));
final MultiPartMIMEWriter.Builder builder = new MultiPartMIMEWriter.Builder();
final RestLiTestAttachmentDataSource toBeDrainedDataSource = RestLiTestAttachmentDataSource.createWithRandomPayload("2");
AttachmentUtils.appendSingleAttachmentToBuilder(builder, toBeDrainedDataSource);
final MultiPartMIMEWriter writer = AttachmentUtils.createMultiPartMIMEWriter(byteStringWriter, "application/json", builder);
final StreamRequest streamRequest = MultiPartMIMEStreamRequestFactory.generateMultiPartMIMEStreamRequest(new URI("/asyncstatuses/?action=streamingAction"), "related", writer, Collections.<String, String>emptyMap(), "POST", ImmutableMap.of(RestConstants.HEADER_ACCEPT, RestConstants.HEADER_VALUE_MULTIPART_RELATED), Collections.emptyList());
final Callback<StreamResponse> callback = new Callback<StreamResponse>() {
@Override
public void onSuccess(StreamResponse streamResponse) {
fail();
}
@Override
public void onError(Throwable e) {
// Verify the exception.
assertTrue(e instanceof StreamException);
StreamException streamException = (StreamException) e;
StreamResponse streamResponse = streamException.getResponse();
assertEquals(streamResponse.getStatus(), 409);
final FullEntityReader fullEntityReader = new FullEntityReader(new Callback<ByteString>() {
@Override
public void onError(Throwable e) {
Assert.fail();
}
@Override
public void onSuccess(ByteString result) {
// We have the body so assert exception made it.
assertTrue(result.length() > 0);
assertTrue(result.asString(Charset.defaultCharset()).contains("Some conflict"));
}
});
streamResponse.getEntityStream().setReader(fullEntityReader);
EasyMock.verify(statusResource);
EasyMock.reset(statusResource);
}
};
_server.handleRequest(streamRequest, new RequestContext(), callback);
// Verify that the request level attachments were drained.
Assert.assertTrue(toBeDrainedDataSource.finished());
// Verify that response attachments were told to abort.
Assert.assertTrue(toBeAbortedDataSource.dataSourceAborted());
}
use of com.linkedin.r2.message.stream.entitystream.Writer in project rest.li by linkedin.
the class TestRestLiServer method testRequestAttachmentsAndResponseAttachments.
@Test
public void testRequestAttachmentsAndResponseAttachments() throws Exception {
// This test verifies the server's ability to accept request attachments and send back response attachments. This is the
// main test to verify the wire protocol for streaming. We send a payload that contains the rest.li payload and some attachments
// and we send a response back with a rest.li payload and some attachments.
// Define the server side resource attachments to be sent back.
final RestLiResponseAttachments.Builder responseAttachmentsBuilder = new RestLiResponseAttachments.Builder();
responseAttachmentsBuilder.appendSingleAttachment(new RestLiTestAttachmentDataSource("1", ByteString.copyString("one", Charset.defaultCharset())));
Capture<ResourceContext> resourceContextCapture = EasyMock.newCapture();
final AsyncStatusCollectionResource statusResource = getMockResource(AsyncStatusCollectionResource.class, capture(resourceContextCapture));
statusResource.streamingAction(EasyMock.<String>anyObject(), EasyMock.<RestLiAttachmentReader>anyObject(), EasyMock.<Callback<Long>>anyObject());
EasyMock.expectLastCall().andAnswer(new IAnswer<Object>() {
@Override
public Object answer() throws Throwable {
// Verify there are still attachments to be read.
final RestLiAttachmentReader attachmentReader = (RestLiAttachmentReader) EasyMock.getCurrentArguments()[1];
Assert.assertFalse(attachmentReader.haveAllAttachmentsFinished());
// Verify the action param.
Assert.assertEquals((String) EasyMock.getCurrentArguments()[0], "someMetadata");
// Set the response attachments
resourceContextCapture.getValue().setResponseAttachments(responseAttachmentsBuilder.build());
// Now respond back to the request.
@SuppressWarnings("unchecked") Callback<Long> callback = (Callback<Long>) EasyMock.getCurrentArguments()[2];
callback.onSuccess(1234l);
return null;
}
});
replay(statusResource);
// Now we create a multipart/related payload.
final String payload = "{\"metadata\": \"someMetadata\"}";
final ByteStringWriter byteStringWriter = new ByteStringWriter(ByteString.copyString(payload, Charset.defaultCharset()));
final MultiPartMIMEWriter.Builder builder = new MultiPartMIMEWriter.Builder();
AttachmentUtils.appendSingleAttachmentToBuilder(builder, new RestLiTestAttachmentDataSource("2", ByteString.copyString("two", Charset.defaultCharset())));
final MultiPartMIMEWriter writer = AttachmentUtils.createMultiPartMIMEWriter(byteStringWriter, "application/json", builder);
final StreamRequest streamRequest = MultiPartMIMEStreamRequestFactory.generateMultiPartMIMEStreamRequest(new URI("/asyncstatuses/?action=streamingAction"), "related", writer, Collections.<String, String>emptyMap(), "POST", ImmutableMap.of(RestConstants.HEADER_ACCEPT, RestConstants.HEADER_VALUE_MULTIPART_RELATED), Collections.emptyList());
final Callback<StreamResponse> callback = new Callback<StreamResponse>() {
@Override
public void onSuccess(StreamResponse streamResponse) {
// Before reading the data make sure top level type is multipart/related
Assert.assertEquals(streamResponse.getStatus(), 200);
try {
ContentType contentType = new ContentType(streamResponse.getHeader(RestConstants.HEADER_CONTENT_TYPE));
Assert.assertEquals(contentType.getBaseType(), RestConstants.HEADER_VALUE_MULTIPART_RELATED);
} catch (ParseException parseException) {
Assert.fail();
}
final CountDownLatch countDownLatch = new CountDownLatch(1);
MultiPartMIMEFullReaderCallback fullReaderCallback = new MultiPartMIMEFullReaderCallback(countDownLatch);
final MultiPartMIMEReader reader = MultiPartMIMEReader.createAndAcquireStream(streamResponse);
reader.registerReaderCallback(fullReaderCallback);
try {
countDownLatch.await(3000, TimeUnit.MILLISECONDS);
} catch (InterruptedException interruptedException) {
Assert.fail();
}
final List<SinglePartMIMEFullReaderCallback> singlePartMIMEReaderCallbacks = fullReaderCallback.getSinglePartMIMEReaderCallbacks();
Assert.assertEquals(singlePartMIMEReaderCallbacks.size(), 2);
// Verify first part is Action response.
Assert.assertEquals(singlePartMIMEReaderCallbacks.get(0).getHeaders().get(RestConstants.HEADER_CONTENT_TYPE), RestConstants.HEADER_VALUE_APPLICATION_JSON);
Assert.assertEquals(singlePartMIMEReaderCallbacks.get(0).getFinishedData().asAvroString(), "{\"value\":1234}");
// Verify the second part matches what the server should have sent back
Assert.assertEquals(singlePartMIMEReaderCallbacks.get(1).getHeaders().get(RestConstants.HEADER_CONTENT_ID), "1");
Assert.assertEquals(singlePartMIMEReaderCallbacks.get(1).getFinishedData().asString(Charset.defaultCharset()), "one");
EasyMock.verify(statusResource);
EasyMock.reset(statusResource);
}
@Override
public void onError(Throwable e) {
fail();
}
};
_server.handleRequest(streamRequest, new RequestContext(), callback);
}
use of com.linkedin.r2.message.stream.entitystream.Writer in project rest.li by linkedin.
the class Http2FrameListener method onHeadersRead.
@Override
public void onHeadersRead(ChannelHandlerContext ctx, int streamId, Http2Headers headers, int padding, boolean endOfStream) throws Http2Exception {
LOG.debug("Received HTTP/2 HEADERS frame, stream={}, end={}, headers={}, padding={}bytes", new Object[] { streamId, endOfStream, headers.size(), padding });
// Ignores response for the upgrade request
if (streamId == Http2CodecUtil.HTTP_UPGRADE_STREAM_ID) {
return;
}
// Refactored duplicate code to new code pipeline.
final StreamResponseBuilder builder = Http2MessageDecoders.ResponseDecoder.buildStreamResponse(headers);
// Gets async pool handle from stream properties
TimeoutAsyncPoolHandle<?> timeoutHandle = Http2PipelinePropertyUtil.remove(ctx, _connection, streamId, Http2ClientPipelineInitializer.CHANNEL_POOL_HANDLE_ATTR_KEY);
if (timeoutHandle == null) {
_lifecycleManager.onError(ctx, false, Http2Exception.connectionError(Http2Error.PROTOCOL_ERROR, "No channel pool handle is associated with this stream", streamId));
return;
}
final StreamResponse response;
if (endOfStream) {
response = builder.build(EntityStreams.emptyStream());
// Release the handle to put the channel back to the pool
timeoutHandle.release();
} else {
// Associate an entity stream writer to the HTTP/2 stream
final TimeoutBufferedWriter writer = new TimeoutBufferedWriter(ctx, streamId, _maxContentLength, timeoutHandle);
if (_connection.stream(streamId).setProperty(_writerKey, writer) != null) {
_lifecycleManager.onError(ctx, false, Http2Exception.connectionError(Http2Error.PROTOCOL_ERROR, "Another writer has already been associated with current stream ID", streamId));
return;
}
// Prepares StreamResponse for the channel pipeline
EntityStream entityStream = EntityStreams.newEntityStream(writer);
response = builder.build(entityStream);
}
// Gets callback from stream properties
TransportCallback<?> callback = Http2PipelinePropertyUtil.remove(ctx, _connection, streamId, Http2ClientPipelineInitializer.CALLBACK_ATTR_KEY);
if (callback != null) {
ctx.fireChannelRead(new ResponseWithCallback<Response, TransportCallback<?>>(response, callback));
}
}
use of com.linkedin.r2.message.stream.entitystream.Writer in project rest.li by linkedin.
the class ClientEntityStreamHandler method channelRead.
@Override
public void channelRead(ChannelHandlerContext ctx, Object msg) {
if (msg instanceof StreamResponseBuilder) {
final StreamResponseBuilder builder = (StreamResponseBuilder) msg;
final Map<String, String> headers = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
headers.putAll(builder.getHeaders());
final Map<String, String> wireAttrs = WireAttributeHelper.removeWireAttributes(headers);
final StreamWriter writer = new StreamWriter(ctx, _maxContentLength);
ctx.channel().attr(NettyChannelAttributes.RESPONSE_WRITER).set(writer);
final StreamResponse response = builder.unsafeSetHeaders(headers).build(EntityStreams.newEntityStream(writer));
final TransportCallback<StreamResponse> callback = ctx.channel().attr(NettyChannelAttributes.RESPONSE_CALLBACK).getAndSet(null);
if (callback != null) {
callback.onResponse(TransportResponseImpl.success(response, wireAttrs));
}
} else if (msg instanceof ByteString) {
final StreamWriter writer = msg == StreamWriter.EOF ? ctx.channel().attr(NettyChannelAttributes.RESPONSE_WRITER).getAndSet(null) : ctx.channel().attr(NettyChannelAttributes.RESPONSE_WRITER).get();
if (writer != null) {
writer.onDataAvailable((ByteString) msg);
}
} else {
ctx.fireChannelRead(msg);
}
}
Aggregations