Java Code Examples for org.elasticsearch.action.ActionListener#wrap()

private ActionListener<GetResponse> onGetDetectorResponse(
    MultiResponsesDelegateActionListener<DetectorProfile> listener,
    String detectorId,
    Set<ProfileName> profiles
) {
    return ActionListener.wrap(getResponse -> {
        if (getResponse != null && getResponse.isExists()) {
            DetectorProfile profile = new DetectorProfile();
            if (profiles.contains(ProfileName.STATE)) {
                profile.setState(DetectorState.DISABLED);
            }
            listener.respondImmediately(profile);
        } else {
            listener.failImmediately(FAIL_TO_FIND_DETECTOR_MSG + detectorId);
        }
    }, exception -> { listener.failImmediately(FAIL_TO_FIND_DETECTOR_MSG + detectorId, exception); });
}
 

private ActionListener<ProfileResponse> onModelResponse(
    String detectorId,
    Set<ProfileName> profiles,
    MultiResponsesDelegateActionListener<DetectorProfile> listener
) {
    return ActionListener.wrap(profileResponse -> {
        DetectorProfile profile = new DetectorProfile();
        if (profiles.contains(ProfileName.COORDINATING_NODE)) {
            profile.setCoordinatingNode(profileResponse.getCoordinatingNode());
        }
        if (profiles.contains(ProfileName.SHINGLE_SIZE)) {
            profile.setShingleSize(profileResponse.getShingleSize());
        }
        if (profiles.contains(ProfileName.TOTAL_SIZE_IN_BYTES)) {
            profile.setTotalSizeInBytes(profileResponse.getTotalSizeInBytes());
        }
        if (profiles.contains(ProfileName.MODELS)) {
            profile.setModelProfile(profileResponse.getModelProfile());
        }

        listener.onResponse(profile);
    }, listener::onFailure);
}
 

/**
 * Makes a single attempt to reconnect to any nodes which are disconnected but should be connected. Does not attempt to reconnect any
 * nodes which are in the process of disconnecting. The onCompletion handler is called after all ongoing connection/disconnection
 * attempts have completed.
 */
private void connectDisconnectedTargets(Runnable onCompletion) {
    final List<Runnable> runnables = new ArrayList<>();
    synchronized (mutex) {
        final Collection<ConnectionTarget> connectionTargets = targetsByNode.values();
        if (connectionTargets.isEmpty()) {
            runnables.add(onCompletion);
        } else {
            LOGGER.trace("connectDisconnectedTargets: {}", targetsByNode);
            final GroupedActionListener<Void> listener = new GroupedActionListener<>(
                ActionListener.wrap(onCompletion), connectionTargets.size());
            for (final ConnectionTarget connectionTarget : connectionTargets) {
                runnables.add(connectionTarget.ensureConnected(listener));
            }
        }
    }
    runnables.forEach(Runnable::run);
}
 

private void awaitPendingActivity(Runnable onCompletion) {
    final List<Runnable> runnables = new ArrayList<>();
    synchronized (mutex) {
        final Collection<ConnectionTarget> connectionTargets = targetsByNode.values();
        if (connectionTargets.isEmpty()) {
            runnables.add(onCompletion);
        } else {
            final GroupedActionListener<Void> listener = new GroupedActionListener<>(
                ActionListener.wrap(onCompletion), connectionTargets.size());
            for (final ConnectionTarget connectionTarget : connectionTargets) {
                runnables.add(connectionTarget.awaitCurrentActivity(listener));
            }
        }
    }
    runnables.forEach(Runnable::run);
}
 

private ActionListener<GetResponse> onGetResponse(String adID, ActionListener<Optional<AnomalyDetector>> listener) {
    return ActionListener.wrap(response -> {
        if (response == null || !response.isExists()) {
            listener.onResponse(Optional.empty());
            return;
        }

        String xc = response.getSourceAsString();
        LOG.info("Fetched anomaly detector: {}", xc);

        try (
            XContentParser parser = XContentType.JSON.xContent().createParser(xContentRegistry, LoggingDeprecationHandler.INSTANCE, xc)
        ) {
            ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser::getTokenLocation);
            AnomalyDetector detector = AnomalyDetector.parse(parser, response.getId());
            currentDetectors.put(adID, new SimpleEntry<>(detector, clock.instant()));
            listener.onResponse(Optional.of(detector));
        } catch (Exception t) {
            LOG.error("Fail to parse detector {}", adID);
            LOG.error("Stack trace:", t);
            listener.onResponse(Optional.empty());
        }
    }, listener::onFailure);
}
 

