Fix FlowSpec Updating Function

gobblin-runtime/src/main/java/org/apache/gobblin/runtime/api/FlowSpec.java

@@ -26,7 +26,6 @@
 import com.linkedin.data.template.StringMap;
 import com.typesafe.config.Config;
 import com.typesafe.config.ConfigFactory;
-import com.typesafe.config.ConfigValueFactory;
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import java.net.URI;
 import java.net.URISyntaxException;
@@ -76,10 +75,11 @@ public class FlowSpec implements Configurable, Spec {
   /** Human-readable description of the flow spec */
   final String description;
 
-  /** Flow config as a typesafe config object*/
+  /** Flow config as a typesafe config object */
   final Config config;
 
-  /** Flow config as a properties collection for backwards compatibility */
+  /** Flow config as a properties collection for backwards compatibility
+   * It can be updated to store properties in addition to ones in the immutable Config object */
   // Note that this property is not strictly necessary as it can be generated from the typesafe
   // config. We use it as a cache until typesafe config is more widely adopted in Gobblin.
   final Properties configAsProperties;
@@ -125,6 +125,18 @@ public static FlowSpec.Builder builder(URI catalogURI, Properties flowProps) {
       throw new RuntimeException("Unable to create a FlowSpec URI: " + e, e);
     }
   }
+
+  /**
+   * Add new property at the specified path to the configAsProperties objects.
+   * Note: this does NOT update the Config so any property added through this function must be retrieved through the
+   * ConfigAsProperties field
+   * @param path
+   * @param value
+   */
+  public void addPropertyToConfigAsProperties(String path, String value) {
+    this.configAsProperties.setProperty(path, value);
+  }
+
   public void addCompilationError(String src, String dst, String errorMessage, int numberOfHops) {
     this.compilationErrors.add(new CompilationError(getConfig(), src, dst, errorMessage, numberOfHops));
   }
@@ -518,15 +530,5 @@ public static int maxFlowSpecUriLength() {
       return URI_SCHEME.length() + ":".length() // URI separator
         + URI_PATH_SEPARATOR.length() + ServiceConfigKeys.MAX_FLOW_NAME_LENGTH + URI_PATH_SEPARATOR.length() + ServiceConfigKeys.MAX_FLOW_GROUP_LENGTH;
     }
-
-    /**
-     * Create a new FlowSpec object with the added property defined by path and value parameters
-     * @param path key for new property
-     * @param value
-     */
-    public static FlowSpec createFlowSpecWithProperty(FlowSpec flowSpec, String path, String value) {
-      Config updatedConfig = flowSpec.getConfig().withValue(path, ConfigValueFactory.fromAnyRef(value));
-      return new Builder(flowSpec.getUri()).withConfig(updatedConfig).build();
-    }
   }
 }

gobblin-runtime/src/test/java/org/apache/gobblin/runtime/api/FlowSpecTest.java

@@ -17,7 +17,6 @@
 
 package org.apache.gobblin.runtime.api;
 
