Search in sources :

Example 1 with StreamResponseBuilder

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

the class TrackerClientTest method testCallTrackingStreamRequest.

@Test
public void testCallTrackingStreamRequest() throws Exception {
    URI uri = URI.create("http://test.qa.com:1234/foo");
    SettableClock clock = new SettableClock();
    AtomicInteger action = new AtomicInteger(0);
    TransportClient tc = new TransportClient() {

        @Override
        public void restRequest(RestRequest request, RequestContext requestContext, Map<String, String> wireAttrs, TransportCallback<RestResponse> callback) {
        }

        @Override
        public void streamRequest(StreamRequest request, RequestContext requestContext, Map<String, String> wireAttrs, TransportCallback<StreamResponse> callback) {
            clock.addDuration(5);
            switch(action.get()) {
                // success
                case 0:
                    callback.onResponse(TransportResponseImpl.success(new StreamResponseBuilder().build(EntityStreams.emptyStream())));
                    break;
                // fail with stream exception
                case 1:
                    callback.onResponse(TransportResponseImpl.error(new StreamException(new StreamResponseBuilder().setStatus(500).build(EntityStreams.emptyStream()))));
                    break;
                // fail with timeout exception
                case 2:
                    callback.onResponse(TransportResponseImpl.error(new RemoteInvocationException(new TimeoutException())));
                    break;
                // fail with other exception
                default:
                    callback.onResponse(TransportResponseImpl.error(new RuntimeException()));
                    break;
            }
        }

        @Override
        public void shutdown(Callback<None> callback) {
        }
    };
    TrackerClient client = createTrackerClient(tc, clock, uri);
    CallTracker callTracker = client.getCallTracker();
    CallTracker.CallStats stats;
    DegraderControl degraderControl = client.getDegraderControl(DefaultPartitionAccessor.DEFAULT_PARTITION_ID);
    DelayConsumeCallback delayConsumeCallback = new DelayConsumeCallback();
    client.streamRequest(new StreamRequestBuilder(uri).build(EntityStreams.emptyStream()), new RequestContext(), new HashMap<>(), delayConsumeCallback);
    clock.addDuration(5);
    // we only recorded the time when stream response arrives, but callcompletion.endcall hasn't been called yet.
    Assert.assertEquals(callTracker.getCurrentCallCountTotal(), 0);
    Assert.assertEquals(callTracker.getCurrentErrorCountTotal(), 0);
    // delay
    clock.addDuration(100);
    delayConsumeCallback.consume();
    clock.addDuration(5000);
    // now that we consumed the entity stream, callcompletion.endcall has been called.
    stats = callTracker.getCallStats();
    Assert.assertEquals(stats.getCallCount(), 1);
    Assert.assertEquals(stats.getErrorCount(), 0);
    Assert.assertEquals(stats.getCallCountTotal(), 1);
    Assert.assertEquals(stats.getErrorCountTotal(), 0);
    Assert.assertEquals(degraderControl.getCurrentComputedDropRate(), 0.0, 0.001);
    action.set(1);
    client.streamRequest(new StreamRequestBuilder(uri).build(EntityStreams.emptyStream()), new RequestContext(), new HashMap<>(), delayConsumeCallback);
    clock.addDuration(5);
    // we endcall with error immediately for stream exception, even before the entity is consumed
    Assert.assertEquals(callTracker.getCurrentCallCountTotal(), 2);
    Assert.assertEquals(callTracker.getCurrentErrorCountTotal(), 1);
    delayConsumeCallback.consume();
    clock.addDuration(5000);
    // no change in tracking after entity is consumed
    stats = callTracker.getCallStats();
    Assert.assertEquals(stats.getCallCount(), 1);
    Assert.assertEquals(stats.getErrorCount(), 1);
    Assert.assertEquals(stats.getCallCountTotal(), 2);
    Assert.assertEquals(stats.getErrorCountTotal(), 1);
    Assert.assertEquals(degraderControl.getCurrentComputedDropRate(), 0.2, 0.001);
    action.set(2);
    client.streamRequest(new StreamRequestBuilder(uri).build(EntityStreams.emptyStream()), new RequestContext(), new HashMap<>(), new TestTransportCallback<>());
    clock.addDuration(5);
    Assert.assertEquals(callTracker.getCurrentCallCountTotal(), 3);
    Assert.assertEquals(callTracker.getCurrentErrorCountTotal(), 2);
    clock.addDuration(5000);
    stats = callTracker.getCallStats();
    Assert.assertEquals(stats.getCallCount(), 1);
    Assert.assertEquals(stats.getErrorCount(), 1);
    Assert.assertEquals(stats.getCallCountTotal(), 3);
    Assert.assertEquals(stats.getErrorCountTotal(), 2);
    Assert.assertEquals(degraderControl.getCurrentComputedDropRate(), 0.4, 0.001);
    action.set(3);
    client.streamRequest(new StreamRequestBuilder(uri).build(EntityStreams.emptyStream()), new RequestContext(), new HashMap<>(), new TestTransportCallback<>());
    clock.addDuration(5);
    Assert.assertEquals(callTracker.getCurrentCallCountTotal(), 4);
    Assert.assertEquals(callTracker.getCurrentErrorCountTotal(), 3);
    clock.addDuration(5000);
    stats = callTracker.getCallStats();
    Assert.assertEquals(stats.getCallCount(), 1);
    Assert.assertEquals(stats.getErrorCount(), 1);
    Assert.assertEquals(stats.getCallCountTotal(), 4);
    Assert.assertEquals(stats.getErrorCountTotal(), 3);
    Assert.assertEquals(degraderControl.getCurrentComputedDropRate(), 0.2, 0.001);
}
Also used : TransportCallback(com.linkedin.r2.transport.common.bridge.common.TransportCallback) TransportClient(com.linkedin.r2.transport.common.bridge.client.TransportClient) StreamResponseBuilder(com.linkedin.r2.message.stream.StreamResponseBuilder) DegraderControl(com.linkedin.util.degrader.DegraderControl) URI(java.net.URI) StreamRequestBuilder(com.linkedin.r2.message.stream.StreamRequestBuilder) StreamRequest(com.linkedin.r2.message.stream.StreamRequest) StreamException(com.linkedin.r2.message.stream.StreamException) RestRequest(com.linkedin.r2.message.rest.RestRequest) Callback(com.linkedin.common.callback.Callback) TransportCallback(com.linkedin.r2.transport.common.bridge.common.TransportCallback) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) SettableClock(com.linkedin.util.clock.SettableClock) RequestContext(com.linkedin.r2.message.RequestContext) RemoteInvocationException(com.linkedin.r2.RemoteInvocationException) HashMap(java.util.HashMap) Map(java.util.Map) TimeoutException(java.util.concurrent.TimeoutException) CallTracker(com.linkedin.util.degrader.CallTracker) Test(org.testng.annotations.Test)

