io.netty.util.concurrent.ScheduledFuture Java Examples

void createEventLoop() {
  EventLoop realEventLoop = super.eventLoop();
  if (realEventLoop == null) {
    return;
  }
  eventLoop = mock(EventLoop.class, delegatesTo(realEventLoop));
  doAnswer(
      new Answer<ScheduledFuture<Void>>() {
        @Override
        public ScheduledFuture<Void> answer(InvocationOnMock invocation) throws Throwable {
          Runnable command = (Runnable) invocation.getArguments()[0];
          Long delay = (Long) invocation.getArguments()[1];
          TimeUnit timeUnit = (TimeUnit) invocation.getArguments()[2];
          return new FakeClockScheduledNettyFuture(eventLoop, command, delay, timeUnit);
        }
      }).when(eventLoop).schedule(any(Runnable.class), anyLong(), any(TimeUnit.class));
}
 

private void setSuccess(AddressedEnvelope<? extends DnsResponse, InetSocketAddress> envelope) {
    parent.queryContextManager.remove(nameServerAddr(), id);

    // Cancel the timeout task.
    final ScheduledFuture<?> timeoutFuture = this.timeoutFuture;
    if (timeoutFuture != null) {
        timeoutFuture.cancel(false);
    }

    Promise<AddressedEnvelope<DnsResponse, InetSocketAddress>> promise = this.promise;
    if (promise.setUncancellable()) {
        @SuppressWarnings("unchecked")
        AddressedEnvelope<DnsResponse, InetSocketAddress> castResponse =
                (AddressedEnvelope<DnsResponse, InetSocketAddress>) envelope.retain();
        if (!promise.trySuccess(castResponse)) {
            // We failed to notify the promise as it was failed before, thus we need to release the envelope
            envelope.release();
        }
    }
}
 

void createEventLoop() {
  EventLoop realEventLoop = super.eventLoop();
  if (realEventLoop == null) {
    return;
  }
  eventLoop = mock(EventLoop.class, delegatesTo(realEventLoop));
  doAnswer(
      new Answer<ScheduledFuture<Void>>() {
        @Override
        public ScheduledFuture<Void> answer(InvocationOnMock invocation) throws Throwable {
          Runnable command = (Runnable) invocation.getArguments()[0];
          Long delay = (Long) invocation.getArguments()[1];
          TimeUnit timeUnit = (TimeUnit) invocation.getArguments()[2];
          return new FakeClockScheduledNettyFuture(eventLoop, command, delay, timeUnit);
        }
      }).when(eventLoop).schedule(any(Runnable.class), anyLong(), any(TimeUnit.class));
}
 

@Before
@SuppressWarnings("unchecked")
public void setup() {
  when(ctx.channel()).thenReturn(channel);

  when(channel.remoteAddress()).thenReturn(remoteAddress);
  when(channel.pipeline()).thenReturn(pipeline);
  when(channel.id()).thenReturn(channelId);

  when(channelId.asLongText()).thenReturn("1");
  when(channelId.asShortText()).thenReturn("1");

  when(pipeline.addLast(any())).thenReturn(pipeline);
  when(pipeline.addFirst(any())).thenReturn(pipeline);

  when(channel.eventLoop()).thenReturn(eventLoop);
  when(eventLoop.schedule(any(Callable.class), anyLong(), any()))
      .thenReturn(mock(ScheduledFuture.class));
}
 

