Repository: apache/kafka
File: clients/src/main/java/org/apache/kafka/clients/consumer/internals/OffsetFetcher.java (Java)

Code before:
        ListOffsetResult result = new ListOffsetResult();
        if (timestampsToSearch.isEmpty())
            return result;

        // In the case that the user supplied a zero timeout to this method, only a single pass of the loop below will
        // be performed before exiting. No TimeoutException will be thrown in that case. In the case that the single
        // pass did not yield a response (either transient or fatal error), make sure to clear the relevant partitions'
        // respective 'end offset requested' flags so that another attempt can be made by the user.
        //
        // If the timeout is not zero, the loop will be executed at least once. In the case that not all the partitions
        // were found, a TimeoutException will be thrown.
        boolean isZeroTimestamp = timer.timeoutMs() == 0L;

        Map<TopicPartition, Long> remainingToSearch = new HashMap<>(timestampsToSearch);
        do {
            RequestFuture<ListOffsetResult> future = sendListOffsetsRequests(remainingToSearch, requireTimestamps);

            future.addListener(new RequestFutureListener<>() {
                @Override
                public void onSuccess(ListOffsetResult value) {
                    synchronized (future) {
                        result.fetchedOffsets.putAll(value.fetchedOffsets);
                        remainingToSearch.keySet().retainAll(value.partitionsToRetry);

                        offsetFetcherUtils.updateSubscriptionState(value.fetchedOffsets, isolationLevel);
                        offsetFetcherUtils.clearPartitionEndOffsetRequests(remainingToSearch.keySet());
                    }
                }

                @Override
                public void onFailure(RuntimeException e) {
                    offsetFetcherUtils.clearPartitionEndOffsetRequests(remainingToSearch.keySet());

                    if (!(e instanceof RetriableException)) {
                        throw future.exception();
                    }
                }
            });

            // if timeout is set to zero, do not try to poll the network client at all
            // and return empty immediately; otherwise try to get the results synchronously
            // and throw timeout exception if it cannot complete in time
            if (isZeroTimestamp)
                return result;

            client.poll(future, timer);

            if (!future.isDone()) {
                break;
            } else if (remainingToSearch.isEmpty()) {
                return result;

Review comment (bug):
here we're clearing the flag for the partitions that didn't get offsets yet. I agree we need this if we don't have any time left to retry. But if there's still time, the do-while will try again. In that case, do we want to clear the flag here? 

I would imagine we don't, because we'll continue retrying while there is time. It could be the case of missing leader info for instance: we want to keep the flag on for those partitions, hit the `client.awaitMetadataUpdate(timer)` below, and try again in the next iteration of the do-while, right?

If so, I imagine we could take the timer into consideration here? (clear the flag for the failed partitions only if timer expired?). Thoughts?

Code after:
                                                 boolean requireTimestamps,
                                                 boolean updatePartitionEndOffsetsFlag) {
        ListOffsetResult result = new ListOffsetResult();
        if (timestampsToSearch.isEmpty())
            return result;

        Map<TopicPartition, Long> remainingToSearch = new HashMap<>(timestampsToSearch);
        do {
            RequestFuture<ListOffsetResult> future = sendListOffsetsRequests(remainingToSearch, requireTimestamps);

            future.addListener(new RequestFutureListener<>() {
                @Override
                public void onSuccess(ListOffsetResult value) {
                    synchronized (future) {
                        result.fetchedOffsets.putAll(value.fetchedOffsets);
                        remainingToSearch.keySet().retainAll(value.partitionsToRetry);

                        offsetFetcherUtils.updateSubscriptionState(value.fetchedOffsets, isolationLevel);

                        if (updatePartitionEndOffsetsFlag)
                            offsetFetcherUtils.clearPartitionEndOffsetRequests(remainingToSearch.keySet());
                    }
                }

                @Override
                public void onFailure(RuntimeException e) {
                    if (updatePartitionEndOffsetsFlag)
                        offsetFetcherUtils.clearPartitionEndOffsetRequests(remainingToSearch.keySet());

                    if (!(e instanceof RetriableException)) {
                        throw future.exception();
                    }
                }
            });

            // if timeout is set to zero, do not try to poll the network client at all
            // and return empty immediately; otherwise try to get the results synchronously
            // and throw timeout exception if it cannot complete in time
            if (timer.timeoutMs() == 0L)
                return result;

            client.poll(future, timer);

            if (!future.isDone()) {
                break;
            } else if (remainingToSearch.isEmpty()) {
                return result;
            } else {
                client.awaitMetadataUpdate(timer);
            }
        } while (timer.notExpired());