All Downloads are FREE. Search and download functionalities are using the official Maven repository.

com.vmware.xenon.services.common.QueryTaskService Maven / Gradle / Ivy

There is a newer version: 1.6.18
Show newest version
/*
 * Copyright (c) 2014-2015 VMware, Inc. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License.  You may obtain a copy of
 * the License at http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed
 * under the License is distributed on an "AS IS" BASIS, without warranties or
 * conditions of any kind, EITHER EXPRESS OR IMPLIED.  See the License for the
 * specific language governing permissions and limitations under the License.
 */

package com.vmware.xenon.services.common;

import static com.vmware.xenon.common.ServiceDocumentQueryResult.ContinuousResult;

import java.net.URI;
import java.util.ArrayList;
import java.util.EnumSet;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.TreeMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.function.BiConsumer;
import java.util.stream.Collectors;

import com.vmware.xenon.common.Operation;
import com.vmware.xenon.common.Operation.CompletionHandler;
import com.vmware.xenon.common.ServiceDocument;
import com.vmware.xenon.common.ServiceDocumentDescription;
import com.vmware.xenon.common.ServiceDocumentQueryResult;
import com.vmware.xenon.common.StatefulService;
import com.vmware.xenon.common.TaskState;
import com.vmware.xenon.common.TaskState.TaskStage;
import com.vmware.xenon.common.UriUtils;
import com.vmware.xenon.common.Utils;
import com.vmware.xenon.services.common.ExampleService.ExampleServiceState;
import com.vmware.xenon.services.common.QueryTask.QuerySpecification;
import com.vmware.xenon.services.common.QueryTask.QuerySpecification.QueryOption;
import com.vmware.xenon.services.common.QueryTask.QueryTerm.MatchType;

public class QueryTaskService extends StatefulService {
    private static final long DEFAULT_EXPIRATION_SECONDS = 600;
    private static final Integer DEFAULT_RESULT_LIMIT = Integer.MAX_VALUE;
    private ServiceDocumentQueryResult results;

    public QueryTaskService() {
        super(QueryTask.class);
        super.toggleOption(ServiceOption.REPLICATION, true);
        super.toggleOption(ServiceOption.OWNER_SELECTION, true);
    }

    @Override
    public void handleStart(Operation startPost) {
        if (!startPost.hasBody()) {
            startPost.fail(new IllegalArgumentException("Body is required"));
            return;
        }

        QueryTask initState = startPost.getBody(QueryTask.class);
        if (initState.taskInfo == null) {
            initState.taskInfo = new TaskState();
        } else if (TaskState.isFinished(initState.taskInfo)) {
            startPost.complete();
            return;
        }

        if (!validateState(initState, startPost)) {
            return;
        }

        if (initState.documentExpirationTimeMicros == 0) {
            // always set expiration so we do not accumulate tasks
            initState.documentExpirationTimeMicros = Utils.fromNowMicrosUtc(
                    TimeUnit.SECONDS.toMicros(DEFAULT_EXPIRATION_SECONDS));
        }
        initState.taskInfo.stage = TaskStage.CREATED;

        if (!initState.taskInfo.isDirect) {
            // complete POST immediately
            startPost.setStatusCode(Operation.STATUS_CODE_ACCEPTED).complete();
            // kick off query processing by patching self to STARTED
            QueryTask patchBody = new QueryTask();
            patchBody.taskInfo = new TaskState();
            patchBody.taskInfo.stage = TaskStage.STARTED;
            patchBody.querySpec = initState.querySpec;
            sendRequest(Operation.createPatch(getUri()).setBody(patchBody));
        } else {
            if (initState.querySpec.options.contains(QueryOption.BROADCAST) ||
                    initState.querySpec.options.contains(QueryOption.READ_AFTER_WRITE_CONSISTENCY)) {
                createAndSendBroadcastQuery(initState, startPost);
            } else {
                forwardQueryToDocumentIndexService(initState, startPost);
            }
        }
    }

