[Rollup] Re-factor Rollup Indexer into a generic indexer for re-usability (#32743)

This extracts a super class out of the rollup indexer called the AsyncTwoPhaseIterator. 
The implementor of it can define the query, transformation of the response, 
indexing and the object to persist the position/state of the indexer.

The stats object used by the indexer to record progress is also now abstract, allowing
the implementation provide custom stats beyond what the indexer provides.  It also
allows the implementation to decide how the stats are presented (leaves toXContent()
up to the implementation).

This should allow new projects to reuse the search-then-index persistent task that Rollup
uses, but without the restrictions/baggage of how Rollup has to work internally to
satisfy time-based rollups.

Author: Hendrik Muhs

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/indexing/AsyncTwoPhaseIndexer.java

@@ -0,0 +1,385 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+package org.elasticsearch.xpack.core.indexing;
+
+import org.apache.log4j.Logger;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.bulk.BulkRequest;
+import org.elasticsearch.action.bulk.BulkResponse;
+import org.elasticsearch.action.index.IndexRequest;
+import org.elasticsearch.action.search.SearchRequest;
+import org.elasticsearch.action.search.SearchResponse;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicReference;
+
+/**
+ * An abstract class that builds an index incrementally. A background job can be launched using {@link #maybeTriggerAsyncJob(long)},
+ * it will create the index from the source index up to the last complete bucket that is allowed to be built (based on job position).
+ * Only one background job can run simultaneously and {@link #onFinish()} is called when the job
+ * finishes. {@link #onFailure(Exception)} is called if the job fails with an exception and {@link #onAbort()} is called if the indexer is
+ * aborted while a job is running. The indexer must be started ({@link #start()} to allow a background job to run when
+ * {@link #maybeTriggerAsyncJob(long)} is called. {@link #stop()} can be used to stop the background job without aborting the indexer.
+ *
+ * In a nutshell this is a 2 cycle engine: 1st it sends a query, 2nd it indexes documents based on the response, sends the next query,
+ * indexes, queries, indexes, ... until a condition lets the engine pause until the source provides new input.
+ *
+ * @param <JobPosition> Type that defines a job position to be defined by the implementation.
+ */
+public abstract class AsyncTwoPhaseIndexer<JobPosition, JobStats extends IndexerJobStats> {
+    private static final Logger logger = Logger.getLogger(AsyncTwoPhaseIndexer.class.getName());
+
+    private final JobStats stats;
+
+    private final AtomicReference<IndexerState> state;
+    private final AtomicReference<JobPosition> position;
+    private final Executor executor;
+
+    protected AsyncTwoPhaseIndexer(Executor executor, AtomicReference<IndexerState> initialState,
+                                   JobPosition initialPosition, JobStats jobStats) {
+        this.executor = executor;
+        this.state = initialState;
+        this.position = new AtomicReference<>(initialPosition);
+        this.stats = jobStats;
+    }
+
+    /**
+     * Get the current state of the indexer.
+     */
+    public IndexerState getState() {
+        return state.get();
+    }
+
+    /**
+     * Get the current position of the indexer.
+     */
+    public JobPosition getPosition() {
+        return position.get();
+    }
+
+    /**
+     * Get the stats of this indexer.
+     */
+    public JobStats getStats() {
+        return stats;
+    }
+
+    /**
+     * Sets the internal state to {@link IndexerState#STARTED} if the previous state
+     * was {@link IndexerState#STOPPED}. Setting the state to STARTED allows a job
+     * to run in the background when {@link #maybeTriggerAsyncJob(long)} is called.
+     *
+     * @return The new state for the indexer (STARTED, INDEXING or ABORTING if the
+     *         job was already aborted).
+     */
+    public synchronized IndexerState start() {
+        state.compareAndSet(IndexerState.STOPPED, IndexerState.STARTED);
+        return state.get();
+    }
+
+    /**
+     * Sets the internal state to {@link IndexerState#STOPPING} if an async job is
+     * running in the background and in such case {@link #onFinish()} will be called
+     * as soon as the background job detects that the indexer is stopped. If there
+     * is no job running when this function is called, the state is directly set to
+     * {@link IndexerState#STOPPED} and {@link #onFinish()} will never be called.
+     *
+     * @return The new state for the indexer (STOPPED, STOPPING or ABORTING if the
+     *         job was already aborted).
+     */
+    public synchronized IndexerState stop() {
+        IndexerState currentState = state.updateAndGet(previousState -> {
+            if (previousState == IndexerState.INDEXING) {
+                return IndexerState.STOPPING;
+            } else if (previousState == IndexerState.STARTED) {
+                return IndexerState.STOPPED;
+            } else {
+                return previousState;
+            }
+        });
+        return currentState;
+    }
+
+    /**
+     * Sets the internal state to {@link IndexerState#ABORTING}. It returns false if
+     * an async job is running in the background and in such case {@link #onAbort}
+     * will be called as soon as the background job detects that the indexer is
+     * aborted. If there is no job running when this function is called, it returns
+     * true and {@link #onAbort()} will never be called.
+     *
+     * @return true if the indexer is aborted, false if a background job is running
+     *         and abort is delayed.
+     */
+    public synchronized boolean abort() {
+        IndexerState prevState = state.getAndUpdate((prev) -> IndexerState.ABORTING);
+        return prevState == IndexerState.STOPPED || prevState == IndexerState.STARTED;
+    }
+
+    /**
+     * Triggers a background job that builds the index asynchronously iff
+     * there is no other job that runs and the indexer is started
+     * ({@link IndexerState#STARTED}.
+     *
+     * @param now
+     *            The current time in milliseconds (used to limit the job to
+     *            complete buckets)
+     * @return true if a job has been triggered, false otherwise
+     */
+    public synchronized boolean maybeTriggerAsyncJob(long now) {
+        final IndexerState currentState = state.get();
+        switch (currentState) {
+        case INDEXING:
+        case STOPPING:
+        case ABORTING:
+            logger.warn("Schedule was triggered for job [" + getJobId() + "], but prior indexer is still running.");
+            return false;
+
+        case STOPPED:
+            logger.debug("Schedule was triggered for job [" + getJobId() + "] but job is stopped.  Ignoring trigger.");
+            return false;
+
+        case STARTED:
+            logger.debug("Schedule was triggered for job [" + getJobId() + "], state: [" + currentState + "]");
+            stats.incrementNumInvocations(1);
+            onStartJob(now);
+
+            if (state.compareAndSet(IndexerState.STARTED, IndexerState.INDEXING)) {
+                // fire off the search. Note this is async, the method will return from here
+                executor.execute(() -> doNextSearch(buildSearchRequest(),
+                        ActionListener.wrap(this::onSearchResponse, exc -> finishWithFailure(exc))));
+                logger.debug("Beginning to index [" + getJobId() + "], state: [" + currentState + "]");
+                return true;
+            } else {
+                logger.debug("Could not move from STARTED to INDEXING state because current state is [" + state.get() + "]");
+                return false;
+            }
+
+        default:
+            logger.warn("Encountered unexpected state [" + currentState + "] while indexing");
+            throw new IllegalStateException("Job encountered an illegal state [" + currentState + "]");
+        }
+    }
+
+    /**
+     * Called to get the Id of the job, used for logging.
+     *
+     * @return a string with the id of the job
+     */
+    protected abstract String getJobId();
+
+    /**
+     * Called to process a response from the 1 search request in order to turn it into a {@link IterationResult}.
+     *
+     * @param searchResponse response from the search phase.
+     * @return Iteration object to be passed to indexing phase.
+     */
+    protected abstract IterationResult<JobPosition> doProcess(SearchResponse searchResponse);
+
+    /**
+     * Called to build the next search request.
+     *
+     * @return SearchRequest to be passed to the search phase.
+     */
+    protected abstract SearchRequest buildSearchRequest();
+
+    /**
+     * Called at startup after job has been triggered using {@link #maybeTriggerAsyncJob(long)} and the
+     * internal state is {@link IndexerState#STARTED}.
+     *
+     * @param now The current time in milliseconds passed through from {@link #maybeTriggerAsyncJob(long)}
+     */
+    protected abstract void onStartJob(long now);
+
+    /**
+     * Executes the {@link SearchRequest} and calls <code>nextPhase</code> with the
+     * response or the exception if an error occurs.
+     *
+     * @param request
+     *            The search request to execute
+     * @param nextPhase
+     *            Listener for the next phase
+     */
+    protected abstract void doNextSearch(SearchRequest request, ActionListener<SearchResponse> nextPhase);
+
+    /**
+     * Executes the {@link BulkRequest} and calls <code>nextPhase</code> with the
+     * response or the exception if an error occurs.
+     *
+     * @param request
+     *            The bulk request to execute
+     * @param nextPhase
+     *            Listener for the next phase
+     */
+    protected abstract void doNextBulk(BulkRequest request, ActionListener<BulkResponse> nextPhase);
+
+    /**
+     * Called periodically during the execution of a background job. Implementation
+     * should persists the state somewhere and continue the execution asynchronously
+     * using <code>next</code>.
+     *
+     * @param state
+     *            The current state of the indexer
+     * @param position
+     *            The current position of the indexer
+     * @param next
+     *            Runnable for the next phase
+     */
+    protected abstract void doSaveState(IndexerState state, JobPosition position, Runnable next);
+
+    /**
+     * Called when a failure occurs in an async job causing the execution to stop.
+     *
+     * @param exc
+     *            The exception
+     */
+    protected abstract void onFailure(Exception exc);
+
+    /**
+     * Called when a background job finishes.
+     */
+    protected abstract void onFinish();
+
+    /**
+     * Called when a background job detects that the indexer is aborted causing the
+     * async execution to stop.
+     */
+    protected abstract void onAbort();
+
+    private void finishWithFailure(Exception exc) {
+        doSaveState(finishAndSetState(), position.get(), () -> onFailure(exc));
+    }
+
+    private IndexerState finishAndSetState() {
+        return state.updateAndGet(prev -> {
+            switch (prev) {
+            case INDEXING:
+                // ready for another job
+                return IndexerState.STARTED;
+
+            case STOPPING:
+                // must be started again
+                return IndexerState.STOPPED;
+
+            case ABORTING:
+                // abort and exit
+                onAbort();
+                return IndexerState.ABORTING; // This shouldn't matter, since onAbort() will kill the task first
+
+            case STOPPED:
+                // No-op. Shouldn't really be possible to get here (should have to go through
+                // STOPPING
+                // first which will be handled) but is harmless to no-op and we don't want to
+                // throw exception here
+                return IndexerState.STOPPED;
+
+            default:
+                // any other state is unanticipated at this point
+                throw new IllegalStateException("Indexer job encountered an illegal state [" + prev + "]");
+            }
+        });
+    }
+
+    private void onSearchResponse(SearchResponse searchResponse) {
+        try {
+            if (checkState(getState()) == false) {
+                return;
+            }
+            if (searchResponse.getShardFailures().length != 0) {
+                throw new RuntimeException("Shard failures encountered while running indexer for job [" + getJobId() + "]: "
+                        + Arrays.toString(searchResponse.getShardFailures()));
+            }
+
+            stats.incrementNumPages(1);
+            IterationResult<JobPosition> iterationResult = doProcess(searchResponse);
+
+            if (iterationResult.isDone()) {
+                logger.debug("Finished indexing for job [" + getJobId() + "], saving state and shutting down.");
+
+                // Change state first, then try to persist. This prevents in-progress
+                // STOPPING/ABORTING from
+                // being persisted as STARTED but then stop the job
+                doSaveState(finishAndSetState(), position.get(), this::onFinish);
+                return;
+            }
+
+            final List<IndexRequest> docs = iterationResult.getToIndex();
+            final BulkRequest bulkRequest = new BulkRequest();
+            docs.forEach(bulkRequest::add);
+
+            // TODO this might be a valid case, e.g. if implementation filters
+            assert bulkRequest.requests().size() > 0;
+
+            doNextBulk(bulkRequest, ActionListener.wrap(bulkResponse -> {
+                // TODO we should check items in the response and move after accordingly to
+                // resume the failing buckets ?
+                if (bulkResponse.hasFailures()) {
+                    logger.warn("Error while attempting to bulk index documents: " + bulkResponse.buildFailureMessage());
+                }
+                stats.incrementNumOutputDocuments(bulkResponse.getItems().length);
+                if (checkState(getState()) == false) {
+                    return;
+                }
+
+                JobPosition newPosition = iterationResult.getPosition();
+                position.set(newPosition);
+
+                onBulkResponse(bulkResponse, newPosition);
+            }, exc -> finishWithFailure(exc)));
+        } catch (Exception e) {
+            finishWithFailure(e);
+        }
+    }
+
+    private void onBulkResponse(BulkResponse response, JobPosition position) {
+        try {
+
+            ActionListener<SearchResponse> listener = ActionListener.wrap(this::onSearchResponse, this::finishWithFailure);
+            // TODO probably something more intelligent than every-50 is needed
+            if (stats.getNumPages() > 0 && stats.getNumPages() % 50 == 0) {
+                doSaveState(IndexerState.INDEXING, position, () -> doNextSearch(buildSearchRequest(), listener));
+            } else {
+                doNextSearch(buildSearchRequest(), listener);
+            }
+        } catch (Exception e) {
+            finishWithFailure(e);
+        }
+    }
+
+    /**
+     * Checks the {@link IndexerState} and returns false if the execution should be
+     * stopped.
+     */
+    private boolean checkState(IndexerState currentState) {
+        switch (currentState) {
+        case INDEXING:
+            // normal state;
+            return true;
+
+        case STOPPING:
+            logger.info("Indexer job encountered [" + IndexerState.STOPPING + "] state, halting indexer.");
+            doSaveState(finishAndSetState(), getPosition(), () -> {
+            });
+            return false;
+
+        case STOPPED:
+            return false;
+
+        case ABORTING:
+            logger.info("Requested shutdown of indexer for job [" + getJobId() + "]");
+            onAbort();
+            return false;
+
+        default:
+            // Anything other than indexing, aborting or stopping is unanticipated
+            logger.warn("Encountered unexpected state [" + currentState + "] while indexing");
+            throw new IllegalStateException("Indexer job encountered an illegal state [" + currentState + "]");
+        }
+    }
+
+}