/**
 * Cleans up stale blobs directly under the repository root as well as all indices paths that aren't referenced by any existing
 * snapshots. This method is only to be called directly after a new {@link RepositoryData} was written to the repository and with
 * parameters {@code foundIndices}, {@code rootBlobs}
 *
 * @param foundIndices all indices blob containers found in the repository before {@code newRepoData} was written
 * @param rootBlobs    all blobs found directly under the repository root
 * @param newRepoData  new repository data that was just written
 * @param listener     listener to invoke with the combined long of all blobs removed in this operation
 */
private void cleanupStaleBlobs(Map<String, BlobContainer> foundIndices, Map<String, BlobMetaData> rootBlobs,
                               RepositoryData newRepoData, ActionListener<Long> listener) {
    final GroupedActionListener<Long> groupedListener = new GroupedActionListener<>(ActionListener.wrap(deleteResults -> {
        long deletes = 0;
        for (Long result : deleteResults) {
            deletes += result;
        }
        listener.onResponse(deletes);
    }, listener::onFailure), 2);

    final Executor executor = threadPool.executor(ThreadPool.Names.SNAPSHOT);
    executor.execute(ActionRunnable.supply(groupedListener, () -> {
        List<String> deletedBlobs = cleanupStaleRootFiles(staleRootBlobs(newRepoData, rootBlobs.keySet()));
        return (long) deletedBlobs.size();
    }));

    final Set<String> survivingIndexIds = newRepoData.getIndices().values().stream().map(IndexId::getId).collect(Collectors.toSet());
    executor.execute(ActionRunnable.supply(groupedListener, () -> cleanupStaleIndices(foundIndices, survivingIndexIds)));
}
 

/**
 * Method that handles the actual authentication of the token. This method will only be called if the token is a
 * supported token. The method validates the credentials of the user and if they match, a {@link User} will be
 * returned as the argument to the {@code listener}'s {@link ActionListener#onResponse(Object)} method. Else
 * {@code null} is returned.
 * @param authenticationToken the token to authenticate
 * @param listener return authentication result by calling {@link ActionListener#onResponse(Object)}
 */
@Override
public void authenticate(AuthenticationToken authenticationToken, ActionListener<AuthenticationResult> listener) {
    try {
        UsernamePasswordToken token = (UsernamePasswordToken) authenticationToken;
        UserHolder userHolder = cache.get(token.principal());
        // NOTE the null check for the password. This is done because a cache is shared between authentication and lookup
        // lookup will not store the password...
        if (userHolder == null || userHolder.password == null) {
            super.authenticate(token, ActionListener.wrap(authenticationResult -> {
                if (authenticationResult.isAuthenticated()) {
                    cache.put(token.principal(),
                            new UserHolder(token.credentials().clone().getChars(), authenticationResult.getUser()));
                }
                listener.onResponse(authenticationResult);
            }, listener::onFailure));
        } else if (token.credentials().equals(new SecureString(userHolder.password))) {
            listener.onResponse(AuthenticationResult.success(userHolder.user));
        } else {
            listener.onResponse(AuthenticationResult.notHandled());
        }
    } catch (Exception e) {
        listener.onFailure(e);
    }
}
 

/**
 * Overridden method that will lookup a user from the cache first. If the user is not in the cache, then the super
 * method is called. A non-null result will be cached.
 * @param username the identifier for the user
 * @param listener used to return the user if one was looked or <code>null</code>
 */