    private boolean validateState(QueryTask initState, Operation startPost) {
        if (initState.querySpec == null) {
            startPost.fail(new IllegalArgumentException("querySpec is required"));
            return false;
        }

        if (initState.querySpec.query == null) {
            startPost.fail(new IllegalArgumentException("querySpec.query is required"));
            return false;
        }

        if (initState.querySpec.options == null || initState.querySpec.options.isEmpty()) {
            return true;
        }

        if (initState.querySpec.options.contains(QueryOption.EXPAND_CONTENT)) {
            final String errFmt = QueryOption.EXPAND_CONTENT + " is not compatible with %s";
            if (initState.querySpec.options.contains(QueryOption.COUNT)) {
                startPost.fail(new IllegalArgumentException(
                        String.format(errFmt, QueryOption.COUNT)));
                return false;
            }
        }

        if (initState.querySpec.options.contains(QueryOption.EXPAND_BINARY_CONTENT)) {
            final String errFmt = QueryOption.EXPAND_BINARY_CONTENT + " is not compatible with %s";
            if (initState.querySpec.options.contains(QueryOption.COUNT)) {
                startPost.fail(new IllegalArgumentException(
                        String.format(errFmt, QueryOption.COUNT)));
                return false;
            }
        }

        if (startPost.isRemote() && initState.querySpec.options
                .contains(QueryOption.EXPAND_BINARY_CONTENT)) {
            final String errFmt = "%s is not allowed for remote clients.";
            startPost.fail(new IllegalArgumentException(
                    String.format(errFmt, QueryOption.EXPAND_BINARY_CONTENT)));
            return false;
        }

        // EXPAND_BINARY_CONTENT option cannot be used along with OWNER_SELECTION as there is no
        // deserialization happens when we fetch the documents and will not be able to get the
        // owner.
        if (initState.querySpec.options.contains(QueryOption.EXPAND_BINARY_CONTENT)
                && (initState.querySpec.options.contains(QueryOption.OWNER_SELECTION)
                || initState.querySpec.options.contains(QueryOption.EXPAND_BUILTIN_CONTENT_ONLY)
                || initState.querySpec.options.contains(QueryOption.EXPAND_CONTENT))) {
            final String errFmt = "%s is not compatible with %s / %s / %s";
            startPost.fail(new IllegalArgumentException(
                    String.format(errFmt, QueryOption.EXPAND_BINARY_CONTENT, QueryOption.OWNER_SELECTION,
                            QueryOption.EXPAND_BUILTIN_CONTENT_ONLY, QueryOption.EXPAND_CONTENT)));
            return false;
        }

        if (initState.querySpec.options.contains(QueryOption.EXPAND_LINKS)) {
            if (!initState.querySpec.options.contains(QueryOption.SELECT_LINKS)) {
                startPost.fail(new IllegalArgumentException(
                        "Must be combined with " + QueryOption.SELECT_LINKS));
                return false;
            }
            // additional option combination validation will be done in the SELECT_LINKS
            // block, since that option must be combined with this one
        }

        if (initState.querySpec.options.contains(QueryOption.GROUP_BY)) {
            final String errFmt = QueryOption.GROUP_BY + " is not compatible with %s";
            if (initState.querySpec.options.contains(QueryOption.COUNT)) {
                startPost.fail(new IllegalArgumentException(
                        String.format(errFmt, QueryOption.COUNT)));
                return false;
            }
            if (initState.querySpec.options.contains(QueryOption.CONTINUOUS)) {
                startPost.fail(new IllegalArgumentException(
                        String.format(errFmt, QueryOption.CONTINUOUS)));
                return false;
            }
            if (initState.querySpec.groupByTerm == null) {
                startPost.fail(new IllegalArgumentException(
                        "querySpec.groupByTerm is required with " + QueryOption.GROUP_BY));
                return false;
            }
            if (initState.querySpec.sortTerm == null) {
                startPost.fail(new IllegalArgumentException(
                        "querySpec.sortTerm is required with " + QueryOption.GROUP_BY));
                return false;
            }
        }

        if (initState.querySpec.options.contains(QueryOption.SELECT_LINKS)) {
            final String errFmt = QueryOption.SELECT_LINKS + " is not compatible with %s";
            if (initState.querySpec.options.contains(QueryOption.COUNT)) {
                startPost.fail(new IllegalArgumentException(
                        String.format(errFmt, QueryOption.COUNT)));
                return false;
            }
            if (initState.querySpec.options.contains(QueryOption.CONTINUOUS)) {
                startPost.fail(new IllegalArgumentException(
                        String.format(errFmt, QueryOption.CONTINUOUS)));
                return false;
            }
            if (initState.querySpec.linkTerms == null || initState.querySpec.linkTerms.isEmpty()) {
                startPost.fail(new IllegalArgumentException(
                        "querySpec.linkTerms must have at least one entry"));
                return false;
            }
        }

        if (initState.querySpec.options.contains(QueryOption.EXPAND_SELECTED_FIELDS)) {
            final String errFmt = QueryOption.EXPAND_SELECTED_FIELDS + " is not compatible with %s";
            if (initState.querySpec.options.contains(QueryOption.EXPAND_CONTENT)) {
                startPost.fail(new IllegalArgumentException(
                        String.format(errFmt, QueryOption.EXPAND_CONTENT)));
                return false;
            }
            if (initState.querySpec.options.contains(QueryOption.EXPAND_BINARY_CONTENT)) {
                startPost.fail(new IllegalArgumentException(
                        String.format(errFmt, QueryOption.EXPAND_BINARY_CONTENT)));
                return false;
            }
            if (initState.querySpec.selectTerms == null || initState.querySpec.selectTerms.isEmpty()) {
                startPost.fail(new IllegalArgumentException(
                        "querySpec.fieldTerms must have at least one entry"));
                return false;
            }
        }

        if (initState.taskInfo.isDirect
                && initState.querySpec.options.contains(QueryOption.CONTINUOUS)) {
            startPost.fail(new IllegalArgumentException("direct query task is not compatible with "
                    + QueryOption.CONTINUOUS));
            return false;
        }

        if ((initState.querySpec.options.contains(QueryOption.BROADCAST) ||
                initState.querySpec.options.contains(QueryOption.READ_AFTER_WRITE_CONSISTENCY))
                && initState.querySpec.options.contains(QueryOption.SORT)
                && initState.querySpec.sortTerm != null
                && !Objects.equals(initState.querySpec.sortTerm.propertyName, ServiceDocument.FIELD_NAME_SELF_LINK)) {
            startPost.fail(new IllegalArgumentException(QueryOption.BROADCAST
                    + " and " + QueryOption.READ_AFTER_WRITE_CONSISTENCY
                    + " only supports sorting on ["
                    + ServiceDocument.FIELD_NAME_SELF_LINK + "]"));
            return false;
        }

        if (initState.querySpec.options.contains(QueryOption.TIME_SNAPSHOT)
                && initState.querySpec.timeSnapshotBoundaryMicros == null) {
            startPost.fail(new IllegalArgumentException(QueryOption.TIME_SNAPSHOT
                    + " will return latest versions of documents only if querySpec.timeSnapshotBoundaryMicros is provided"));
            return false;
        }

        if (!initState.querySpec.options.contains(QueryOption.TIME_SNAPSHOT)
                && initState.querySpec.timeSnapshotBoundaryMicros != null) {
            startPost.fail(new IllegalArgumentException("Either enable " + QueryOption.TIME_SNAPSHOT
                    + " for retreiving latest versions of documents, for the given querySpec.timeSnapshotBoundaryMicros or do not provide querySpec.timeSnapshotBoundaryMicros"));
            return false;
        }
        if (initState.querySpec.options.contains(QueryOption.READ_AFTER_WRITE_CONSISTENCY)
                && initState.querySpec.options.contains(QueryOption.COUNT)) {
            startPost.fail(new IllegalArgumentException("Options " + QueryOption.READ_AFTER_WRITE_CONSISTENCY
                    + " and " + QueryOption.COUNT + "are not compatible"));
            return false;
        }
        return true;
    }

