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

Example 71 with StreamResponse

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

the class TrackerClientTest method testClientStreamRequest.

@Test(groups = { "small", "back-end" })
public void testClientStreamRequest() throws URISyntaxException {
    URI uri = URI.create("http://test.qa.com:1234/foo");
    double weight = 3d;
    TestClient wrappedClient = new TestClient(true);
    Clock clock = new SettableClock();
    Map<Integer, PartitionData> partitionDataMap = new HashMap<Integer, PartitionData>(2);
    partitionDataMap.put(DefaultPartitionAccessor.DEFAULT_PARTITION_ID, new PartitionData(3d));
    TrackerClient client = new TrackerClient(uri, partitionDataMap, wrappedClient, clock, null);
    assertEquals(client.getUri(), uri);
    Double clientWeight = client.getPartitionWeight(DefaultPartitionAccessor.DEFAULT_PARTITION_ID);
    assertEquals(clientWeight, weight);
    assertEquals(client.getWrappedClient(), wrappedClient);
    StreamRequest streamRequest = new StreamRequestBuilder(uri).build(EntityStreams.emptyStream());
    Map<String, String> restWireAttrs = new HashMap<String, String>();
    TestTransportCallback<StreamResponse> restCallback = new TestTransportCallback<StreamResponse>();
    client.streamRequest(streamRequest, new RequestContext(), restWireAttrs, restCallback);
    assertFalse(restCallback.response.hasError());
    assertSame(wrappedClient.streamRequest, streamRequest);
    assertEquals(wrappedClient.restWireAttrs, restWireAttrs);
}

Example 72 with StreamResponse

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

the class RewriteClientTestStreamRequest method testPathAppend.

@Test
public void testPathAppend() {
    URI uri = URI.create("http://test.linkedin.com:9876/test");
    String serviceName = "HistoryService";
    TestClient wrappedClient = new TestClient();
    RewriteClient client = new RewriteClient(serviceName, uri, wrappedClient);
    assertEquals(client.getUri(), uri);
    assertEquals(client.getServiceName(), serviceName);
    assertEquals(client.getWrappedClient(), wrappedClient);
    StreamRequest streamRequest;
    Map<String, String> restWireAttrs = new HashMap<String, String>();
    TestTransportCallback<StreamResponse> restCallback = new TestTransportCallback<StreamResponse>();
    streamRequest = getRequest("d2://HistoryService");
    client.streamRequest(streamRequest, new RequestContext(), restWireAttrs, restCallback);
    checkRewrite(wrappedClient, streamRequest, restCallback, "http://test.linkedin.com:9876/test");
    streamRequest = getRequest("d2://HistoryService/");
    client.streamRequest(streamRequest, new RequestContext(), restWireAttrs, restCallback);
    checkRewrite(wrappedClient, streamRequest, restCallback, "http://test.linkedin.com:9876/test/");
    streamRequest = getRequest("d2://HistoryService//");
    client.streamRequest(streamRequest, new RequestContext(), restWireAttrs, restCallback);
    checkRewrite(wrappedClient, streamRequest, restCallback, "http://test.linkedin.com:9876/test//");
    streamRequest = getRequest("d2://HistoryService/foo");
    client.streamRequest(streamRequest, new RequestContext(), restWireAttrs, restCallback);
    checkRewrite(wrappedClient, streamRequest, restCallback, "http://test.linkedin.com:9876/test/foo");
    streamRequest = getRequest("d2://HistoryService/foo/");
    client.streamRequest(streamRequest, new RequestContext(), restWireAttrs, restCallback);
    checkRewrite(wrappedClient, streamRequest, restCallback, "http://test.linkedin.com:9876/test/foo/");
}

Example 73 with StreamResponse

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

the class DispatcherRequestFilter method onStreamRequest.

@Override
public void onStreamRequest(StreamRequest req, RequestContext requestContext, Map<String, String> wireAttrs, NextFilter<StreamRequest, StreamResponse> nextFilter) {
    Connector connector = null;
    try {
        final AtomicBoolean responded = new AtomicBoolean(false);
        TransportCallback<StreamResponse> callback = createStreamCallback(requestContext, nextFilter, responded);
        connector = new Connector(responded, nextFilter, requestContext, wireAttrs);
        req.getEntityStream().setReader(connector);
        EntityStream newStream = EntityStreams.newEntityStream(connector);
        _dispatcher.handleStreamRequest(req.builder().build(newStream), wireAttrs, requestContext, callback);
    } catch (Exception e) {
        nextFilter.onError(e, requestContext, new HashMap<String, String>());
        if (connector != null) {
            connector.cancel();
        }
    }
}

Example 74 with StreamResponse

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

the class Messages method toRestResponse.

/**
   * Converts a StreamResponse to RestResponse
   * @param streamResponse the stream request to be converted
   * @param callback the callback to be invoked when the rest response is constructed
   * @param addContentLengthHeader whether the rest response should have content-length header
   */
public static void toRestResponse(StreamResponse streamResponse, final Callback<RestResponse> callback, final boolean addContentLengthHeader) {
    final RestResponseBuilder builder = new RestResponseBuilder(streamResponse);
    Callback<ByteString> assemblyCallback = new Callback<ByteString>() {

        @Override
        public void onError(Throwable e) {
            callback.onError(e);
        }

        @Override
        public void onSuccess(ByteString result) {
            if (addContentLengthHeader) {
                builder.setHeader(HttpConstants.CONTENT_LENGTH, String.valueOf(result.length()));
            }
            RestResponse restResponse = builder.setEntity(result).build();
            callback.onSuccess(restResponse);
        }
    };
    streamResponse.getEntityStream().setReader(new FullEntityReader(assemblyCallback));
}

Example 75 with StreamResponse

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

the class ServletHelper method writeResponseHeadersToServletResponse.

static StreamResponse writeResponseHeadersToServletResponse(TransportResponse<StreamResponse> response, HttpServletResponse resp) {
    Map<String, String> wireAttrs = response.getWireAttributes();
    for (Map.Entry<String, String> e : WireAttributeHelper.toWireAttributes(wireAttrs).entrySet()) {
        resp.setHeader(e.getKey(), e.getValue());
    }
    StreamResponse streamResponse;
    if (response.hasError()) {
        Throwable e = response.getError();
        if (e instanceof StreamException) {
            streamResponse = ((StreamException) e).getResponse();
        } else {
            streamResponse = Messages.toStreamResponse(RestStatus.responseForError(RestStatus.INTERNAL_SERVER_ERROR, e));
        }
    } else {
        streamResponse = response.getResponse();
    }
    resp.setStatus(streamResponse.getStatus());
    Map<String, String> headers = streamResponse.getHeaders();
    for (Map.Entry<String, String> e : headers.entrySet()) {
        // TODO multi-valued headers
        resp.setHeader(e.getKey(), e.getValue());
    }
    for (String cookie : streamResponse.getCookies()) {
        resp.addHeader(HttpConstants.RESPONSE_COOKIE_HEADER_NAME, cookie);
    }
    return streamResponse;
}