Search job state storage for Data Lake Preview queries (#21525)

* Skeleton for search job state storage

AutoValues instead of records, first, humble methods in the service

Additional fields for SearchJobState

Update methods in the service, with unit tests

Added TIMEOUT to status

Periodical for old jobs removal and expiration.

Expiration is for now only a known concept for Data Lake jobs...

Helper creators for new and done/finished jobs

* Reflecting a naming choice dilemma in the comment for future

* SearchType.Result instead of message list in the SearchJobState object

graylog2-server/src/main/java/org/graylog/plugins/views/ViewsBindings.java

@@ -58,6 +58,7 @@
 import org.graylog.plugins.views.search.filter.QueryStringFilter;
 import org.graylog.plugins.views.search.filter.StreamCategoryFilter;
 import org.graylog.plugins.views.search.filter.StreamFilter;
+import org.graylog.plugins.views.search.jobs.periodical.SearchJobStateCleanupPeriodical;
 import org.graylog.plugins.views.search.querystrings.LastUsedQueryStringsService;
 import org.graylog.plugins.views.search.querystrings.MongoLastUsedQueryStringsService;
 import org.graylog.plugins.views.search.rest.DashboardsResource;
@@ -239,6 +240,7 @@ protected void configure() {
         registerVisualizationConfigSubtypes();
 
         addPeriodical(SearchesCleanUpJob.class);
+        addPeriodical(SearchJobStateCleanupPeriodical.class);
 
         addMigration(V20181220133700_AddViewsAdminRole.class);
         addMigration(V20190304102700_MigrateMessageListStructure.class);

graylog2-server/src/main/java/org/graylog/plugins/views/search/SearchJobIdentifier.java

@@ -17,9 +17,17 @@
 package org.graylog.plugins.views.search;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
+import org.mongojack.Id;
+import org.mongojack.ObjectId;
 
-public record SearchJobIdentifier(@JsonProperty("id") String id,
-                                  @JsonProperty("search_id") String searchId,
-                                  @JsonProperty("owner") String owner,
-                                  @JsonProperty("executing_node") String executingNodeId) {}
+public record SearchJobIdentifier(@JsonProperty("id")
+                                  @ObjectId
+                                  @Id
+                                  String id,
+                                  @JsonProperty(SEARCH_ID_FIELD) String searchId,
+                                  @JsonProperty(OWNER_FIELD) String owner,
+                                  @JsonProperty("executing_node") String executingNodeId) {
+    public static final String SEARCH_ID_FIELD = "search_id";
+    public static final String OWNER_FIELD = "owner";
+}
 

graylog2-server/src/main/java/org/graylog/plugins/views/search/jobs/SearchJobState.java

@@ -0,0 +1,142 @@
+/*
+ * Copyright (C) 2020 Graylog, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the Server Side Public License, version 1,
+ * as published by MongoDB, Inc.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * Server Side Public License for more details.
+ *
+ * You should have received a copy of the Server Side Public License
+ * along with this program. If not, see
+ * <http://www.mongodb.com/licensing/server-side-public-license>.
+ */
+package org.graylog.plugins.views.search.jobs;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.annotation.JsonUnwrapped;
+import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
+import com.google.auto.value.AutoValue;
+import org.graylog.plugins.views.search.SearchJobIdentifier;
+import org.graylog.plugins.views.search.SearchType;
+import org.graylog2.database.MongoEntity;
+import org.joda.time.DateTime;
+import org.joda.time.DateTimeZone;
+import org.mongojack.Id;
+import org.mongojack.ObjectId;
+
+import javax.annotation.Nullable;
+
+@AutoValue
+@JsonDeserialize(builder = SearchJobState.Builder.class)
+public abstract class SearchJobState implements MongoEntity {
+
+    public static final String CREATED_AT_FIELD = "created_at";
+    public static final String UPDATED_AT_FIELD = "updated_at";
+    public static final String STATUS_FIELD = "status";
+    public static final String TYPE_FIELD = "type";
+    public static final String RESULT_FIELD = "result";
+
+    @JsonUnwrapped
+    public abstract SearchJobIdentifier identifier();
+
+    @JsonProperty(STATUS_FIELD)
+    public abstract SearchJobStatus status();
+
+    @JsonProperty(TYPE_FIELD)
+    public abstract SearchJobType type();
+
+    @JsonProperty("error_message")
+    public abstract String errorMessage();
+
+    @JsonProperty("progress")
+    public abstract int progress();
+
+    @JsonProperty(RESULT_FIELD)
+    @Nullable
+    public abstract SearchType.Result result();
+
+    @JsonProperty(CREATED_AT_FIELD)
+    public abstract DateTime createdAt();
+
+    @JsonProperty(UPDATED_AT_FIELD)
+    public abstract DateTime updatedAt();
+
+    public abstract Builder toBuilder();
+
+    public static Builder builder() {
+        return Builder.create();
+    }
+
+    @AutoValue.Builder
+    public abstract static class Builder {
+
+        @JsonUnwrapped
+        public abstract Builder identifier(final SearchJobIdentifier identifier);
+
+        @JsonProperty(STATUS_FIELD)
+        public abstract Builder status(final SearchJobStatus status);
+
+        @JsonProperty(TYPE_FIELD)
+        public abstract Builder type(final SearchJobType type);
+
+        @JsonProperty("error_message")
+        public abstract Builder errorMessage(final String errorMessage);
+
+        @JsonProperty("progress")
+        public abstract Builder progress(final int progress);
+
+        @JsonProperty(RESULT_FIELD)
+        public abstract Builder result(final SearchType.Result result);
+
+        @JsonProperty(CREATED_AT_FIELD)
+        public abstract Builder createdAt(final DateTime createdAt);
+
+        @JsonProperty(UPDATED_AT_FIELD)
+        public abstract Builder updatedAt(final DateTime updatedAt);
+
+        public abstract SearchJobState build();
+
+        @JsonCreator
+        public static Builder create() {
+            return new AutoValue_SearchJobState.Builder()
+                    .progress(0)
+                    .type(SearchJobType.DATA_LAKE)
+                    .errorMessage("");
+        }
+    }
+
+    public static SearchJobState createNewJob(final SearchJobIdentifier searchJobIdentifier) {
+        return SearchJobState.builder()
+                .identifier(searchJobIdentifier)
+                .result(null)
+                .status(SearchJobStatus.RUNNING)
+                .progress(0)
+                .createdAt(DateTime.now(DateTimeZone.UTC))
+                .updatedAt(DateTime.now(DateTimeZone.UTC))
+                .build();
+    }
+
+    public static SearchJobState createDoneJobFrom(final SearchJobState existingSearchJob,
+                                                   final SearchType.Result result) {
+        return existingSearchJob.toBuilder()
+                .result(result)
+                .status(SearchJobStatus.DONE)
+                .progress(100)
+                .updatedAt(DateTime.now(DateTimeZone.UTC))
+                .build();
+    }
+
+    @Override
+    @ObjectId
+    @Id
+    @JsonIgnore
+    public String id() {
+        return identifier().id();
+    }
+}

graylog2-server/src/main/java/org/graylog/plugins/views/search/jobs/SearchJobStateService.java

@@ -0,0 +1,127 @@
+/*
+ * Copyright (C) 2020 Graylog, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the Server Side Public License, version 1,
+ * as published by MongoDB, Inc.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * Server Side Public License for more details.
+ *
+ * You should have received a copy of the Server Side Public License
+ * along with this program. If not, see
+ * <http://www.mongodb.com/licensing/server-side-public-license>.
+ */
+package org.graylog.plugins.views.search.jobs;
+
+import com.mongodb.client.MongoCollection;
+import com.mongodb.client.model.Filters;
+import com.mongodb.client.model.Indexes;
+import com.mongodb.client.model.ReplaceOptions;
+import com.mongodb.client.model.Sorts;
+import com.mongodb.client.model.Updates;
+import com.mongodb.client.result.DeleteResult;
+import com.mongodb.client.result.InsertOneResult;
+import com.mongodb.client.result.UpdateResult;
+import jakarta.inject.Inject;
+import org.graylog2.database.MongoCollections;
+import org.graylog2.database.utils.MongoUtils;
+import org.joda.time.DateTime;
+import org.joda.time.DateTimeZone;
+
+import java.util.Optional;
+
+import static com.mongodb.client.model.Filters.and;
+import static com.mongodb.client.model.Filters.eq;
+import static com.mongodb.client.model.Filters.lte;
+import static org.graylog.plugins.views.search.SearchJobIdentifier.OWNER_FIELD;
+import static org.graylog.plugins.views.search.SearchJobIdentifier.SEARCH_ID_FIELD;
+import static org.graylog.plugins.views.search.jobs.SearchJobState.CREATED_AT_FIELD;
+import static org.graylog.plugins.views.search.jobs.SearchJobState.STATUS_FIELD;
+import static org.graylog.plugins.views.search.jobs.SearchJobState.TYPE_FIELD;
+import static org.graylog.plugins.views.search.jobs.SearchJobType.DATA_LAKE;
+
+public class SearchJobStateService {
+
+    public static final String COLLECTION_NAME = "search_job_states";
+
+    private final MongoCollection<SearchJobState> collection;
+    private final MongoUtils<SearchJobState> mongoUtils;
+
+    @Inject
+    public SearchJobStateService(final MongoCollections mongoCollections) {
+        this.collection = mongoCollections.collection(COLLECTION_NAME, SearchJobState.class);
+        this.mongoUtils = mongoCollections.utils(collection);
+        this.collection.createIndex(Indexes.ascending(SEARCH_ID_FIELD));
+        this.collection.createIndex(Indexes.ascending(OWNER_FIELD));
+        this.collection.createIndex(Indexes.ascending(CREATED_AT_FIELD));
+    }
+
+    public Optional<SearchJobState> get(final String id) {
+        return mongoUtils.getById(id);
+    }
+
+    public Optional<SearchJobState> getLatestForUser(final String user) {
+        return Optional.ofNullable(
+                collection.find(Filters.eq(OWNER_FIELD, user))
+                        .sort(Sorts.descending(CREATED_AT_FIELD))
+                        .first()
+        );
+    }
+
+    public boolean delete(final String id) {
+        return mongoUtils.deleteById(id);
+    }
+
+    public long deleteOlderThan(final DateTime dateTime) {
+        final DeleteResult deleteResult = collection.deleteMany(lte(CREATED_AT_FIELD, dateTime));
+        return deleteResult.getDeletedCount();
+    }
+
+    public long expireOlderThan(final DateTime dateTime) {
+        final UpdateResult updateResult = collection.updateMany(
+                and(
+                        eq(TYPE_FIELD, DATA_LAKE),
+                        //eq(STATUS_FIELD, DONE) //TODO: should all jobs be expired, or only DONE ones?
+                        lte(CREATED_AT_FIELD, dateTime)
+                ),
+                Updates.set(STATUS_FIELD, SearchJobStatus.EXPIRED)
+
+        );
+        return updateResult.getModifiedCount();
+    }
+
+    public SearchJobState create(final SearchJobState searchJobState) {
+        if (searchJobState.identifier().id() != null) {
+            throw new IllegalStateException("ID should be null for a call to create a new SearchJobState");
+        }
+        final InsertOneResult insertOneResult = this.collection.insertOne(searchJobState);
+        return get(MongoUtils.insertedIdAsString(insertOneResult)).orElseThrow(() -> new IllegalStateException("Unable to retrieve saved search job state!"));
+    }
+
+    public boolean update(final SearchJobState searchJobState) {
+        if (searchJobState.identifier().id() == null) {
+            throw new IllegalStateException("Missing ID of SearchJobState to update");
+        }
+        final UpdateResult updateResult = collection.replaceOne(
+                MongoUtils.idEq(searchJobState.identifier().id()),
+                searchJobState.toBuilder().updatedAt(DateTime.now(DateTimeZone.UTC)).build(),
+                new ReplaceOptions().upsert(true)
+        );
+        return updateResult.getModifiedCount() > 0;
+    }
+
+    public boolean changeStatus(final String searchJobStateID,
+                                final SearchJobStatus searchJobStatus) {
+        return get(searchJobStateID)
+                .map(searchJobState -> searchJobState.toBuilder()
+                        .status(searchJobStatus)
+                        .updatedAt(DateTime.now(DateTimeZone.UTC))
+                        .build())
+                .map(this::update)
+                .orElse(false);
+
+    }
+}

graylog2-server/src/main/java/org/graylog/plugins/views/search/jobs/SearchJobStatus.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright (C) 2020 Graylog, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the Server Side Public License, version 1,
+ * as published by MongoDB, Inc.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * Server Side Public License for more details.
+ *
+ * You should have received a copy of the Server Side Public License
+ * along with this program. If not, see
+ * <http://www.mongodb.com/licensing/server-side-public-license>.
+ */
+package org.graylog.plugins.views.search.jobs;
+
+public enum SearchJobStatus {
+    RUNNING,
+    DONE,
+    CANCELLATION_REQUESTED,
+    CANCELLED,
+    TIMEOUT,
+    EXPIRED,
+    ERROR,
+    RESET
+}

graylog2-server/src/main/java/org/graylog/plugins/views/search/jobs/SearchJobType.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright (C) 2020 Graylog, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the Server Side Public License, version 1,
+ * as published by MongoDB, Inc.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * Server Side Public License for more details.
+ *
+ * You should have received a copy of the Server Side Public License
+ * along with this program. If not, see
+ * <http://www.mongodb.com/licensing/server-side-public-license>.
+ */
+package org.graylog.plugins.views.search.jobs;
+
+public enum SearchJobType {
+    DATA_LAKE,
+    //SEARCH_ENGINE or INDEXER - TBD in future
+}