    private void createAndSendBroadcastQuery(QueryTask origQueryTask, Operation startPost) {
        QueryTask queryTask = Utils.clone(origQueryTask);
        queryTask.setDirect(true);

        if (queryTask.querySpec.options.contains(QueryOption.BROADCAST)) {
            queryTask.querySpec.options.remove(QueryOption.BROADCAST);
        }

        if (queryTask.querySpec.options.contains(QueryOption.READ_AFTER_WRITE_CONSISTENCY)) {
            queryTask.querySpec.options.remove(QueryOption.READ_AFTER_WRITE_CONSISTENCY);
            queryTask.querySpec.options.add(QueryOption.EXPAND_CONTENT);
        }

        if (!queryTask.querySpec.options.contains(QueryOption.SORT)) {
            queryTask.querySpec.options.add(QueryOption.SORT);
            queryTask.querySpec.sortOrder = QuerySpecification.SortOrder.ASC;
            queryTask.querySpec.sortTerm = new QueryTask.QueryTerm();
            queryTask.querySpec.sortTerm.propertyType = ServiceDocumentDescription.TypeName.STRING;
            queryTask.querySpec.sortTerm.propertyName = ServiceDocument.FIELD_NAME_SELF_LINK;
        }

        URI localQueryTaskFactoryUri = UriUtils.buildUri(this.getHost(),
                ServiceUriPaths.CORE_LOCAL_QUERY_TASKS);
        URI forwardingService = UriUtils.buildBroadcastRequestUri(localQueryTaskFactoryUri,
                queryTask.nodeSelectorLink);

        queryTask.documentSelfLink = null;

        Operation op = Operation
                .createPost(forwardingService)
                .setBody(queryTask)
                .setReferer(this.getUri())
                .setConnectionSharing(true)
                .setCompletion((o, e) -> {
                    if (e != null) {
                        failTask(e, startPost, null);
                        return;
                    }

                    NodeGroupBroadcastResponse rsp = o.getBody((NodeGroupBroadcastResponse.class));
                    if (!rsp.failures.isEmpty()) {
                        if (rsp.jsonResponses.size() < rsp.membershipQuorum ||
                                queryTask.querySpec.options.contains(QueryOption.OWNER_SELECTION)) {
                            failTask(new IllegalStateException(
                                            "Broadcast response received failures: " + Utils.toJsonHtml(rsp)),
                                    startPost, null);
                            return;
                        } else {
                            logWarning(
                                    "task will proceed, received %d responses (for quorum size %d)"
                                            + "even though %d errors were received: %s",
                                    rsp.jsonResponses.size(), rsp.membershipQuorum,
                                    rsp.failures.size(), Utils.toJson(rsp.failures));
                        }
                    }

                    collectBroadcastQueryResults(rsp.jsonResponses, origQueryTask, rsp,
                            (response, exception) -> {
                                if (exception != null) {
                                    failTask(new IllegalStateException(
                                            "Failures received: " + Utils.toJsonHtml(exception)),
                                            startPost, null);
                                    return;
                                }
                                queryTask.taskInfo.stage = TaskStage.FINISHED;
                                queryTask.results = response;
                                if (startPost != null) {
                                    // direct query, complete original POST
                                    startPost.setBodyNoCloning(queryTask).complete();
                                } else {
                                    // self patch with results
                                    sendRequest(Operation.createPatch(getUri()).setBodyNoCloning(queryTask));
                                }
                            });
                });
        this.getHost().sendRequest(op);
    }