@Override
public void lookupUser(String username, ActionListener<User> listener) {
    // a separate cache could be used for lookups to simplify the checking needed in the authenticate method but this
    // requires the lookup cache to also be cleared by the clear cache API
    UserHolder userHolder = cache.get(username);
    if (userHolder != null) {
        listener.onResponse(userHolder.user);
    } else {
        super.lookupUser(username, ActionListener.wrap(lookedUpUser -> {
            if (lookedUpUser != null) {
                UserHolder holder = new UserHolder(null, lookedUpUser);
                cache.put(username, holder);
            }
            listener.onResponse(lookedUpUser);
        }, listener::onFailure));
    }
}
 

void prepareTargetForTranslog(boolean fileBasedRecovery,
                              int totalTranslogOps,
                              ActionListener<TimeValue> listener) {
    StopWatch stopWatch = new StopWatch().start();
    final ActionListener<Void> wrappedListener = ActionListener.wrap(
        nullVal -> {
            stopWatch.stop();
            final TimeValue tookTime = stopWatch.totalTime();
            logger.trace("recovery [phase1]: remote engine start took [{}]", tookTime);
            listener.onResponse(tookTime);
        },
        e -> listener.onFailure(new RecoveryEngineException(shard.shardId(), 1, "prepare target for translog failed", e)));
    // Send a request preparing the new shard's translog to receive operations. This ensures the shard engine is started and disables
    // garbage collection (not the JVM's GC!) of tombstone deletes.
    logger.trace("recovery [phase1]: prepare remote engine for translog");
    cancellableThreads.execute(
        () -> recoveryTarget.prepareForTranslogOperations(
            fileBasedRecovery,
            totalTranslogOps,
            wrappedListener)
    );
}
 

public static <T extends AcknowledgedResponse> ActionListener<T> waitForShards(ActionListener<T> delegate,
                                                                               ActiveShardsObserver observer,
                                                                               TimeValue timeout,
                                                                               Runnable onShardsNotAcknowledged,
                                                                               Supplier<String[]> indices) {

    return ActionListener.wrap(
        resp -> {
            if (resp.isAcknowledged()) {
                observer.waitForActiveShards(
                    indices.get(),
                    ActiveShardCount.DEFAULT,
                    timeout,
                    shardsAcknowledged -> {
                        if (!shardsAcknowledged) {
                            onShardsNotAcknowledged.run();
                        }
                        delegate.onResponse(resp);
                    },
                    delegate::onFailure
                );
            } else {
                delegate.onResponse(resp);
            }
        },
        delegate::onFailure
    );
}
 

@Override
protected void masterOperation(final ClusterRerouteRequest request, final ClusterState state, final ActionListener<ClusterRerouteResponse> listener) {
    ActionListener<ClusterRerouteResponse> logWrapper = ActionListener.wrap(
        response -> {
            if (request.dryRun() == false) {
                response.getExplanations().getYesDecisionMessages().forEach(logger::info);
            }
            listener.onResponse(response);
        },
        listener::onFailure
    );

    clusterService.submitStateUpdateTask("cluster_reroute (api)", new ClusterRerouteResponseAckedClusterStateUpdateTask(logger,
        allocationService, request, logWrapper));
}
 

/**
 * Connect to all the given nodes, but do not disconnect from any extra nodes. Calls the completion handler on completion of all
 * connection attempts to _new_ nodes, but not on attempts to re-establish connections to nodes that are already known.
 */
public void connectToNodes(DiscoveryNodes discoveryNodes, Runnable onCompletion) {

    if (discoveryNodes.getSize() == 0) {
        onCompletion.run();
        return;
    }

    final GroupedActionListener<Void> listener
        = new GroupedActionListener<>(ActionListener.wrap(onCompletion), discoveryNodes.getSize());

    final List<Runnable> runnables = new ArrayList<>(discoveryNodes.getSize());
    synchronized (mutex) {
        for (final DiscoveryNode discoveryNode : discoveryNodes) {
            ConnectionTarget connectionTarget = targetsByNode.get(discoveryNode);
            final boolean isNewNode;
            if (connectionTarget == null) {
                // new node, set up target and listener
                connectionTarget = new ConnectionTarget(discoveryNode);
                targetsByNode.put(discoveryNode, connectionTarget);
                isNewNode = true;
            } else {
                // existing node, but maybe we're disconnecting from it, in which case it was recently removed from the cluster
                // state and has now been re-added so we should wait for the re-connection
                isNewNode = connectionTarget.isPendingDisconnection();
            }

            if (isNewNode) {
                runnables.add(connectionTarget.connect(listener));
            } else {
                // known node, try and ensure it's connected but do not wait
                runnables.add(connectionTarget.connect(null));
                runnables.add(() -> listener.onResponse(null));
            }
        }
    }
    runnables.forEach(Runnable::run);
}
 