-import com.typesafe.config.Config;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.util.Properties;
@@ -30,15 +29,16 @@
 public class FlowSpecTest {
 
   /**
-   * Tests that the addProperty() function to ensure the new flowSpec returned has the original properties and updated
+   * Tests that the addPropertyToConfigAsProperties() function results in a modified configAsProperties object in the
+   * FlowSpec containing both the original properties and the added flowExecutionId property
    * ones
    * @throws URISyntaxException
    */
   @Test
-  public void testAddProperty() throws URISyntaxException {
+  public void testAddPropertyToConfigAsProperties() throws URISyntaxException {
     String flowGroup = "myGroup";
     String flowName = "myName";
-    String flowExecutionId = "myId";
+    String flowExecutionId = "1234";
     FlowId flowId = new FlowId().setFlowGroup(flowGroup).setFlowName(flowName);
     URI flowUri = FlowSpec.Utils.createFlowSpecUri(flowId);
 
@@ -48,19 +48,13 @@ public void testAddProperty() throws URISyntaxException {
     properties.setProperty(ConfigurationKeys.FLOW_NAME_KEY, flowName);
     properties.setProperty(ConfigurationKeys.FLOW_IS_REMINDER_EVENT_KEY, "true");
 
-    FlowSpec originalFlowSpec = FlowSpec.builder(flowUri).withConfigAsProperties(properties).build();
-    FlowSpec updatedFlowSpec = FlowSpec.Utils.createFlowSpecWithProperty(originalFlowSpec, ConfigurationKeys.FLOW_EXECUTION_ID_KEY, flowExecutionId);
+    FlowSpec flowSpec = FlowSpec.builder(flowUri).withConfigAsProperties(properties).build();
+    flowSpec.addPropertyToConfigAsProperties(ConfigurationKeys.FLOW_EXECUTION_ID_KEY, flowExecutionId);
 
-    Properties updatedProperties = updatedFlowSpec.getConfigAsProperties();
+    Properties updatedProperties = flowSpec.getConfigAsProperties();
     Assert.assertEquals(updatedProperties.getProperty(ConfigurationKeys.FLOW_EXECUTION_ID_KEY), flowExecutionId);
     Assert.assertEquals(updatedProperties.getProperty(ConfigurationKeys.FLOW_GROUP_KEY), flowGroup);
     Assert.assertEquals(updatedProperties.getProperty(ConfigurationKeys.FLOW_NAME_KEY), flowName);
     Assert.assertEquals(updatedProperties.getProperty(ConfigurationKeys.FLOW_IS_REMINDER_EVENT_KEY), "true");
-
-    Config updatedConfig = updatedFlowSpec.getConfig();
-    Assert.assertEquals(updatedConfig.getString(ConfigurationKeys.FLOW_EXECUTION_ID_KEY), flowExecutionId);
-    Assert.assertEquals(updatedConfig.getString(ConfigurationKeys.FLOW_GROUP_KEY), flowGroup);
-    Assert.assertEquals(updatedConfig.getString(ConfigurationKeys.FLOW_NAME_KEY), flowName);
-    Assert.assertEquals(updatedConfig.getString(ConfigurationKeys.FLOW_IS_REMINDER_EVENT_KEY), "true");
   }
 }

gobblin-service/src/main/java/org/apache/gobblin/service/modules/orchestration/DagManager.java

@@ -484,7 +484,9 @@ public synchronized void setActive(boolean active) {
           log.error("Exception encountered when shutting down DagManager threads.", e);
         }
       }
-    } catch (IOException e) {
+    } catch (Throwable e) {
+      // All exceptions should fail leader transition obviously to avoid case where transition to active fails to
+      // complete but is not apparent
       log.error("Exception encountered when activating the new DagManager", e);
       throw new RuntimeException(e);
     }