    private void collectBroadcastQueryResults(Map jsonResponses,
            QueryTask queryTask, NodeGroupBroadcastResponse nodeGroupResponse,
            BiConsumer onCompletion) {
        long startTimeNanos = System.nanoTime();

        List queryResults = new ArrayList<>();
        for (Map.Entry entry : jsonResponses.entrySet()) {
            QueryTask rsp = Utils.fromJson(entry.getValue(), QueryTask.class);
            queryResults.add(rsp.results);
        }

        if (queryResults.size() > 0) {
            long timeElapsed = System.nanoTime() - startTimeNanos;
            timeElapsed /= 1000;
            queryTask.taskInfo.durationMicros = timeElapsed +
                    queryResults.stream().map(r -> r.queryTimeMicros).max(Long::compare).orElse(0L);
        }

        boolean isPaginatedQuery = queryTask.querySpec.resultLimit != null
                && queryTask.querySpec.resultLimit < Integer.MAX_VALUE
                && !queryTask.querySpec.options.contains(QueryOption.TOP_RESULTS);

        if (!isPaginatedQuery) {
            boolean isAscOrder = queryTask.querySpec.sortOrder == null
                    || queryTask.querySpec.sortOrder == QuerySpecification.SortOrder.ASC;
            ServiceDocumentQueryResult result = new ServiceDocumentQueryResult();
            result.queryTimeMicros = queryTask.taskInfo.durationMicros;
            QueryTaskUtils.processQueryResults(getHost(), queryResults, isAscOrder,
                    queryTask.querySpec.options, nodeGroupResponse, result, onCompletion);
        } else {
            URI broadcastPageServiceUri = UriUtils.buildUri(this.getHost(), UriUtils.buildUriPath(
                    ServiceUriPaths.CORE_QUERY_BROADCAST_PAGE, String.valueOf(Utils.getNowMicrosUtc())));

            URI forwarderUri = UriUtils.buildForwardToPeerUri(broadcastPageServiceUri, getHost().getId(),
                    queryTask.nodeSelectorLink != null ? queryTask.nodeSelectorLink : ServiceUriPaths.DEFAULT_NODE_SELECTOR,
                        EnumSet.noneOf(ServiceOption.class));

            ServiceDocument postBody = new ServiceDocument();
            postBody.documentSelfLink = broadcastPageServiceUri.getPath();
            postBody.documentExpirationTimeMicros = queryTask.documentExpirationTimeMicros;

            Operation startPost = Operation
                    .createPost(broadcastPageServiceUri)
                    .setBody(postBody)
                    .setCompletion((o, e) -> {
                        if (e != null) {
                            failTask(e, o, null);
                        }
                    });

            List nextPageLinks = queryResults.stream()
                    .filter(r -> r.nextPageLink != null)
                    .map(r -> r.nextPageLink)
                    .collect(Collectors.toList());

            queryTask.results = new ServiceDocumentQueryResult();
            queryTask.results.documentCount = 0L;
            if (queryTask.querySpec.options.contains(QueryOption.EXPAND_CONTENT) ||
                    queryTask.querySpec.options.contains(QueryOption.EXPAND_SELECTED_FIELDS)) {
                queryTask.results.documents = new HashMap<>();
                queryTask.results.documentLinks = new ArrayList<>();
            }

            if (!nextPageLinks.isEmpty()) {
                queryTask.results.nextPageLink = forwarderUri.getPath() + UriUtils.URI_QUERY_CHAR +
                        forwarderUri.getQuery();
                this.getHost().startService(startPost,
                        // first broadcast query result page will not have prevPageLink back to this query result page
                        new BroadcastQueryPageService(queryTask.querySpec, nextPageLinks,
                                queryTask.documentExpirationTimeMicros, nodeGroupResponse,
                                queryTask.results.nextPageLink, null, null));
            } else {
                queryTask.results.nextPageLink = null;
            }
            onCompletion.accept(queryTask.results, null);
        }
    }

