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);  }
@Test public void testThrottled() throws Exception { final SettableFuture registerInvoked = SettableFuture.create(); final SettableFuture registerWaiter = SettableFuture.create(); final long id = 1l; when(mockResourceLimiter.registerOperationWithHeapSize(anyLong())) .then( new Answer<Long>() { @Override public Long answer(InvocationOnMock invocation) throws Throwable { registerInvoked.set(""); registerWaiter.get(1, TimeUnit.SECONDS); return id; } }); final ThrottlingClientInterceptor underTest = new ThrottlingClientInterceptor(mockResourceLimiter); final int numMessages = 5; Future<?> future = executorService.submit( new Runnable() { @Override public void run() { ClientCall call = underTest.interceptCall(methodDescriptor, CallOptions.DEFAULT, mockChannel); call.start(mockListener, new Metadata()); call.request(numMessages); call.sendMessage(request); } }); registerInvoked.get(1, TimeUnit.SECONDS); verify(mockResourceLimiter, times(1)) .registerOperationWithHeapSize(eq((long) request.getSerializedSize())); verify(mockChannel, times(0)).newCall(any(MethodDescriptor.class), any(CallOptions.class)); registerWaiter.set(""); future.get(1, TimeUnit.SECONDS); verify(mockChannel, times(1)).newCall(any(MethodDescriptor.class), any(CallOptions.class)); verify(mockClientCall, times(1)).request(eq(numMessages)); ArgumentCaptor<ClientCall.Listener> listenerArgumentCaptor = ArgumentCaptor.forClass(ClientCall.Listener.class); verify(mockClientCall, times(1)).start(listenerArgumentCaptor.capture(), any(Metadata.class)); Status status = Status.OK; Metadata trailers = new Metadata(); listenerArgumentCaptor.getValue().onClose(status, trailers); verify(mockResourceLimiter, times(1)).markCanBeCompleted(eq(id)); verify(mockListener, times(1)).onClose(same(status), same(trailers)); }
@Test public void testInterrupted() throws Exception { when(mockResourceLimiter.registerOperationWithHeapSize(anyLong())) .thenThrow(new InterruptedException("Fake interrupted error")); final ThrottlingClientInterceptor underTest = new ThrottlingClientInterceptor(mockResourceLimiter); ClientCall call = underTest.interceptCall(methodDescriptor, CallOptions.DEFAULT, mockChannel); call.start(mockListener, new Metadata()); call.sendMessage(request); call.halfClose(); call.request(1); ArgumentCaptor<Status> statusCaptor = ArgumentCaptor.forClass(Status.class); verify(mockListener, times(1)).onClose(statusCaptor.capture(), any(Metadata.class)); Assert.assertEquals(statusCaptor.getValue().getCode(), Code.INTERNAL); }
@Test public void testCallProxy() { final ThrottlingClientInterceptor underTest = new ThrottlingClientInterceptor(mockResourceLimiter); ClientCall call = underTest.interceptCall(methodDescriptor, CallOptions.DEFAULT, mockChannel); call.start(mockListener, new Metadata()); call.sendMessage(request); call.halfClose(); call.request(1); verify(mockClientCall).start(any(ClientCall.Listener.class), any(Metadata.class)); verify(mockClientCall).sendMessage(request); verify(mockClientCall).halfClose(); verify(mockClientCall).request(1); }
@Test public void testCancel() throws Exception { final ThrottlingClientInterceptor underTest = new ThrottlingClientInterceptor(mockResourceLimiter); ClientCall call = underTest.interceptCall(methodDescriptor, CallOptions.DEFAULT, mockChannel); call.start(mockListener, new Metadata()); call.sendMessage(request); call.cancel("fake cancel", null); verify(mockClientCall).cancel(eq("fake cancel"), (Throwable) any()); }
@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); }