add migration actions with context

graylog2-server/src/main/java/org/graylog/plugins/views/storage/migration/state/actions/MigrationActions.java

@@ -16,10 +16,12 @@
  */
 package org.graylog.plugins.views.storage.migration.state.actions;
 
+import org.graylog.plugins.views.storage.migration.state.machine.MigrationStateMachineContext;
+
 /**
  * Set of callbacks used during the migration in different states.
  */
-public interface MigrationActions extends WithArgs {
+public interface MigrationActions {
     boolean runDirectoryCompatibilityCheck();
 
     boolean isOldClusterStopped();
@@ -47,4 +49,7 @@ public interface MigrationActions extends WithArgs {
 
     void resetMigration();
 
+    void setStateMachineContext(MigrationStateMachineContext context);
+
+    MigrationStateMachineContext getStateMachineContext();
 }

graylog2-server/src/main/java/org/graylog/plugins/views/storage/migration/state/actions/MigrationActionsImpl.java

@@ -17,6 +17,8 @@
 package org.graylog.plugins.views.storage.migration.state.actions;
 
 import jakarta.inject.Inject;
+import jakarta.inject.Singleton;
+import org.graylog.plugins.views.storage.migration.state.machine.MigrationStateMachineContext;
 import org.graylog.plugins.views.storage.migration.state.persistence.DatanodeMigrationConfiguration;
 import org.graylog.security.certutil.CaService;
 import org.graylog.security.certutil.ca.exceptions.KeyStoreStorageException;
@@ -30,11 +32,13 @@
 
 import java.util.Map;
 
+@Singleton
 public class MigrationActionsImpl implements MigrationActions {
 
     private final ClusterConfigService clusterConfigService;
     private final NodeService<DataNodeDto> nodeService;
     private final CaService caService;
+    private MigrationStateMachineContext stateMachineContext;
     private final DataNodeProvisioningService dataNodeProvisioningService;
 
     @Inject
@@ -130,4 +134,15 @@ public void provisionDataNodes() {
     public boolean provisioningFinished() {
         return nodeService.allActive().values().stream().allMatch(node -> node.getDataNodeStatus() == DataNodeStatus.AVAILABLE);
     }
+
+    @Override
+    public void setStateMachineContext(MigrationStateMachineContext context) {
+        this.stateMachineContext = context;
+    }
+
+    @Override
+    public MigrationStateMachineContext getStateMachineContext() {
+        return stateMachineContext;
+    }
+
 }

graylog2-server/src/main/java/org/graylog/plugins/views/storage/migration/state/machine/MigrationStateMachine.java

@@ -16,14 +16,20 @@
  */
 package org.graylog.plugins.views.storage.migration.state.machine;
 
+import org.graylog.plugins.views.storage.migration.state.rest.CurrentStateInformation;
+
 import java.util.List;
 import java.util.Map;
 
 public interface MigrationStateMachine {
-    MigrationState trigger(MigrationStep step, Map<String, Object> args);
+
+    CurrentStateInformation trigger(MigrationStep step, Map<String, Object> args);
+
     MigrationState getState();
 
     List<MigrationStep> nextSteps();
 
+    MigrationStateMachineContext getContext();
+
     String serialize();
 }

graylog2-server/src/main/java/org/graylog/plugins/views/storage/migration/state/machine/MigrationStateMachineBuilder.java

@@ -19,6 +19,7 @@
 import com.github.oxo42.stateless4j.StateMachine;
 import com.github.oxo42.stateless4j.StateMachineConfig;
 import com.github.oxo42.stateless4j.delegates.Trace;
+import com.google.common.annotations.VisibleForTesting;
 import org.graylog.plugins.views.storage.migration.state.actions.MigrationActions;
 import org.graylog.plugins.views.storage.migration.state.persistence.DatanodeMigrationConfiguration;
 import org.graylog.plugins.views.storage.migration.state.persistence.DatanodeMigrationPersistence;
@@ -32,6 +33,11 @@ public class MigrationStateMachineBuilder {
 
     @NotNull
     private static StateMachineConfig<MigrationState, MigrationStep> configureStates(MigrationActions migrationActions) {
+
+        // All actions which can fail should be performed on transition to make sure that there is no state change on error.
+        // For async tasks, the task should be triggered in the transition with a following intermediary step which
+        // has a guard to continue only on task completion.
+
         StateMachineConfig<MigrationState, MigrationStep> config = new StateMachineConfig<>();
 
         config.configure(MigrationState.NEW)
@@ -54,11 +60,12 @@ private static StateMachineConfig<MigrationState, MigrationStep> configureStates
 
         // Major decision - remote reindexing or rolling upgrade(in-place)?
         config.configure(MigrationState.MIGRATION_SELECTION_PAGE)
-                .permit(MigrationStep.SELECT_ROLLING_UPGRADE_MIGRATION, MigrationState.ROLLING_UPGRADE_MIGRATION_WELCOME_PAGE, migrationActions::rollingUpgradeSelected)
-                .permit(MigrationStep.SELECT_REMOTE_REINDEX_MIGRATION, MigrationState.REMOTE_REINDEX_WELCOME_PAGE, migrationActions::reindexUpgradeSelected);
+                .permit(MigrationStep.SELECT_ROLLING_UPGRADE_MIGRATION, MigrationState.ROLLING_UPGRADE_MIGRATION_WELCOME_PAGE)
+                .permit(MigrationStep.SELECT_REMOTE_REINDEX_MIGRATION, MigrationState.REMOTE_REINDEX_WELCOME_PAGE);
 
         // remote reindexing branch of the migration
         config.configure(MigrationState.REMOTE_REINDEX_WELCOME_PAGE)
+                .onEntry(migrationActions::reindexUpgradeSelected)
                 .permit(MigrationStep.DISCOVER_NEW_DATANODES, MigrationState.PROVISION_DATANODE_CERTIFICATES_PAGE, () -> {
                     LOG.info("Remote Reindexing selected");
                 });
@@ -79,10 +86,11 @@ private static StateMachineConfig<MigrationState, MigrationStep> configureStates
 
         // in place / rolling upgrade branch of the migration
         config.configure(MigrationState.ROLLING_UPGRADE_MIGRATION_WELCOME_PAGE)
+                .onEntry(migrationActions::rollingUpgradeSelected)
                 .permit(MigrationStep.INSTALL_DATANODES_ON_EVERY_NODE, MigrationState.DIRECTORY_COMPATIBILITY_CHECK_PAGE2, () -> LOG.info("Showing directory compatibility check page"));
 
         config.configure(MigrationState.DIRECTORY_COMPATIBILITY_CHECK_PAGE2)
-                .permit(MigrationStep.SHOW_PROVISION_ROLLING_UPGRADE_NODES_WITH_CERTIFICATES, MigrationState.PROVISION_ROLLING_UPGRADE_NODES_WITH_CERTIFICATES, migrationActions::directoryCompatibilityCheckOk);
+                .permitIf(MigrationStep.SHOW_PROVISION_ROLLING_UPGRADE_NODES_WITH_CERTIFICATES, MigrationState.PROVISION_ROLLING_UPGRADE_NODES_WITH_CERTIFICATES, migrationActions::directoryCompatibilityCheckOk);
 
         config.configure(MigrationState.PROVISION_ROLLING_UPGRADE_NODES_WITH_CERTIFICATES)
                 .permit(MigrationStep.CALCULATE_JOURNAL_SIZE, MigrationState.JOURNAL_SIZE_DOWNTIME_WARNING, migrationActions::provisionDataNodes);
@@ -105,9 +113,16 @@ private static StateMachineConfig<MigrationState, MigrationStep> configureStates
         return config;
     }
 
+    @VisibleForTesting
+    static StateMachine<MigrationState, MigrationStep> buildWithTestState(MigrationState state, MigrationActions migrationActions) {
+        final StateMachineConfig<MigrationState, MigrationStep> config = configureStates(migrationActions);
+        return new StateMachine<>(state, config);
+    }
+
     private static StateMachine<MigrationState, MigrationStep> fromState(MigrationState state, Trace<MigrationState, MigrationStep> persistence, MigrationActions migrationActions) {
         final StateMachineConfig<MigrationState, MigrationStep> config = configureStates(migrationActions);
         final StateMachine<MigrationState, MigrationStep> stateMachine = new StateMachine<>(state, config);
+        stateMachine.fireInitialTransition(); // asserts that onEntry will be performed when loading persisted state
         stateMachine.setTrace(persistence);
         return stateMachine;
     }

graylog2-server/src/main/java/org/graylog/plugins/views/storage/migration/state/machine/MigrationStateMachineContext.java

@@ -0,0 +1,72 @@
+/*
+ * 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.storage.migration.state.machine;
+
+import com.fasterxml.jackson.annotation.JsonAutoDetect;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+
+@JsonAutoDetect(fieldVisibility = JsonAutoDetect.Visibility.PROTECTED_AND_PUBLIC)
+public class MigrationStateMachineContext {
+
+    @JsonProperty
+    protected MigrationStep currentStep;
+    @JsonProperty
+    protected Map<MigrationStep, Map<String, Object>> actionArguments;
+    @JsonProperty
+    protected Map<String, Object> extendedState;
+
+    public MigrationStateMachineContext() {
+        this.actionArguments = new HashMap<>();
+        this.extendedState = new HashMap<>();
+    }
+
+    public void setCurrentStep(MigrationStep currentStep) {
+        this.currentStep = currentStep;
+    }
+
+    public <T> T getActionArgument(String name, Class<T> type) {
+        Map<String, Object> args = this.actionArguments.get(currentStep);
+        if (Objects.isNull(args)) {
+            throw new IllegalArgumentException("Missing arguments for step " + currentStep);
+        }
+        if (!args.containsKey(name)) {
+            throw new IllegalArgumentException("Missing argument " + name + " for step " + currentStep);
+        }
+        Object arg = args.get(name);
+        if (!type.isInstance(arg)) {
+            throw new IllegalArgumentException("Argument " + name + " must be of type " + type);
+        }
+        return (T) arg;
+    }
+
+    public void addActionArguments(MigrationStep step, Map<String, Object> args) {
+        this.actionArguments.put(step, args);
+    }
+
+    public void addExtendedState(String key, Object value) {
+        this.extendedState.put(key, value);
+    }
+
+    public Object getExtendedState(String key) {
+        return this.extendedState.get(key);
+    }
+
+}

graylog2-server/src/main/java/org/graylog/plugins/views/storage/migration/state/machine/MigrationStateMachineImpl.java

@@ -18,31 +18,50 @@
 
 import com.github.oxo42.stateless4j.StateMachine;
 import org.graylog.plugins.views.storage.migration.state.actions.MigrationActions;
+import org.graylog.plugins.views.storage.migration.state.persistence.DatanodeMigrationPersistence;
+import org.graylog.plugins.views.storage.migration.state.rest.CurrentStateInformation;
 
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 
 public class MigrationStateMachineImpl implements MigrationStateMachine {
     private final StateMachine<MigrationState, MigrationStep> stateMachine;
     private final MigrationActions migrationActions;
+    private final DatanodeMigrationPersistence persistenceService;
+    private MigrationStateMachineContext context;
 
-    public MigrationStateMachineImpl(StateMachine<MigrationState, MigrationStep> stateMachine, MigrationActions migrationActions) {
+    public MigrationStateMachineImpl(StateMachine<MigrationState, MigrationStep> stateMachine, MigrationActions migrationActions, DatanodeMigrationPersistence persistenceService) {
         this.stateMachine = stateMachine;
         this.migrationActions = migrationActions;
+        this.persistenceService = persistenceService;
+        this.context = persistenceService.getStateMachineContext().orElse(new MigrationStateMachineContext());
     }
 
     @Override
-    public MigrationState trigger(MigrationStep step, Map<String, Object> args) {
+    public CurrentStateInformation trigger(MigrationStep step, Map<String, Object> args) {
+        context.setCurrentStep(step);
+        if (Objects.nonNull(args) && !args.isEmpty()) {
+            context.addActionArguments(step, args);
+        }
+        migrationActions.setStateMachineContext(context);
+        String errorMessage = null;
         try {
-            migrationActions.setArgs(args);
             stateMachine.fire(step);
-        } finally {
-            migrationActions.clearArgs();
+        } catch (Exception e) {
+            errorMessage = Objects.nonNull(e.getMessage()) ? e.getMessage() : e.toString();
         }
-        return stateMachine.getState();
+        context = migrationActions.getStateMachineContext();
+        persistenceService.saveStateMachineContext(context);
+        return new CurrentStateInformation(getState(), nextSteps(), errorMessage);
+    }
+
+    @Override
+    public MigrationStateMachineContext getContext() {
+        return context;
     }
 
     @Override

graylog2-server/src/main/java/org/graylog/plugins/views/storage/migration/state/machine/MigrationStateMachineProvider.java

@@ -38,6 +38,6 @@ public MigrationStateMachineProvider(DatanodeMigrationPersistence persistenceSer
     @Override
     public MigrationStateMachine get() {
         final StateMachine<MigrationState, MigrationStep> stateMachine = MigrationStateMachineBuilder.buildFromPersistedState(persistenceService, migrationActions);
-        return new MigrationStateMachineImpl(stateMachine, migrationActions);
+        return new MigrationStateMachineImpl(stateMachine, migrationActions, persistenceService);
     }
 }

graylog2-server/src/main/java/org/graylog/plugins/views/storage/migration/state/persistence/DatanodeMigrationConfigurationImpl.java

@@ -18,6 +18,7 @@
 
 import jakarta.inject.Inject;
 import jakarta.inject.Singleton;
+import org.graylog.plugins.views.storage.migration.state.machine.MigrationStateMachineContext;
 import org.graylog2.plugin.cluster.ClusterConfigService;
 
 import java.util.Optional;
@@ -41,4 +42,14 @@ public Optional<DatanodeMigrationConfiguration> getConfiguration() {
     public void saveConfiguration(DatanodeMigrationConfiguration configuration) {
         clusterConfigService.write(configuration);
     }
+
+    @Override
+    public Optional<MigrationStateMachineContext> getStateMachineContext() {
+        return Optional.ofNullable(clusterConfigService.get(MigrationStateMachineContext.class));
+    }
+
+    @Override
+    public void saveStateMachineContext(MigrationStateMachineContext context) {
+        clusterConfigService.write(context);
+    }
 }

graylog2-server/src/main/java/org/graylog/plugins/views/storage/migration/state/persistence/DatanodeMigrationPersistence.java

@@ -16,9 +16,17 @@
  */
 package org.graylog.plugins.views.storage.migration.state.persistence;
 
+import org.graylog.plugins.views.storage.migration.state.machine.MigrationStateMachineContext;
+
 import java.util.Optional;
 
 public interface DatanodeMigrationPersistence {
+
     Optional<DatanodeMigrationConfiguration> getConfiguration();
+
     void saveConfiguration(DatanodeMigrationConfiguration configuration);
+
+    Optional<MigrationStateMachineContext> getStateMachineContext();
+
+    void saveStateMachineContext(MigrationStateMachineContext context);
 }

graylog2-server/src/main/java/org/graylog/plugins/views/storage/migration/state/rest/CurrentStateInformation.java

@@ -16,10 +16,20 @@
  */
 package org.graylog.plugins.views.storage.migration.state.rest;
 
+import org.apache.commons.lang3.StringUtils;
 import org.graylog.plugins.views.storage.migration.state.machine.MigrationState;
 import org.graylog.plugins.views.storage.migration.state.machine.MigrationStep;
 
 import java.util.List;
 
-public record CurrentStateInformation(MigrationState state, List<MigrationStep> nextSteps) {
+public record CurrentStateInformation(MigrationState state, List<MigrationStep> nextSteps, String errorMessage) {
+
+    public CurrentStateInformation(MigrationState state, List<MigrationStep> nextSteps) {
+        this(state, nextSteps, null);
+    }
+
+    public boolean hasErrors() {
+        return StringUtils.isNotBlank(errorMessage);
+    }
+
 }