    @Override
    public void handleGet(Operation get) {
        QueryTask currentState = Utils.clone(getState(get));
        ServiceDocumentQueryResult r = this.results;
        if (r == null || currentState == null) {
            get.setBodyNoCloning(currentState).complete();
            return;
        }

        // Infrastructure special case, do not cut and paste in services:
        // the results might contain non clonable JSON serialization artifacts so we go through
        // all these steps to use cached results, avoid cloning, etc This is NOT what services
        // should be doing but due to a unfortunate combination of KRYO and GSON, we cant
        // use results as the body, since it will not clone properly
        currentState.results = new ServiceDocumentQueryResult();
        r.copyTo(currentState.results);

        resetQuerySpecNativeContext(currentState);

        if (r.documentLinks != null) {
            currentState.results.documentLinks = new ArrayList<>(r.documentLinks);
        }
        if (r.documents != null) {
            currentState.results.documents = new HashMap<>(r.documents);
        }
        if (r.selectedLinksPerDocument != null) {
            currentState.results.selectedLinksPerDocument = new HashMap<>(r.selectedLinksPerDocument);
        }
        if (r.selectedLinks != null) {
            currentState.results.selectedLinks = new HashSet<>(r.selectedLinks);
        }
        if (r.selectedDocuments != null) {
            currentState.results.selectedDocuments = new HashMap<>(r.selectedDocuments);
        }
        if (r.nextPageLinksPerGroup != null) {
            currentState.results.nextPageLinksPerGroup = new TreeMap<>(r.nextPageLinksPerGroup);
        }
        if (r.continuousResults != null) {
            ContinuousResult continuousResult = new ContinuousResult();
            continuousResult.documentCountAdded = r.continuousResults.documentCountAdded;
            continuousResult.documentCountUpdated = r.continuousResults.documentCountUpdated;
            continuousResult.documentCountDeleted = r.continuousResults.documentCountDeleted;
            currentState.results.continuousResults = continuousResult;
        }

        get.setBodyNoCloning(currentState).complete();
    }