private ActionListener<DetectorProfile> getProfileActionListener(RestChannel channel, String detectorId) {
    return ActionListener
        .wrap(
            profile -> { channel.sendResponse(new BytesRestResponse(RestStatus.OK, profile.toXContent(channel.newBuilder()))); },
            exception -> { channel.sendResponse(buildInternalServerErrorResponse(exception, exception.getMessage())); }
        );
}
 

/**
 * Listener sends response once Node Stats and Cluster Stats are gathered
 *
 * @param channel Channel
 * @return ActionListener for ADStatsResponse
 */
private ActionListener<ADStatsResponse> getRestStatsListener(RestChannel channel) {
    return ActionListener
        .wrap(
            adStatsResponse -> {
                channel.sendResponse(new BytesRestResponse(RestStatus.OK, adStatsResponse.toXContent(channel.newBuilder())));
            },
            exception -> channel.sendResponse(new BytesRestResponse(RestStatus.INTERNAL_SERVER_ERROR, exception.getMessage()))
        );
}
 

private ActionListener<SearchResponse> onGetLatestAnomalyResult(ActionListener<DetectorProfile> listener, String detectorId) {
    return ActionListener.wrap(searchResponse -> {
        SearchHits hits = searchResponse.getHits();
        if (hits.getHits().length == 0L) {
            listener.onResponse(new DetectorProfile());
        } else {
            SearchHit hit = hits.getAt(0);

            try (
                XContentParser parser = XContentType.JSON
                    .xContent()
                    .createParser(xContentRegistry, LoggingDeprecationHandler.INSTANCE, hit.getSourceAsString())
            ) {
                ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser::getTokenLocation);
                AnomalyResult result = parser.namedObject(AnomalyResult.class, AnomalyResult.PARSE_FIELD_NAME, null);
                DetectorProfile profile = new DetectorProfile();
                if (result.getError() != null) {
                    profile.setError(result.getError());
                }
                listener.onResponse(profile);

            } catch (IOException | XContentParseException | NullPointerException e) {
                logger.error("Fail to parse anomaly result with " + hit.toString());
                listener.onFailure(new RuntimeException("Fail to find detector error: " + detectorId, e));
            }
        }
    }, exception -> {
        if (exception instanceof IndexNotFoundException) {
            listener.onResponse(new DetectorProfile());
        } else {
            logger.error("Fail to find any anomaly result after AD job enabled time for detector {}", detectorId);
            listener.onFailure(new RuntimeException("Fail to find detector error: " + detectorId, exception));
        }
    });
}
 

public static ActionListener<CreateIndexResponse> createActionListener(
    CheckedConsumer<CreateIndexResponse, ? extends Exception> consumer,
    Consumer<Exception> failureConsumer
) {
    return ActionListener.wrap(consumer, failureConsumer);
}
 

private ActionListener<Optional<AnomalyDetector>> onGetDetector(
    ActionListener<AnomalyResultResponse> listener,
    String adID,
    AnomalyResultRequest request
) {
    return ActionListener.wrap(detector -> {
        if (!detector.isPresent()) {
            listener.onFailure(new EndRunException(adID, "AnomalyDetector is not available.", true));
            return;
        }
        AnomalyDetector anomalyDetector = detector.get();

        String thresholdModelID = modelManager.getThresholdModelId(adID);
        Optional<DiscoveryNode> asThresholdNode = hashRing.getOwningNode(thresholdModelID);
        if (!asThresholdNode.isPresent()) {
            listener.onFailure(new InternalFailure(adID, "Threshold model node is not available."));
            return;
        }

        DiscoveryNode thresholdNode = asThresholdNode.get();

        if (!shouldStart(listener, adID, anomalyDetector, thresholdNode.getId(), thresholdModelID)) {
            return;
        }

        long delayMillis = Optional
            .ofNullable((IntervalTimeConfiguration) anomalyDetector.getWindowDelay())
            .map(t -> t.toDuration().toMillis())
            .orElse(0L);
        long dataStartTime = request.getStart() - delayMillis;
        long dataEndTime = request.getEnd() - delayMillis;

        featureManager
            .getCurrentFeatures(
                anomalyDetector,
                dataStartTime,
                dataEndTime,
                onFeatureResponse(adID, anomalyDetector, listener, thresholdModelID, thresholdNode, dataStartTime, dataEndTime)
            );
    }, exception -> handleExecuteException(exception, listener, adID));

}
 