@Test
public void doChannelRead_cancels_timeout_check_if_response_finishes_before_timeout_check_occurs() throws Exception {
    // given
    ScheduledFuture timeoutCheckMock = mock(ScheduledFuture.class);
    doReturn(timeoutCheckMock).when(eventLoopMock).schedule(any(Runnable.class), any(Long.class), any(TimeUnit.class));
    handlerSpy.doChannelRead(ctxMock, msg);
    ArgumentCaptor<BiConsumer> timeoutCheckCancellationLogicArgumentCaptor = ArgumentCaptor.forClass(BiConsumer.class);
    // The 2nd whenComplete is for cancelling the timeout check if the response finishes before the timeout
    verify(futureThatWillBeAttachedToSpy, times(2)).whenComplete(timeoutCheckCancellationLogicArgumentCaptor.capture());
    BiConsumer<ResponseInfo<?>, Throwable> timeoutCheckCancellationLogic = timeoutCheckCancellationLogicArgumentCaptor.getAllValues().get(1);

    // when: the timeout check scheduled future is not yet complete when the response finishes
    doReturn(false).when(timeoutCheckMock).isDone();
    timeoutCheckCancellationLogic.accept(mock(ResponseInfo.class), null);

    // then: timeout check scheduled future should be cancelled
    verify(timeoutCheckMock).cancel(false);
}
 

@Test
public void doChannelRead_does_nothing_to_timeout_check_if_timeout_check_is_already_completed_when_response_completes() throws Exception {
    // given
    ScheduledFuture timeoutCheckMock = mock(ScheduledFuture.class);
    doReturn(timeoutCheckMock).when(eventLoopMock).schedule(any(Runnable.class), any(Long.class), any(TimeUnit.class));
    handlerSpy.doChannelRead(ctxMock, msg);
    ArgumentCaptor<BiConsumer> timeoutCheckCancellationLogicArgumentCaptor = ArgumentCaptor.forClass(BiConsumer.class);
    // The 2nd whenComplete is for cancelling the timeout check if the response finishes before the timeout
    verify(futureThatWillBeAttachedToSpy, times(2)).whenComplete(timeoutCheckCancellationLogicArgumentCaptor.capture());
    BiConsumer<ResponseInfo<?>, Throwable> timeoutCheckCancellationLogic = timeoutCheckCancellationLogicArgumentCaptor.getAllValues().get(1);

    // when: the timeout check scheduled future is already done
    doReturn(true).when(timeoutCheckMock).isDone();
    timeoutCheckCancellationLogic.accept(mock(ResponseInfo.class), null);

    // then: nothing should be done
    verify(timeoutCheckMock).isDone();
    verify(timeoutCheckMock, times(0)).cancel(any(Boolean.class));
    verifyNoMoreInteractions(timeoutCheckMock);
}
 

@Override
public void channelActive(ChannelHandlerContext ctx) throws Exception {
    super.channelActive(ctx);
    Channel ch = ctx.channel();
    channel.set(ch);
    log.info("Connection established with endpoint {} on channel {}", connectionName, ch);
    ch.writeAndFlush(new WireCommands.Hello(WireCommands.WIRE_VERSION, WireCommands.OLDEST_COMPATIBLE_VERSION), ch.voidPromise());
    registeredFutureLatch.release(null); //release all futures waiting for channel registration to complete.
    // WireCommands.KeepAlive messages are sent for every network connection to a SegmentStore.
    ScheduledFuture<?> old = keepAliveFuture.getAndSet(ch.eventLoop()
                                                         .scheduleWithFixedDelay(keepAlive,
                                                                                 KEEP_ALIVE_TIMEOUT_SECONDS,
                                                                                 KEEP_ALIVE_TIMEOUT_SECONDS,
                                                                                 TimeUnit.SECONDS));
    if (old != null) {
        old.cancel(false);
    }
}
 

/**
 * Invoke all the {@link ReplyProcessor#connectionDropped()} for all the registered flows once the
 * connection is disconnected.
 *
 * @see io.netty.channel.ChannelInboundHandler#channelUnregistered(ChannelHandlerContext)
 */
@Override
public void channelUnregistered(ChannelHandlerContext ctx) throws Exception {
    ScheduledFuture<?> future = keepAliveFuture.get();
    if (future != null) {
        future.cancel(false);
    }
    channel.set(null);
    log.info("Connection drop observed with endpoint {}", connectionName);
    flowIdReplyProcessorMap.forEach((flowId, rp) -> {
        try {
            log.debug("Connection dropped for flow id {}", flowId);
            rp.connectionDropped();
        } catch (Exception e) {
            // Suppressing exception which prevents all ReplyProcessor.connectionDropped from being invoked.
            log.warn("Encountered exception invoking ReplyProcessor for flow id {}", flowId, e);
        }
    });
    registeredFutureLatch.releaseExceptionally(new ConnectionClosedException());
    super.channelUnregistered(ctx);
}
 

private O limitedExecute(ClientRequestContext ctx, I req) throws Exception {
    final CompletableFuture<O> resFuture = new CompletableFuture<>();
    final O deferred = newDeferredResponse(ctx, resFuture);
    final PendingTask currentTask = new PendingTask(ctx, req, resFuture);

    pendingRequests.add(currentTask);
    drain();

    if (!currentTask.isRun() && timeoutMillis != 0) {
        // Current request was not delegated. Schedule a timeout.
        final ScheduledFuture<?> timeoutFuture = ctx.eventLoop().schedule(
                () -> resFuture.completeExceptionally(
                        UnprocessedRequestException.of(RequestTimeoutException.get())),
                timeoutMillis, TimeUnit.MILLISECONDS);
        currentTask.set(timeoutFuture);
    }

    return deferred;
}
 

/**
 * Schedules next retry.
 */
protected static void scheduleNextRetry(ClientRequestContext ctx,
                                        Consumer<? super Throwable> actionOnException,
                                        Runnable retryTask, long nextDelayMillis) {
    try {
        if (nextDelayMillis == 0) {
            ctx.contextAwareEventLoop().execute(retryTask);
        } else {
            @SuppressWarnings("unchecked")
            final ScheduledFuture<Void> scheduledFuture = (ScheduledFuture<Void>) ctx
                    .contextAwareEventLoop().schedule(retryTask, nextDelayMillis, TimeUnit.MILLISECONDS);
            scheduledFuture.addListener(future -> {
                if (future.isCancelled()) {
                    // future is cancelled when the client factory is closed.
                    actionOnException.accept(new IllegalStateException(
                            ClientFactory.class.getSimpleName() + " has been closed."));
                } else if (future.cause() != null) {
                    // Other unexpected exceptions.
                    actionOnException.accept(future.cause());
                }
            });
        }
    } catch (Throwable t) {
        actionOnException.accept(t);
    }
}
 

protected final <T extends RRemoteServiceResponse> RPromise<T> pollResponse(long timeout,
        RequestId requestId, boolean insertFirst) {
    RPromise<T> responseFuture = new RedissonPromise<T>();

    ResponseEntry entry;
    synchronized (responses) {
        entry = responses.computeIfAbsent(responseQueueName, k -> new ResponseEntry());

        addCancelHandling(requestId, responseFuture);

        ScheduledFuture<?> responseTimeoutFuture = createResponseTimeout(timeout, requestId, responseFuture);

        Map<RequestId, List<Result>> entryResponses = entry.getResponses();
        List<Result> list = entryResponses.computeIfAbsent(requestId, k -> new ArrayList<>(3));

        Result res = new Result(responseFuture, responseTimeoutFuture);
        if (insertFirst) {
            list.add(0, res);
        } else {
            list.add(res);
        }
    }

    pollResponse(entry);
    return responseFuture;
}
 

@Override
public void close() {
  isRunning.set(false);
  for (ScheduledFuture<?> f : scheduledFutures) {
    f.cancel(true);
  }

  // wait for scheduled futures to cancel
  while (scheduledFutures.stream().anyMatch((f) -> !f.isDone())) {
    Uninterruptibles.sleepUninterruptibly(250, TimeUnit.MILLISECONDS);
  }

  // handle remaining items in the queue
  while (counter.get() > 0) {
    drainMessageQ();
  }

  connectionPool.close();
}
 

/**
 * Validate that the channel returned by the underlying channel pool is healthy. If so, complete the result future with the
 * channel returned by the underlying pool. If not, close the channel and try to get a different one.
 *
 * @param delegateFuture A completed promise as a result of invoking delegate.acquire().
 * @param resultFuture The future that should be completed with the healthy, acquired channel.
 * @param timeoutFuture The future for the timeout task. This future will be cancelled when a channel is acquired.
 */
private void ensureAcquiredChannelIsHealthy(Promise<Channel> delegateFuture,
                                            Promise<Channel> resultFuture,
                                            ScheduledFuture<?> timeoutFuture) {
    // If our delegate failed to connect, forward down the failure. Don't try again.
    if (!delegateFuture.isSuccess()) {
        timeoutFuture.cancel(false);
        resultFuture.tryFailure(delegateFuture.cause());
        return;
    }

    // If our delegate gave us an unhealthy connection, close it and try to get a new one.
    Channel channel = delegateFuture.getNow();
    if (!isHealthy(channel)) {
        channel.close();
        delegate.release(channel);
        tryAcquire(resultFuture, timeoutFuture);
        return;
    }

    // Cancel the timeout (best effort), and return back the healthy channel.
    timeoutFuture.cancel(false);
    if (!resultFuture.trySuccess(channel)) {
        // If we couldn't give the channel to the result future (because it failed for some other reason),
        // just return it to the pool.
        release(channel);
    }
}
 

public void updateTimeout() {
    ScheduledFuture<?> key = this.timeoutKey;
    if (key != null) {
        if (!key.cancel(false)) {
            return;
        } else {
            this.timeoutKey = null;
        }
    }
    if (timeout > 0 && !complete) {
        this.timeoutKey = exchange.getIoThread().schedule(timeoutTask, timeout, TimeUnit.MILLISECONDS);
    }
}
 

/**
 * Try to acquire a channel from the underlying pool. This will keep retrying the acquisition until the provided result
 * future is completed.
 *
 * @param resultFuture The future that should be completed with the acquired channel. If this is completed external to this
 * function, this function will stop trying to acquire a channel.
 * @param timeoutFuture The future for the timeout task. This future will be cancelled when a channel is acquired.
 */
private void tryAcquire(Promise<Channel> resultFuture, ScheduledFuture<?> timeoutFuture) {
    // Something else completed the future (probably a timeout). Stop trying to get a channel.
    if (resultFuture.isDone()) {
        return;
    }

    Promise<Channel> delegateFuture = eventLoopGroup.next().newPromise();
    delegate.acquire(delegateFuture);
    delegateFuture.addListener(f -> ensureAcquiredChannelIsHealthy(delegateFuture, resultFuture, timeoutFuture));
}
 

@Override
public Future<Channel> acquire(Promise<Channel> resultFuture) {
    // Schedule a task to time out this acquisition, in case we can't acquire a channel fast enough.
    ScheduledFuture<?> timeoutFuture =
            eventLoopGroup.schedule(() -> timeoutAcquire(resultFuture), acquireTimeoutMillis, TimeUnit.MILLISECONDS);

    tryAcquire(resultFuture, timeoutFuture);
    return resultFuture;
}
 

private void cancelRequestTimeout()
{
    ScheduledFuture<?> timeout = this.timeout.get();
    if (timeout != null) {
        timeout.cancel(false);
    }
}
 

@Override
public PipelineContinuationBehavior doChannelActive(ChannelHandlerContext ctx) throws Exception {
    // New channel opening. See if we have too many open channels.
    int actualOpenChannelsCount = openChannelsGroup.size();
    if (actualOpenChannelsCount >= maxOpenChannelsThreshold) {
        Channel channel = ctx.channel();

        // Mark this channel as needing to be closed.
        ctx.channel().attr(TOO_MANY_OPEN_CONNECTIONS_THIS_CHANNEL_SHOULD_CLOSE).set(actualOpenChannelsCount);

        // Schedule a double-check event to make sure the channel gets closed.
        ScheduledFuture doubleCheckScheduledFuture = ctx.channel().eventLoop().schedule(() -> {
            if (channel.isOpen())
                channel.close();
        }, 100, TimeUnit.MILLISECONDS);

        // Add a channel close future listener to cancel the double-check scheduled event immediately if the channel
        //      is closed quickly. Even though the double-check event will execute in 100 milliseconds that's 100
        //      milliseconds of potential garbage accumulating when it shouldn't. Could be a lot for a high traffic
        //      server (which this likely is if the open channels limit is being hit).
        channel.closeFuture().addListener(future -> {
            if (!doubleCheckScheduledFuture.isDone())
                doubleCheckScheduledFuture.cancel(false);
        });
    }
    else {
        // Not at the threshold. Add this channel to the open channel group.
        openChannelsGroup.add(ctx.channel());
    }

    return PipelineContinuationBehavior.CONTINUE;
}
 

@Before
public void beforeMethod() {
    channelMock = mock(Channel.class);
    ctxMock = mock(ChannelHandlerContext.class);
    tooManyOpenConnectionsAttributeMock = mock(Attribute.class);
    doReturn(channelMock).when(ctxMock).channel();
    doReturn(tooManyOpenConnectionsAttributeMock).when(channelMock)
                                                 .attr(TOO_MANY_OPEN_CONNECTIONS_THIS_CHANNEL_SHOULD_CLOSE);
    doReturn(true).when(channelMock).isOpen();

    eventLoopMock = mock(EventLoop.class);
    closeFutureMock = mock(ChannelFuture.class);
    doReturn(eventLoopMock).when(channelMock).eventLoop();
    doReturn(closeFutureMock).when(channelMock).closeFuture();

    doubleCheckScheduledFutureMock = mock(ScheduledFuture.class);
    doubleCheckRunnableCaptor = ArgumentCaptor.forClass(Runnable.class);
    closeFutureListenerCaptor = ArgumentCaptor.forClass(GenericFutureListener.class);

    doReturn(doubleCheckScheduledFutureMock).when(eventLoopMock)
                                            .schedule(any(Runnable.class), anyLong(), any(TimeUnit.class));
    doReturn(false).when(doubleCheckScheduledFutureMock).isDone();

    channelGroupMock = mock(ChannelGroup.class);
    maxOpenChannelsThreshold = 42;

    handler = new OpenChannelLimitHandler(channelGroupMock, maxOpenChannelsThreshold);
}
 

private <T extends RRemoteServiceResponse> ScheduledFuture<?> createResponseTimeout(long timeout, RequestId requestId, RPromise<T> responseFuture) {
    return commandExecutor.getConnectionManager().getGroup().schedule(new Runnable() {
                @Override
                public void run() {
                    synchronized (responses) {
                        ResponseEntry entry = responses.get(responseQueueName);
                        if (entry == null) {
                            return;
                        }

                        RemoteServiceTimeoutException ex = new RemoteServiceTimeoutException("No response after " + timeout + "ms");
                        if (!responseFuture.tryFailure(ex)) {
                            return;
                        }

                        List<Result> list = entry.getResponses().get(requestId);
                        list.remove(0);
                        if (list.isEmpty()) {
                            entry.getResponses().remove(requestId);
                        }
                        if (entry.getResponses().isEmpty()) {
                            responses.remove(responseQueueName, entry);
                        }
                    }
                }
            }, timeout, TimeUnit.MILLISECONDS);
}
 

@Override
public ScheduledFuture<?> scheduleWithFixedDelay(Runnable command,
		long initialDelay,
		long delay,
		TimeUnit unit) {
	return next().scheduleWithFixedDelay(command, initialDelay, delay, unit);
}
 

@Test
public void testScheduledCancelled() throws Exception {
    EmbeddedChannel ch = new EmbeddedChannel(new ChannelInboundHandlerAdapter());
    ScheduledFuture<?> future = ch.eventLoop().schedule(new Runnable() {
        @Override
        public void run() { }
    }, 1, TimeUnit.DAYS);
    ch.finish();
    Assert.assertTrue(future.isCancelled());
}
 

@Test
public void testScheduledCancelled() throws Exception {
    EmbeddedChannel ch = new EmbeddedChannel(new ChannelInboundHandlerAdapter());
    ScheduledFuture<?> future = ch.eventLoop().schedule(new Runnable() {
        @Override
        public void run() { }
    }, 1, TimeUnit.DAYS);
    ch.finish();
    assertTrue(future.isCancelled());
}
 

@Override
public void run() {
    isRun = true;

    final ScheduledFuture<?> timeoutFuture = get();
    if (timeoutFuture != null) {
        if (timeoutFuture.isDone() || !timeoutFuture.cancel(false)) {
            // Timeout task ran already or is determined to run.
            numActiveRequests.decrementAndGet();
            return;
        }
    }

    try (SafeCloseable ignored = ctx.replace()) {
        try {
            final O actualRes = unwrap().execute(ctx, req);
            actualRes.whenComplete().handleAsync((unused, cause) -> {
                numActiveRequests.decrementAndGet();
                drain();
                return null;
            }, ctx.eventLoop());
            resFuture.complete(actualRes);
        } catch (Throwable t) {
            numActiveRequests.decrementAndGet();
            resFuture.completeExceptionally(t);
        }
    }
}
 

/**
 * Stops polling DNS servers for service updates.
 */
@Override
protected final void doCloseAsync(CompletableFuture<?> future) {
    final ScheduledFuture<?> scheduledFuture = this.scheduledFuture;
    if (scheduledFuture != null) {
        scheduledFuture.cancel(true);
    }
    future.complete(null);
}
 

PendingResponse(HttpMethod method,
                HttpResponseWriter res,
                @Nullable ScheduledFuture<?> pingFuture,
                ScheduledFuture<?> timeoutFuture) {
    this.method = method;
    this.res = res;
    this.pingFuture = pingFuture;
    this.timeoutFuture = timeoutFuture;
}
 

private void configurePipeline(ChannelPipeline p, Set<SessionProtocol> protocols,
                               @Nullable ProxiedAddresses proxiedAddresses) {
    if (protocols.size() == 1) {
        switch (protocols.iterator().next()) {
            case HTTP:
                configureHttp(p, proxiedAddresses);
                break;
            case HTTPS:
                configureHttps(p, proxiedAddresses);
                break;
            default:
                // Should never reach here.
                throw new Error();
        }
        return;
    }

    // More than one protocol were specified. Detect the protocol.

    final ScheduledFuture<?> protocolDetectionTimeoutFuture;
    // FIXME(trustin): Add a dedicated timeout option to ServerConfig.
    final long requestTimeoutMillis = config.defaultVirtualHost().requestTimeoutMillis();
    if (requestTimeoutMillis > 0) {
        // Close the connection if the protocol detection is not finished in time.
        final Channel ch = p.channel();
        protocolDetectionTimeoutFuture = ch.eventLoop().schedule(
                (Runnable) ch::close, requestTimeoutMillis, TimeUnit.MILLISECONDS);
    } else {
        protocolDetectionTimeoutFuture = null;
    }

    p.addLast(new ProtocolDetectionHandler(protocols, proxiedAddresses, protocolDetectionTimeoutFuture));
}
 

ProtocolDetectionHandler(Set<SessionProtocol> protocols, @Nullable ProxiedAddresses proxiedAddresses,
                         @Nullable ScheduledFuture<?> timeoutFuture) {
    candidates = EnumSet.copyOf(protocols);
    if (protocols.contains(PROXY)) {
        proxiedCandidates = EnumSet.copyOf(candidates);
        proxiedCandidates.remove(PROXY);
    } else {
        proxiedCandidates = null;
    }
    this.proxiedAddresses = proxiedAddresses;
    this.timeoutFuture = timeoutFuture;
}
 

@Override
protected void doCloseAsync(CompletableFuture<?> future) {
    closed = true;
    final ScheduledFuture<?> scheduledFuture = this.scheduledFuture;
    if (scheduledFuture != null) {
        scheduledFuture.cancel(true);
    }
    super.doCloseAsync(future);
}
 

@Override
public ScheduledFuture<?> scheduleAtFixedRate(Runnable command,
		long initialDelay,
		long period,
		TimeUnit unit) {
	return next().scheduleAtFixedRate(command, initialDelay, period, unit);
}