    private void resetQuerySpecNativeContext(QueryTask currentState) {
        currentState.querySpec.context.nativePage = null;
        currentState.querySpec.context.nativeQuery = null;
        currentState.querySpec.context.nativeSort = null;
        currentState.querySpec.context.nativeSearcher = null;
    }

    @Override
    public void handlePatch(Operation patch) {
        if (patch.isFromReplication()) {
            patch.complete();
            return;
        }

        QueryTask state = getState(patch);

        if (state == null) {
            // service has likely expired
            patch.fail(new IllegalStateException("service state missing"));
            return;
        }

        QueryTask patchBody = patch.getBody(QueryTask.class);
        TaskState newTaskState = patchBody.taskInfo;

        this.results = patchBody.results;

        if (newTaskState == null) {
            patch.fail(new IllegalArgumentException("taskInfo is required"));
            return;
        }

        if (newTaskState.stage == null) {
            patch.fail(new IllegalArgumentException("stage is required"));
            return;
        }

        if (state.querySpec.options.contains(QueryOption.CONTINUOUS)) {
            if (handlePatchForContinuousQuery(state, patchBody, patch)) {
                return;
            }
        }

        if (newTaskState.stage.ordinal() <= state.taskInfo.stage.ordinal()) {
            patch.fail(new IllegalArgumentException(
                    "new stage must be greater than current"));
            return;
        }

        state.taskInfo = newTaskState;
        if (newTaskState.stage == TaskStage.STARTED) {
            patch.setStatusCode(Operation.STATUS_CODE_ACCEPTED);
        } else if (newTaskState.stage == TaskStage.FAILED
                || newTaskState.stage == TaskStage.CANCELLED) {
            if (newTaskState.failure == null) {
                patch.fail(new IllegalArgumentException(
                        "failure must be specified"));
                return;
            }
            logWarning("query failed: %s", newTaskState.failure.message);
        }
        patch.complete();

        if (newTaskState.stage == TaskStage.STARTED) {
            if (patchBody.querySpec.options.contains(QueryOption.BROADCAST) ||
                    patchBody.querySpec.options.contains(QueryOption.READ_AFTER_WRITE_CONSISTENCY)) {
                createAndSendBroadcastQuery(state, null);
            } else {
                forwardQueryToDocumentIndexService(state, null);
            }
        }
    }