@@ -507,6 +509,9 @@ public void handleLaunchFlowEvent(DagActionStore.DagAction launchAction) {
           this.flowCompilationValidationHelper.createExecutionPlanIfValid(spec);
       if (optionalJobExecutionPlanDag.isPresent()) {
         addDag(optionalJobExecutionPlanDag.get(), true, true);
+      } else {
+        log.warn("Failed flow compilation of spec causing launch flow event to be skipped on startup. Flow {}", flowId);
+        this.dagManagerMetrics.incrementFailedLaunchCount();
       }
       // Upon handling the action, delete it so on leadership change this is not duplicated
       this.dagActionStore.get().deleteDagAction(launchAction);

gobblin-service/src/main/java/org/apache/gobblin/service/modules/utils/FlowCompilationValidationHelper.java

@@ -52,7 +52,7 @@
  */
 @Slf4j
 @Data
-public final class FlowCompilationValidationHelper {
+public class FlowCompilationValidationHelper {
   private final SharedFlowMetricsSingleton sharedFlowMetricsSingleton;
   private final SpecCompiler specCompiler;
   private final UserQuotaManager quotaManager;
@@ -144,7 +144,7 @@ public Optional<Dag<JobExecutionPlan>> validateAndHandleConcurrentExecution(Conf
    * @param allowConcurrentExecution
    * @return true if the {@link FlowSpec} allows concurrent executions or if no other instance of the flow is currently RUNNING.
    */
-  private boolean isExecutionPermitted(FlowStatusGenerator flowStatusGenerator, String flowName, String flowGroup,
+  protected boolean isExecutionPermitted(FlowStatusGenerator flowStatusGenerator, String flowName, String flowGroup,
       boolean allowConcurrentExecution) {
     return allowConcurrentExecution || !flowStatusGenerator.isFlowRunning(flowName, flowGroup);
   }

gobblin-service/src/main/java/org/apache/gobblin/service/monitoring/DagActionStoreChangeMonitor.java

@@ -203,8 +203,8 @@ protected void submitFlowToDagManagerHelper(String flowGroup, String flowName, S
       URI flowUri = FlowSpec.Utils.createFlowSpecUri(flowId);
       spec = (FlowSpec) flowCatalog.getSpecs(flowUri);
       // Adds flowExecutionId to config to ensure they are consistent across hosts
-      FlowSpec updatedSpec = FlowSpec.Utils.createFlowSpecWithProperty(spec, ConfigurationKeys.FLOW_EXECUTION_ID_KEY, flowExecutionId);
-      this.orchestrator.submitFlowToDagManager(updatedSpec);
+      spec.addPropertyToConfigAsProperties(ConfigurationKeys.FLOW_EXECUTION_ID_KEY, flowExecutionId);
+      this.orchestrator.submitFlowToDagManager(spec);
     } catch (URISyntaxException e) {
       log.warn("Could not create URI object for flowId {}. Exception {}", flowId, e.getMessage());
       this.failedFlowLaunchSubmissions.mark();

gobblin-service/src/test/java/org/apache/gobblin/service/modules/core/IdentityFlowToJobSpecCompilerTest.java

@@ -109,7 +109,7 @@ private void cleanUpDir(String dir) throws Exception {
     }
   }
 
-  private TopologySpec initTopologySpec() {
+  public static TopologySpec initTopologySpec() {
     Properties properties = new Properties();
     properties.put("specStore.fs.dir", TOPOLOGY_SPEC_STORE_DIR);
     properties.put("specExecInstance.capabilities", TEST_SOURCE_NAME + ":" + TEST_SINK_NAME);
@@ -125,11 +125,11 @@ private TopologySpec initTopologySpec() {
     return topologySpecBuilder.build();
   }
 
-  private FlowSpec initFlowSpec() {
+  public static FlowSpec initFlowSpec() {
     return initFlowSpec(TEST_FLOW_GROUP, TEST_FLOW_NAME, TEST_SOURCE_NAME, TEST_SINK_NAME);
   }
 
-  private FlowSpec initFlowSpec(String flowGroup, String flowName, String source, String destination) {
+  private static FlowSpec initFlowSpec(String flowGroup, String flowName, String source, String destination) {
     Properties properties = new Properties();
     properties.put(ConfigurationKeys.JOB_SCHEDULE_KEY, "* * * * *");
     properties.put(ConfigurationKeys.FLOW_GROUP_KEY, flowGroup);

gobblin-service/src/test/java/org/apache/gobblin/service/modules/utils/FlowCompilationValidationHelperTest.java

@@ -0,0 +1,105 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.gobblin.service.modules.utils;
+
+import com.google.common.base.Optional;
+import java.io.IOException;
+import java.util.Properties;
+import org.apache.gobblin.configuration.ConfigurationKeys;
+import org.apache.gobblin.metrics.event.EventSubmitter;
+import org.apache.gobblin.runtime.api.FlowSpec;
+import org.apache.gobblin.service.modules.core.IdentityFlowToJobSpecCompilerTest;
+import org.apache.gobblin.service.modules.flow.IdentityFlowToJobSpecCompiler;
+import org.apache.gobblin.service.modules.flow.SpecCompiler;
+import org.apache.gobblin.service.modules.flowgraph.Dag;
+import org.apache.gobblin.service.modules.orchestration.UserQuotaManager;
+import org.apache.gobblin.service.modules.spec.JobExecutionPlan;
+import org.apache.gobblin.service.monitoring.FlowStatusGenerator;
+import org.apache.gobblin.util.ConfigUtils;
+import org.junit.Assert;
+import org.testng.annotations.Test;
+
+import static org.mockito.Mockito.*;
+
+
+/**
+ * Test functionality provided by the helper class re-used between the DagManager and Orchestrator for flow compilation.
+ */
+public class FlowCompilationValidationHelperTest {
+
+  class MockFlowCompilationValidationHelper extends FlowCompilationValidationHelper {
+
+    public MockFlowCompilationValidationHelper(SharedFlowMetricsSingleton sharedFlowMetricsSingleton,
+        SpecCompiler specCompiler, UserQuotaManager quotaManager, Optional<EventSubmitter> eventSubmitter,
+        FlowStatusGenerator flowStatusGenerator, boolean isFlowConcurrencyEnabled) {
+      super(sharedFlowMetricsSingleton, specCompiler, quotaManager, eventSubmitter, flowStatusGenerator,
+          isFlowConcurrencyEnabled);
+    }
+
+    /*
+      In overriden function simply return if concurrent execution is allowed or not because flowStatusGenerator will be
+      mocked
+     */
+    @Override
+    protected boolean isExecutionPermitted(FlowStatusGenerator flowStatusGenerator, String flowName, String flowGroup,
+        boolean allowConcurrentExecution) {
+      return allowConcurrentExecution;
+    }
+
+  }
+
+  /*
+  Creates a mock {@link FlowCompilationValidationHelper} which has a valid {@link SpecCompiler} but mocks other
+  components.
+   */
+  MockFlowCompilationValidationHelper createMockFlowCompilationValidationHelper(boolean isFlowConcurrencyEnabled) {
+    SpecCompiler specCompiler = new IdentityFlowToJobSpecCompiler(ConfigUtils.propertiesToConfig(new Properties()));
+    specCompiler.onAddSpec(IdentityFlowToJobSpecCompilerTest.initTopologySpec());
+    return new MockFlowCompilationValidationHelper(mock(SharedFlowMetricsSingleton.class), specCompiler,
+        mock(UserQuotaManager.class), Optional.of(mock(EventSubmitter.class)), mock(FlowStatusGenerator.class),
+        isFlowConcurrencyEnabled);
+  }
+
+  /**
+   * Verifies that a flow spec that compiles to create a valid job execution plan dag will also generate one after
+   * having a flow execution id key-value pair added to its config
+   * @throws IOException
+   * @throws InterruptedException
+   */
+  @Test
+  public void compileFlowSpec() throws IOException, InterruptedException {
+    MockFlowCompilationValidationHelper mockFlowCompilationValidationHelper = createMockFlowCompilationValidationHelper(true);
+    FlowSpec flowSpec = IdentityFlowToJobSpecCompilerTest.initFlowSpec();
+    Optional<Dag<JobExecutionPlan>> dagOptional = mockFlowCompilationValidationHelper.createExecutionPlanIfValid(flowSpec);
+    // Assert FlowSpec compilation results in non-null or empty dag
+    Assert.assertTrue(dagOptional.isPresent());
+    Assert.assertNotNull(dagOptional.get());
+    Assert.assertTrue(dagOptional.get().getNodes().size() == 1);
+    Assert.assertEquals(dagOptional.get().getStartNodes().size(), 1);
+
+    // Update flow spec and check if still compiles and passes checks
+    flowSpec.addPropertyToConfigAsProperties(ConfigurationKeys.FLOW_EXECUTION_ID_KEY, "54321");
+    dagOptional = mockFlowCompilationValidationHelper.createExecutionPlanIfValid(flowSpec);
+    // Assert FlowSpec compilation results in non-null or empty dag
+    Assert.assertTrue(dagOptional.isPresent());
+    Assert.assertNotNull(dagOptional.get());
+    Assert.assertTrue(dagOptional.get().getNodes().size() == 1);
+    Assert.assertEquals(dagOptional.get().getStartNodes().size(), 1);
+  }
+
+}