@Test public void testEarlyCancel() throws Exception { final ThrottlingClientInterceptor underTest = new ThrottlingClientInterceptor(mockResourceLimiter); ClientCall call = underTest.interceptCall(methodDescriptor, CallOptions.DEFAULT, mockChannel); call.start(mockListener, new Metadata()); call.cancel("fake cancel", null); ArgumentCaptor<Status> statusCaptor = ArgumentCaptor.forClass(Status.class); verify(mockListener, times(1)).onClose(statusCaptor.capture(), any(Metadata.class)); Assert.assertEquals(statusCaptor.getValue().getCode(), Code.CANCELLED); }
@Override public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall( final MethodDescriptor<ReqT, RespT> method, final CallOptions callOptions, final Channel delegateChannel) { if (resourceLimiter == null) { return delegateChannel.newCall(method, callOptions); } return new ClientCall<ReqT, RespT>() { private boolean cancelledEarly; private ClientCall.Listener<RespT> delegateListener = null; private ClientCall<ReqT, RespT> delegateCall; private Metadata headers; private int numMessagesRequested = 0; private Long id = null; @Override public void start(ClientCall.Listener<RespT> listener, Metadata headers) { Preconditions.checkState(!cancelledEarly, "Call already cancelled"); Preconditions.checkState( this.delegateListener == null && this.headers == null, "Call already started"); this.delegateListener = Preconditions.checkNotNull(listener); this.headers = Preconditions.checkNotNull(headers); } @Override public void request(int numMessages) { if (delegateCall != null) { delegateCall.request(numMessages); return; } Preconditions.checkState(!cancelledEarly, "Call already cancelled"); numMessagesRequested += numMessages; } @Override public void cancel(@Nullable String message, @Nullable Throwable cause) { if (delegateCall != null) { delegateCall.cancel(message, cause); return; } cancelledEarly = true; if (message == null && cause == null) { cause = new CancellationException("Cancelled without a message or cause"); } if (delegateListener != null) { delegateListener.onClose( Status.CANCELLED.withDescription(message).withCause(cause), new Metadata()); } } @Override public void sendMessage(ReqT message) { Preconditions.checkState( delegateCall == null, "ThrottlingClientInterceptor only supports unary operations"); Preconditions.checkState( delegateListener != null && headers != null, "start() has to be called before sendMessage()."); Preconditions.checkState(!cancelledEarly, "Call already cancelled"); try { id = resourceLimiter.registerOperationWithHeapSize( ((MessageLite) message).getSerializedSize()); } catch (InterruptedException e) { delegateListener.onClose( Status.INTERNAL.withDescription("Operation was interrupted"), new Metadata()); return; } delegateCall = delegateChannel.newCall(method, callOptions); SimpleForwardingClientCallListener<RespT> markCompletionListener = new SimpleForwardingClientCallListener<RespT>(this.delegateListener) { @Override public void onClose(io.grpc.Status status, Metadata trailers) { resourceLimiter.markCanBeCompleted(id); delegate().onClose(status, trailers); } }; delegateCall.start(markCompletionListener, headers); delegateCall.request(numMessagesRequested); delegateCall.sendMessage(message); delegateCall.halfClose(); } @Override public void halfClose() { } }; }
ThrottlingClientInterceptor implements ClientInterceptor { @Override public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall( final MethodDescriptor<ReqT, RespT> method, final CallOptions callOptions, final Channel delegateChannel) { if (resourceLimiter == null) { return delegateChannel.newCall(method, callOptions); } return new ClientCall<ReqT, RespT>() { private boolean cancelledEarly; private ClientCall.Listener<RespT> delegateListener = null; private ClientCall<ReqT, RespT> delegateCall; private Metadata headers; private int numMessagesRequested = 0; private Long id = null; @Override public void start(ClientCall.Listener<RespT> listener, Metadata headers) { Preconditions.checkState(!cancelledEarly, "Call already cancelled"); Preconditions.checkState( this.delegateListener == null && this.headers == null, "Call already started"); this.delegateListener = Preconditions.checkNotNull(listener); this.headers = Preconditions.checkNotNull(headers); } @Override public void request(int numMessages) { if (delegateCall != null) { delegateCall.request(numMessages); return; } Preconditions.checkState(!cancelledEarly, "Call already cancelled"); numMessagesRequested += numMessages; } @Override public void cancel(@Nullable String message, @Nullable Throwable cause) { if (delegateCall != null) { delegateCall.cancel(message, cause); return; } cancelledEarly = true; if (message == null && cause == null) { cause = new CancellationException("Cancelled without a message or cause"); } if (delegateListener != null) { delegateListener.onClose( Status.CANCELLED.withDescription(message).withCause(cause), new Metadata()); } } @Override public void sendMessage(ReqT message) { Preconditions.checkState( delegateCall == null, "ThrottlingClientInterceptor only supports unary operations"); Preconditions.checkState( delegateListener != null && headers != null, "start() has to be called before sendMessage()."); Preconditions.checkState(!cancelledEarly, "Call already cancelled"); try { id = resourceLimiter.registerOperationWithHeapSize( ((MessageLite) message).getSerializedSize()); } catch (InterruptedException e) { delegateListener.onClose( Status.INTERNAL.withDescription("Operation was interrupted"), new Metadata()); return; } delegateCall = delegateChannel.newCall(method, callOptions); SimpleForwardingClientCallListener<RespT> markCompletionListener = new SimpleForwardingClientCallListener<RespT>(this.delegateListener) { @Override public void onClose(io.grpc.Status status, Metadata trailers) { resourceLimiter.markCanBeCompleted(id); delegate().onClose(status, trailers); } }; delegateCall.start(markCompletionListener, headers); delegateCall.request(numMessagesRequested); delegateCall.sendMessage(message); delegateCall.halfClose(); } @Override public void halfClose() { } }; } }
ThrottlingClientInterceptor implements ClientInterceptor { @Override public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall( final MethodDescriptor<ReqT, RespT> method, final CallOptions callOptions, final Channel delegateChannel) { if (resourceLimiter == null) { return delegateChannel.newCall(method, callOptions); } return new ClientCall<ReqT, RespT>() { private boolean cancelledEarly; private ClientCall.Listener<RespT> delegateListener = null; private ClientCall<ReqT, RespT> delegateCall; private Metadata headers; private int numMessagesRequested = 0; private Long id = null; @Override public void start(ClientCall.Listener<RespT> listener, Metadata headers) { Preconditions.checkState(!cancelledEarly, "Call already cancelled"); Preconditions.checkState( this.delegateListener == null && this.headers == null, "Call already started"); this.delegateListener = Preconditions.checkNotNull(listener); this.headers = Preconditions.checkNotNull(headers); } @Override public void request(int numMessages) { if (delegateCall != null) { delegateCall.request(numMessages); return; } Preconditions.checkState(!cancelledEarly, "Call already cancelled"); numMessagesRequested += numMessages; } @Override public void cancel(@Nullable String message, @Nullable Throwable cause) { if (delegateCall != null) { delegateCall.cancel(message, cause); return; } cancelledEarly = true; if (message == null && cause == null) { cause = new CancellationException("Cancelled without a message or cause"); } if (delegateListener != null) { delegateListener.onClose( Status.CANCELLED.withDescription(message).withCause(cause), new Metadata()); } } @Override public void sendMessage(ReqT message) { Preconditions.checkState( delegateCall == null, "ThrottlingClientInterceptor only supports unary operations"); Preconditions.checkState( delegateListener != null && headers != null, "start() has to be called before sendMessage()."); Preconditions.checkState(!cancelledEarly, "Call already cancelled"); try { id = resourceLimiter.registerOperationWithHeapSize( ((MessageLite) message).getSerializedSize()); } catch (InterruptedException e) { delegateListener.onClose( Status.INTERNAL.withDescription("Operation was interrupted"), new Metadata()); return; } delegateCall = delegateChannel.newCall(method, callOptions); SimpleForwardingClientCallListener<RespT> markCompletionListener = new SimpleForwardingClientCallListener<RespT>(this.delegateListener) { @Override public void onClose(io.grpc.Status status, Metadata trailers) { resourceLimiter.markCanBeCompleted(id); delegate().onClose(status, trailers); } }; delegateCall.start(markCompletionListener, headers); delegateCall.request(numMessagesRequested); delegateCall.sendMessage(message); delegateCall.halfClose(); } @Override public void halfClose() { } }; } ThrottlingClientInterceptor(ResourceLimiter resourceLimiter); }
ThrottlingClientInterceptor implements ClientInterceptor { @Override public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall( final MethodDescriptor<ReqT, RespT> method, final CallOptions callOptions, final Channel delegateChannel) { if (resourceLimiter == null) { return delegateChannel.newCall(method, callOptions); } return new ClientCall<ReqT, RespT>() { private boolean cancelledEarly; private ClientCall.Listener<RespT> delegateListener = null; private ClientCall<ReqT, RespT> delegateCall; private Metadata headers; private int numMessagesRequested = 0; private Long id = null; @Override public void start(ClientCall.Listener<RespT> listener, Metadata headers) { Preconditions.checkState(!cancelledEarly, "Call already cancelled"); Preconditions.checkState( this.delegateListener == null && this.headers == null, "Call already started"); this.delegateListener = Preconditions.checkNotNull(listener); this.headers = Preconditions.checkNotNull(headers); } @Override public void request(int numMessages) { if (delegateCall != null) { delegateCall.request(numMessages); return; } Preconditions.checkState(!cancelledEarly, "Call already cancelled"); numMessagesRequested += numMessages; } @Override public void cancel(@Nullable String message, @Nullable Throwable cause) { if (delegateCall != null) { delegateCall.cancel(message, cause); return; } cancelledEarly = true; if (message == null && cause == null) { cause = new CancellationException("Cancelled without a message or cause"); } if (delegateListener != null) { delegateListener.onClose( Status.CANCELLED.withDescription(message).withCause(cause), new Metadata()); } } @Override public void sendMessage(ReqT message) { Preconditions.checkState( delegateCall == null, "ThrottlingClientInterceptor only supports unary operations"); Preconditions.checkState( delegateListener != null && headers != null, "start() has to be called before sendMessage()."); Preconditions.checkState(!cancelledEarly, "Call already cancelled"); try { id = resourceLimiter.registerOperationWithHeapSize( ((MessageLite) message).getSerializedSize()); } catch (InterruptedException e) { delegateListener.onClose( Status.INTERNAL.withDescription("Operation was interrupted"), new Metadata()); return; } delegateCall = delegateChannel.newCall(method, callOptions); SimpleForwardingClientCallListener<RespT> markCompletionListener = new SimpleForwardingClientCallListener<RespT>(this.delegateListener) { @Override public void onClose(io.grpc.Status status, Metadata trailers) { resourceLimiter.markCanBeCompleted(id); delegate().onClose(status, trailers); } }; delegateCall.start(markCompletionListener, headers); delegateCall.request(numMessagesRequested); delegateCall.sendMessage(message); delegateCall.halfClose(); } @Override public void halfClose() { } }; } ThrottlingClientInterceptor(ResourceLimiter resourceLimiter); @Override ClientCall<ReqT, RespT> interceptCall(
      final MethodDescriptor<ReqT, RespT> method,
      final CallOptions callOptions,
      final Channel delegateChannel); }
ThrottlingClientInterceptor implements ClientInterceptor { @Override public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall( final MethodDescriptor<ReqT, RespT> method, final CallOptions callOptions, final Channel delegateChannel) { if (resourceLimiter == null) { return delegateChannel.newCall(method, callOptions); } return new ClientCall<ReqT, RespT>() { private boolean cancelledEarly; private ClientCall.Listener<RespT> delegateListener = null; private ClientCall<ReqT, RespT> delegateCall; private Metadata headers; private int numMessagesRequested = 0; private Long id = null; @Override public void start(ClientCall.Listener<RespT> listener, Metadata headers) { Preconditions.checkState(!cancelledEarly, "Call already cancelled"); Preconditions.checkState( this.delegateListener == null && this.headers == null, "Call already started"); this.delegateListener = Preconditions.checkNotNull(listener); this.headers = Preconditions.checkNotNull(headers); } @Override public void request(int numMessages) { if (delegateCall != null) { delegateCall.request(numMessages); return; } Preconditions.checkState(!cancelledEarly, "Call already cancelled"); numMessagesRequested += numMessages; } @Override public void cancel(@Nullable String message, @Nullable Throwable cause) { if (delegateCall != null) { delegateCall.cancel(message, cause); return; } cancelledEarly = true; if (message == null && cause == null) { cause = new CancellationException("Cancelled without a message or cause"); } if (delegateListener != null) { delegateListener.onClose( Status.CANCELLED.withDescription(message).withCause(cause), new Metadata()); } } @Override public void sendMessage(ReqT message) { Preconditions.checkState( delegateCall == null, "ThrottlingClientInterceptor only supports unary operations"); Preconditions.checkState( delegateListener != null && headers != null, "start() has to be called before sendMessage()."); Preconditions.checkState(!cancelledEarly, "Call already cancelled"); try { id = resourceLimiter.registerOperationWithHeapSize( ((MessageLite) message).getSerializedSize()); } catch (InterruptedException e) { delegateListener.onClose( Status.INTERNAL.withDescription("Operation was interrupted"), new Metadata()); return; } delegateCall = delegateChannel.newCall(method, callOptions); SimpleForwardingClientCallListener<RespT> markCompletionListener = new SimpleForwardingClientCallListener<RespT>(this.delegateListener) { @Override public void onClose(io.grpc.Status status, Metadata trailers) { resourceLimiter.markCanBeCompleted(id); delegate().onClose(status, trailers); } }; delegateCall.start(markCompletionListener, headers); delegateCall.request(numMessagesRequested); delegateCall.sendMessage(message); delegateCall.halfClose(); } @Override public void halfClose() { } }; } ThrottlingClientInterceptor(ResourceLimiter resourceLimiter); @Override ClientCall<ReqT, RespT> interceptCall(
      final MethodDescriptor<ReqT, RespT> method,
      final CallOptions callOptions,
      final Channel delegateChannel);  }