    private boolean handlePatchForContinuousQuery(QueryTask state, QueryTask patchBody,
            Operation patch) {
        switch (state.taskInfo.stage) {
        case STARTED:
            // handled below
            break;
        default:
            return false;
        }

        // handle transitions from the STARTED stage
        switch (patchBody.taskInfo.stage) {
        case CREATED:
            return false;
        case STARTED:
            if (patchBody.results.continuousResults == null) {
                patchBody.results.continuousResults = new ContinuousResult();
            }
            // if the new state is STARTED, and we are in STARTED, this is just a update notification
            // from the index that either the initial query completed, or a new update passed the
            // query filter. Subscribers can subscribe to this task and see what changed.
            if (state.results == null) {
                // This would be the first time when the query task has STARTED. Store the
                // count of the documents.
                state.results = patchBody.results;
                if (state.results.documentCount == null) {
                    state.results.documentCount = 0L;
                }
            } else {
                // After it has STARTED, now adjust the count based on
                // documentUpdateAction.
                if (this.results.documents != null) {
                    this.results.documents.values().stream().forEach((doc) -> {
                        ServiceDocument serviceDocument = (ServiceDocument) doc;
                        if (serviceDocument.documentUpdateAction.equals(Action.DELETE.name())) {
                            --state.results.documentCount;
                            ++state.results.continuousResults.documentCountDeleted;
                        } else if (serviceDocument.documentUpdateAction.equals(Action.POST.name())
                                && serviceDocument.documentVersion == 0) {
                            ++state.results.documentCount;
                            ++state.results.continuousResults.documentCountAdded;
                        } else if (serviceDocument.documentUpdateAction.equals(Action.PATCH.name())
                                || serviceDocument.documentUpdateAction.equals(Action.PUT.name())) {
                            ++state.results.continuousResults.documentCountUpdated;
                        }
                    });
                    // Clear from the results documents / documentLinks as only count is requested. Use the
                    // documents array to update the count locally.
                    if (state.querySpec.options.contains(QueryOption.COUNT)) {
                        state.results.documents = null;
                        state.results.documentLinks = null;
                        this.results.documents.clear();
                        this.results.documentLinks.clear();
                        this.results.documentCount = state.results.documentCount;
                        this.results.continuousResults = state.results.continuousResults;
                    }
                    patchBody.results.continuousResults = state.results.continuousResults;
                } else if (this.results.documentCount != null) {
                    state.results.documentCount += this.results.documentCount;
                }
            }
            break;
        case CANCELLED:
        case FAILED:
        case FINISHED:
            cancelContinuousQueryOnIndex(state);
            break;
        default:
            break;
        }

        if (state.querySpec.options.contains(QueryOption.COUNT)) {
            patch.setBodyNoCloning(state).complete();
        } else {
            if (patchBody.results.continuousResults == null) {
                patchBody.results.continuousResults = new ContinuousResult();
            }
            patch.complete();
        }
        return true;
    }

    private void forwardQueryToDocumentIndexService(QueryTask task, Operation directOp) {
        try {
            if (task.querySpec.resultLimit == null) {
                task.querySpec.resultLimit = DEFAULT_RESULT_LIMIT;
            }

            Operation localPatch = Operation
                    .createPatch(this, task.indexLink)
                    .setBodyNoCloning(task)
                    .setCompletion((o, e) -> {
                        if (e == null) {
                            task.results = (ServiceDocumentQueryResult) o.getBodyRaw();
                        }

                        handleQueryCompletion(task, e, directOp);
                    });

            sendRequest(localPatch);
        } catch (Exception e) {
            handleQueryCompletion(task, e, directOp);
        }
    }

    private void scheduleTaskExpiration(QueryTask task) {
        if (task.taskInfo.isDirect) {
            getHost().stopService(this);
            return;
        }

        if (getHost().isStopping()) {
            return;
        }

        Operation delete = Operation.createDelete(getUri()).setBody(new ServiceDocument());
        long delta = task.documentExpirationTimeMicros - Utils.getSystemNowMicrosUtc();
        delta = Math.max(1, delta);
        getHost().scheduleCore(() -> {
            if (task.querySpec.options.contains(QueryOption.CONTINUOUS)) {
                cancelContinuousQueryOnIndex(task);
            }
            sendRequest(delete);
        }, delta, TimeUnit.MICROSECONDS);
    }

    private void cancelContinuousQueryOnIndex(QueryTask task) {
        QueryTask body = new QueryTask();
        body.documentSelfLink = task.documentSelfLink;
        body.taskInfo.stage = TaskStage.CANCELLED;
        body.querySpec = task.querySpec;
        body.documentKind = task.documentKind;
        Operation cancelActiveQueryPatch = Operation
                .createPatch(this, task.indexLink)
                .setBodyNoCloning(body);
        sendRequest(cancelActiveQueryPatch);
    }

    private void failTask(Throwable e, Operation directOp, CompletionHandler c) {
        QueryTask t = new QueryTask();
        // self patch to failure
        t.taskInfo.stage = TaskStage.FAILED;
        t.taskInfo.failure = Utils.toServiceErrorResponse(e);
        if (directOp != null) {
            directOp.setBody(t).fail(e);
            return;
        }

        sendRequest(Operation.createPatch(getUri()).setBody(t).setCompletion(c));
    }