/**
 * All the exceptions thrown by AD is a subclass of AnomalyDetectionException.
 *  ClientException is a subclass of AnomalyDetectionException. All exception visible to
 *   Client is under ClientVisible. Two classes directly extends ClientException:
 *   - InternalFailure for "root cause unknown failure. Maybe transient." We can continue the
 *    detector running.
 *   - EndRunException for "failures that might impact the customer." The method endNow() is
 *    added to indicate whether the client should immediately terminate running a detector.
 *      + endNow() returns true for "unrecoverable issue". We want to terminate the detector run
 *       immediately.
 *      + endNow() returns false for "maybe unrecoverable issue but worth retrying a few more
 *       times." We want to wait for a few more times on different requests before terminating
 *        the detector run.
 *
 *  AD may not be able to get an anomaly grade but can find a feature vector.  Consider the
 *   case when the shingle is not ready.  In that case, AD just put NaN as anomaly grade and
 *    return the feature vector. If AD cannot even find a feature vector, AD throws
 *     EndRunException if there is an issue or returns empty response (all the numeric fields
 *      are Double.NaN and feature array is empty.  Do so so that customer can write painless
 *       script.) otherwise.
 *
 *  Also, AD is responsible for logging the stack trace.  To avoid bloating our logs, alerting
 *   should always just log the message of an AnomalyDetectionException exception by default.
 *
 *  Known cause of EndRunException with endNow returning false:
 *   + training data for cold start not available
 *   + cold start cannot succeed
 *   + unknown prediction error
 *   + memory circuit breaker tripped
 *
 *  Known cause of EndRunException with endNow returning true:
 *   + a model's memory size reached limit
 *   + models' total memory size reached limit
 *   + Having trouble querying feature data due to
 *    * index does not exist
 *    * all features have been disabled
 *   + anomaly detector is not available
 *   + AD plugin is disabled
 *
 *  Known cause of InternalFailure:
 *   + threshold model node is not available
 *   + cluster read/write is blocked
 *   + cold start hasn't been finished
 *   + fail to get all of rcf model nodes' responses
 *   + fail to get threshold model node's response
 *   + RCF/Threshold model node failing to get checkpoint to restore model before timeout
 *   + Detection is throttle because previous detection query is running
 *
 */
@Override
protected void doExecute(Task task, ActionRequest actionRequest, ActionListener<AnomalyResultResponse> listener) {

    AnomalyResultRequest request = AnomalyResultRequest.fromActionRequest(actionRequest);
    ActionListener<AnomalyResultResponse> original = listener;
    listener = ActionListener.wrap(original::onResponse, e -> {
        adStats.getStat(StatNames.AD_EXECUTE_FAIL_COUNT.getName()).increment();
        original.onFailure(e);
    });

    String adID = request.getAdID();

    if (!EnabledSetting.isADPluginEnabled()) {
        throw new EndRunException(adID, CommonErrorMessages.DISABLED_ERR_MSG, true);
    }

    adStats.getStat(StatNames.AD_EXECUTE_REQUEST_COUNT.getName()).increment();

    if (adCircuitBreakerService.isOpen()) {
        listener.onFailure(new LimitExceededException(adID, CommonErrorMessages.MEMORY_CIRCUIT_BROKEN_ERR_MSG, false));
        return;
    }

    try {
        stateManager.getAnomalyDetector(adID, onGetDetector(listener, adID, request));
    } catch (Exception ex) {
        handleExecuteException(ex, listener, adID);
    }
}