Example 2 with StreamResponseBuilder

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

the class RAPStreamResponseHandler method channelRead0.

@Override
protected void channelRead0(ChannelHandlerContext ctx, StreamResponse response) throws Exception {
    final Map<String, String> headers = new HashMap<String, String>(response.getHeaders());
    final Map<String, String> wireAttrs = new HashMap<String, String>(WireAttributeHelper.removeWireAttributes(headers));
    final StreamResponse newResponse = new StreamResponseBuilder(response).unsafeSetHeaders(headers).build(response.getEntityStream());
    // In general there should always be a callback to handle a received message,
    // but it could have been removed due to a previous exception or closure on the
    // channel
    TransportCallback<StreamResponse> callback = ctx.channel().attr(CALLBACK_ATTR_KEY).getAndRemove();
    if (callback != null) {
        LOG.debug("{}: handling a response", ctx.channel().remoteAddress());
        callback.onResponse(TransportResponseImpl.success(newResponse, wireAttrs));
    } else {
        LOG.debug("{}: dropped a response", ctx.channel().remoteAddress());
    }
}
Also used : StreamResponseBuilder(com.linkedin.r2.message.stream.StreamResponseBuilder) HashMap(java.util.HashMap) StreamResponse(com.linkedin.r2.message.stream.StreamResponse)

Example 3 with StreamResponseBuilder

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

the class RAPResponseDecoder method channelRead0.

@Override
protected void channelRead0(final ChannelHandlerContext ctx, HttpObject msg) throws Exception {
    if (msg instanceof HttpResponse) {
        HttpResponse m = (HttpResponse) msg;
        _shouldCloseConnection = !HttpUtil.isKeepAlive(m);
        if (HttpUtil.is100ContinueExpected(m)) {
            ctx.writeAndFlush(CONTINUE).addListener(new ChannelFutureListener() {

                @Override
                public void operationComplete(ChannelFuture future) throws Exception {
                    if (!future.isSuccess()) {
                        ctx.fireExceptionCaught(future.cause());
                    }
                }
            });
        }
        if (!m.decoderResult().isSuccess()) {
            ctx.fireExceptionCaught(m.decoderResult().cause());
            return;
        }
        // remove chunked encoding.
        if (HttpUtil.isTransferEncodingChunked(m)) {
            HttpUtil.setTransferEncodingChunked(m, false);
        }
        Timeout<None> timeout = ctx.channel().attr(TIMEOUT_ATTR_KEY).getAndRemove();
        if (timeout == null) {
            LOG.debug("dropped a response after channel inactive or exception had happened.");
            return;
        }
        final TimeoutBufferedWriter writer = new TimeoutBufferedWriter(ctx, _maxContentLength, BUFFER_HIGH_WATER_MARK, BUFFER_LOW_WATER_MARK, timeout);
        EntityStream entityStream = EntityStreams.newEntityStream(writer);
        _chunkedMessageWriter = writer;
        StreamResponseBuilder builder = new StreamResponseBuilder();
        builder.setStatus(m.status().code());
        for (Map.Entry<String, String> e : m.headers()) {
            String key = e.getKey();
            String value = e.getValue();
            if (key.equalsIgnoreCase(HttpConstants.RESPONSE_COOKIE_HEADER_NAME)) {
                builder.addCookie(value);
            } else {
                builder.unsafeAddHeaderValue(key, value);
            }
        }
        ctx.fireChannelRead(builder.build(entityStream));
    } else if (msg instanceof HttpContent) {
        HttpContent chunk = (HttpContent) msg;
        TimeoutBufferedWriter currentWriter = _chunkedMessageWriter;
        // Sanity check
        if (currentWriter == null) {
            throw new IllegalStateException("received " + HttpContent.class.getSimpleName() + " without " + HttpResponse.class.getSimpleName());
        }
        if (!chunk.decoderResult().isSuccess()) {
            this.exceptionCaught(ctx, chunk.decoderResult().cause());
        }
        currentWriter.processHttpChunk(chunk);
        if (chunk instanceof LastHttpContent) {
            _chunkedMessageWriter = null;
        }
    } else {
        // something must be wrong, but let's proceed so that
        // handler after us has a chance to process it.
        ctx.fireChannelRead(msg);
    }
}
Also used : ChannelFuture(io.netty.channel.ChannelFuture) StreamResponseBuilder(com.linkedin.r2.message.stream.StreamResponseBuilder) FullHttpResponse(io.netty.handler.codec.http.FullHttpResponse) DefaultFullHttpResponse(io.netty.handler.codec.http.DefaultFullHttpResponse) HttpResponse(io.netty.handler.codec.http.HttpResponse) ByteString(com.linkedin.data.ByteString) LastHttpContent(io.netty.handler.codec.http.LastHttpContent) ChannelFutureListener(io.netty.channel.ChannelFutureListener) TimeoutException(java.util.concurrent.TimeoutException) ClosedChannelException(java.nio.channels.ClosedChannelException) RemoteInvocationException(com.linkedin.r2.RemoteInvocationException) TooLongFrameException(io.netty.handler.codec.TooLongFrameException) IOException(java.io.IOException) EntityStream(com.linkedin.r2.message.stream.entitystream.EntityStream) None(com.linkedin.common.util.None) Map(java.util.Map) LastHttpContent(io.netty.handler.codec.http.LastHttpContent) HttpContent(io.netty.handler.codec.http.HttpContent)

Example 4 with StreamResponseBuilder

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

the class DegraderTrackerClientTest method testCallTrackingStreamRequest.

@Test
public void testCallTrackingStreamRequest() throws Exception {
    URI uri = URI.create("http://test.qa.com:1234/foo");
    SettableClock clock = new SettableClock();
    AtomicInteger action = new AtomicInteger(0);
    TransportClient tc = new TransportClient() {

        @Override
        public void restRequest(RestRequest request, RequestContext requestContext, Map<String, String> wireAttrs, TransportCallback<RestResponse> callback) {
        }

        @Override
        public void streamRequest(StreamRequest request, RequestContext requestContext, Map<String, String> wireAttrs, TransportCallback<StreamResponse> callback) {
            clock.addDuration(5);
            switch(action.get()) {
                // success
                case 0:
                    callback.onResponse(TransportResponseImpl.success(new StreamResponseBuilder().build(EntityStreams.emptyStream())));
                    break;
                // fail with stream exception
                case 1:
                    callback.onResponse(TransportResponseImpl.error(new StreamException(new StreamResponseBuilder().setStatus(500).build(EntityStreams.emptyStream()))));
                    break;
                // fail with timeout exception
                case 2:
                    callback.onResponse(TransportResponseImpl.error(new RemoteInvocationException(new TimeoutException())));
                    break;
                // fail with other exception
                default:
                    callback.onResponse(TransportResponseImpl.error(new RuntimeException()));
                    break;
            }
        }

        @Override
        public void shutdown(Callback<None> callback) {
        }
    };
    DegraderTrackerClientImpl client = (DegraderTrackerClientImpl) createTrackerClient(tc, clock, uri);
    CallTracker callTracker = client.getCallTracker();
    CallTracker.CallStats stats;
    DegraderControl degraderControl = client.getDegraderControl(DefaultPartitionAccessor.DEFAULT_PARTITION_ID);
    DelayConsumeCallback delayConsumeCallback = new DelayConsumeCallback();
    client.streamRequest(new StreamRequestBuilder(uri).build(EntityStreams.emptyStream()), new RequestContext(), new HashMap<>(), delayConsumeCallback);
    clock.addDuration(5);
    // we only recorded the time when stream response arrives, but callcompletion.endcall hasn't been called yet.
    Assert.assertEquals(callTracker.getCurrentCallCountTotal(), 0);
    Assert.assertEquals(callTracker.getCurrentErrorCountTotal(), 0);
    // delay
    clock.addDuration(100);
    delayConsumeCallback.consume();
    clock.addDuration(5000);
    // now that we consumed the entity stream, callcompletion.endcall has been called.
    stats = callTracker.getCallStats();
    Assert.assertEquals(stats.getCallCount(), 1);
    Assert.assertEquals(stats.getErrorCount(), 0);
    Assert.assertEquals(stats.getCallCountTotal(), 1);
    Assert.assertEquals(stats.getErrorCountTotal(), 0);
    Assert.assertEquals(degraderControl.getCurrentComputedDropRate(), 0.0, 0.001);
    action.set(1);
    client.streamRequest(new StreamRequestBuilder(uri).build(EntityStreams.emptyStream()), new RequestContext(), new HashMap<>(), delayConsumeCallback);
    clock.addDuration(5);
    // we endcall with error immediately for stream exception, even before the entity is consumed
    Assert.assertEquals(callTracker.getCurrentCallCountTotal(), 2);
    Assert.assertEquals(callTracker.getCurrentErrorCountTotal(), 1);
    delayConsumeCallback.consume();
    clock.addDuration(5000);
    // no change in tracking after entity is consumed
    stats = callTracker.getCallStats();
    Assert.assertEquals(stats.getCallCount(), 1);
    Assert.assertEquals(stats.getErrorCount(), 1);
    Assert.assertEquals(stats.getCallCountTotal(), 2);
    Assert.assertEquals(stats.getErrorCountTotal(), 1);
    Assert.assertEquals(degraderControl.getCurrentComputedDropRate(), 0.2, 0.001);
    action.set(2);
    client.streamRequest(new StreamRequestBuilder(uri).build(EntityStreams.emptyStream()), new RequestContext(), new HashMap<>(), new TestTransportCallback<>());
    clock.addDuration(5);
    Assert.assertEquals(callTracker.getCurrentCallCountTotal(), 3);
    Assert.assertEquals(callTracker.getCurrentErrorCountTotal(), 2);
    clock.addDuration(5000);
    stats = callTracker.getCallStats();
    Assert.assertEquals(stats.getCallCount(), 1);
    Assert.assertEquals(stats.getErrorCount(), 1);
    Assert.assertEquals(stats.getCallCountTotal(), 3);
    Assert.assertEquals(stats.getErrorCountTotal(), 2);
    Assert.assertEquals(degraderControl.getCurrentComputedDropRate(), 0.4, 0.001);
    action.set(3);
    client.streamRequest(new StreamRequestBuilder(uri).build(EntityStreams.emptyStream()), new RequestContext(), new HashMap<>(), new TestTransportCallback<>());
    clock.addDuration(5);
    Assert.assertEquals(callTracker.getCurrentCallCountTotal(), 4);
    Assert.assertEquals(callTracker.getCurrentErrorCountTotal(), 3);
    clock.addDuration(5000);
    stats = callTracker.getCallStats();
    Assert.assertEquals(stats.getCallCount(), 1);
    Assert.assertEquals(stats.getErrorCount(), 1);
    Assert.assertEquals(stats.getCallCountTotal(), 4);
    Assert.assertEquals(stats.getErrorCountTotal(), 3);
    Assert.assertEquals(degraderControl.getCurrentComputedDropRate(), 0.2, 0.001);
}
Also used : TransportCallback(com.linkedin.r2.transport.common.bridge.common.TransportCallback) TransportClient(com.linkedin.r2.transport.common.bridge.client.TransportClient) StreamResponseBuilder(com.linkedin.r2.message.stream.StreamResponseBuilder) DegraderControl(com.linkedin.util.degrader.DegraderControl) URI(java.net.URI) StreamRequestBuilder(com.linkedin.r2.message.stream.StreamRequestBuilder) StreamRequest(com.linkedin.r2.message.stream.StreamRequest) StreamException(com.linkedin.r2.message.stream.StreamException) RestRequest(com.linkedin.r2.message.rest.RestRequest) Callback(com.linkedin.common.callback.Callback) TransportCallback(com.linkedin.r2.transport.common.bridge.common.TransportCallback) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) SettableClock(com.linkedin.util.clock.SettableClock) RequestContext(com.linkedin.r2.message.RequestContext) RemoteInvocationException(com.linkedin.r2.RemoteInvocationException) HashMap(java.util.HashMap) Map(java.util.Map) TimeoutException(java.util.concurrent.TimeoutException) CallTracker(com.linkedin.util.degrader.CallTracker) Test(org.testng.annotations.Test)

Example 5 with StreamResponseBuilder

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

the class TestReplaceableFilter method testStreamingPassThrough.

@Test
public void testStreamingPassThrough() {
    ReplaceableFilter filter = getFilter();
    CountRequestChunksFilter countRequestChunksFilter = new CountRequestChunksFilter();
    CountResponseChunksFilter countResponseChunksFilter = new CountResponseChunksFilter();
    FilterChain filterChain = FilterChains.empty().addLast(countResponseChunksFilter).addLast(filter).addLast(countRequestChunksFilter);
    int requestChunks = 100;
    int responseChunks = 200;
    StreamRequest streamRequest = new StreamRequestBuilder(URI.create("/test")).build(EntityStreams.newEntityStream(new Writer() {

        private WriteHandle _wh;

        int _count = 0;

        @Override
        public void onInit(WriteHandle wh) {
            _wh = wh;
        }

        @Override
        public void onWritePossible() {
            while (_wh.remaining() > 0) {
                if (_count < requestChunks) {
                    _wh.write(ByteString.copy(new byte[10]));
                    _count++;
                } else {
                    _wh.done();
                }
            }
        }

        @Override
        public void onAbort(Throwable e) {
        }
    }));
    StreamResponse streamResponse = new StreamResponseBuilder().build(EntityStreams.newEntityStream(new Writer() {

        private WriteHandle _wh;

        int _count = 0;

        @Override
        public void onInit(WriteHandle wh) {
            _wh = wh;
        }

        @Override
        public void onWritePossible() {
            while (_wh.remaining() > 0) {
                if (_count < responseChunks) {
                    _wh.write(ByteString.copy(new byte[10]));
                    _count++;
                } else {
                    _wh.done();
                }
            }
        }

        @Override
        public void onAbort(Throwable e) {
        }
    }));
    FilterUtil.fireStreamRequest(filterChain, streamRequest);
    FilterUtil.fireStreamResponse(filterChain, streamResponse, new RequestContext(), new HashMap<>());
    Assert.assertEquals(countRequestChunksFilter.getRequestChunks(), requestChunks);
    Assert.assertEquals(countResponseChunksFilter.getResponseChunks(), responseChunks);
}
Also used : StreamResponseBuilder(com.linkedin.r2.message.stream.StreamResponseBuilder) WriteHandle(com.linkedin.r2.message.stream.entitystream.WriteHandle) ReplaceableFilter(com.linkedin.r2.caprep.ReplaceableFilter) FilterChain(com.linkedin.r2.filter.FilterChain) StreamResponse(com.linkedin.r2.message.stream.StreamResponse) StreamRequestBuilder(com.linkedin.r2.message.stream.StreamRequestBuilder) StreamRequest(com.linkedin.r2.message.stream.StreamRequest) RequestContext(com.linkedin.r2.message.RequestContext) Writer(com.linkedin.r2.message.stream.entitystream.Writer) BaseFilterTest(com.linkedin.r2.testutils.filter.BaseFilterTest) Test(org.testng.annotations.Test)

Aggregations

StreamResponseBuilder (com.linkedin.r2.message.stream.StreamResponseBuilder)26 StreamResponse (com.linkedin.r2.message.stream.StreamResponse)18 TransportCallback (com.linkedin.r2.transport.common.bridge.common.TransportCallback)10 StreamException (com.linkedin.r2.message.stream.StreamException)9 ByteString (com.linkedin.data.ByteString)8 EntityStream (com.linkedin.r2.message.stream.entitystream.EntityStream)8 Map (java.util.Map)8 HashMap (java.util.HashMap)7 Test (org.testng.annotations.Test)7 Callback (com.linkedin.common.callback.Callback)6 RequestContext (com.linkedin.r2.message.RequestContext)6 StreamRequest (com.linkedin.r2.message.stream.StreamRequest)6 RemoteInvocationException (com.linkedin.r2.RemoteInvocationException)4 StreamEncodingType (com.linkedin.r2.filter.compression.streaming.StreamEncodingType)4 StreamingCompressor (com.linkedin.r2.filter.compression.streaming.StreamingCompressor)4 RestRequest (com.linkedin.r2.message.rest.RestRequest)4 StreamRequestBuilder (com.linkedin.r2.message.stream.StreamRequestBuilder)4 TimeoutException (java.util.concurrent.TimeoutException)4 None (com.linkedin.common.util.None)3 RestResponseBuilder (com.linkedin.r2.message.rest.RestResponseBuilder)3