    private boolean handleQueryRetry(QueryTask task, Operation directOp) {
        if (task.querySpec.expectedResultCount == null) {
            return false;
        }

        if (task.results.documentCount >= task.querySpec.expectedResultCount) {
            return false;
        }

        // Fail the task now if we would expire within the next maint interval.
        // Otherwise self patch can fail if the document has expired and clients
        // need a chance to GET the FAILED state.
        long exp = task.documentExpirationTimeMicros - getHost().getMaintenanceIntervalMicros();
        if (exp < Utils.getSystemNowMicrosUtc()) {
            failTask(new TimeoutException(), directOp, (o, e) -> {
                scheduleTaskExpiration(task);
            });
            return true;
        }

        getHost().scheduleCore(() -> {
            forwardQueryToDocumentIndexService(task, directOp);
        }, getMaintenanceIntervalMicros(), TimeUnit.MICROSECONDS);

        return true;
    }

    private void handleQueryCompletion(QueryTask task, Throwable e, Operation directOp) {
        boolean scheduleExpiration = true;

        try {
            task.querySpec.context.nativeQuery = null;

            if (e != null) {
                failTask(e, directOp, null);
                return;
            }

            if (handleQueryRetry(task, directOp)) {
                scheduleExpiration = false;
                return;
            }

            if (task.querySpec.options.contains(QueryOption.CONTINUOUS)) {
                // A continuous query does not cache results: since it receive updates
                // at any time, a GET on the query will cause the query to be re-computed. This is
                // costly, so it should be avoided.
                task.taskInfo.stage = TaskStage.STARTED;
            } else {
                this.results = task.results;
                task.taskInfo.stage = TaskStage.FINISHED;
                task.taskInfo.durationMicros = task.results.queryTimeMicros;
            }

            if (task.documentOwner == null) {
                task.documentOwner = getHost().getId();
            }

            scheduleExpiration = !task.querySpec.options.contains(QueryOption.EXPAND_LINKS);
            if (directOp != null) {
                resetQuerySpecNativeContext(task);
                if (!task.querySpec.options.contains(QueryOption.EXPAND_LINKS)) {
                    directOp.setBodyNoCloning(task).complete();
                    return;
                }
                directOp.nestCompletion((o, ex) -> {
                    directOp.setStatusCode(o.getStatusCode())
                            .setBodyNoCloning(o.getBodyRaw()).complete();
                    scheduleTaskExpiration(task);
                });
                QueryTaskUtils.expandLinks(getHost(), task, directOp);
            } else {
                if (!task.querySpec.options.contains(QueryOption.EXPAND_LINKS)) {
                    sendRequest(Operation.createPatch(getUri()).setBodyNoCloning(task));
                    return;
                }

                CompletionHandler c = (o, ex) -> {
                    scheduleTaskExpiration(task);
                    if (ex != null) {
                        failTask(ex, null, null);
                        return;
                    }
                    sendRequest(Operation.createPatch(getUri()).setBodyNoCloning(task));
                };
                Operation dummyOp = Operation.createGet(getHost().getUri()).setCompletion(c)
                        .setReferer(getUri());
                QueryTaskUtils.expandLinks(getHost(), task, dummyOp);
            }
        } finally {
            if (scheduleExpiration) {
                scheduleTaskExpiration(task);
            }
        }
    }

    @Override
    public ServiceDocument getDocumentTemplate() {
        ServiceDocument td = super.getDocumentTemplate();
        QueryTask template = (QueryTask) td;
        QuerySpecification q = new QueryTask.QuerySpecification();

        QueryTask.Query kindClause = new QueryTask.Query().setTermPropertyName(
                ServiceDocument.FIELD_NAME_KIND).setTermMatchValue(
                Utils.buildKind(ExampleServiceState.class));

        QueryTask.Query nameClause = new QueryTask.Query();
        nameClause.setTermPropertyName("name")
                .setTermMatchValue("query-target")
                .setTermMatchType(MatchType.PHRASE);

        q.query.addBooleanClause(kindClause).addBooleanClause(nameClause);
        template.querySpec = q;

        QueryTask exampleTask = new QueryTask();
        template.indexLink = exampleTask.indexLink;
        return template;
    }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy