[GOBBLIN-2173] Avoid Adhoc flow spec addition for non leasable entity

gobblin-restli/gobblin-flow-config-service/gobblin-flow-config-service-server/src/main/java/org/apache/gobblin/service/modules/restli/FlowConfigsV2ResourceHandler.java

@@ -60,6 +60,7 @@
 import org.apache.gobblin.metrics.ServiceMetricNames;
 import org.apache.gobblin.runtime.api.FlowSpec;
 import org.apache.gobblin.runtime.api.FlowSpecSearchObject;
+import org.apache.gobblin.runtime.api.LeaseUnavailableException;
 import org.apache.gobblin.runtime.api.SpecNotFoundException;
 import org.apache.gobblin.runtime.spec_catalog.AddSpecResponse;
 import org.apache.gobblin.runtime.spec_catalog.FlowCatalog;
@@ -256,7 +257,10 @@ public CreateKVResponse<ComplexResourceKey<FlowId, FlowStatusId>, FlowConfig> cr
       responseMap = this.flowCatalog.put(flowSpec, true);
     } catch (QuotaExceededException e) {
       throw new RestLiServiceException(HttpStatus.S_503_SERVICE_UNAVAILABLE, e.getMessage());
-    } catch (Throwable e) {
+    } catch(LeaseUnavailableException e){
+      throw new RestLiServiceException(HttpStatus.S_409_CONFLICT, e.getMessage());
+    }
+    catch (Throwable e) {
       // TODO: Compilation errors should fall under throwable exceptions as well instead of checking for strings
       log.warn(String.format("Failed to add flow configuration %s.%s to catalog due to", flowConfig.getId().getFlowGroup(), flowConfig.getId().getFlowName()), e);
       throw new RestLiServiceException(HttpStatus.S_500_INTERNAL_SERVER_ERROR, e.getMessage());

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

@@ -0,0 +1,27 @@
+/*
+ * 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.runtime.api;
+
+/**
+ * An {@link RuntimeException} thrown when lease cannot be acquired on provided entity.
+ */
+public class LeaseUnavailableException extends RuntimeException {
+  public LeaseUnavailableException(String message) {
+    super(message);
+  }
+}

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

@@ -103,6 +103,13 @@ public interface DagManagementStateStore {
    */
   void updateDagNode(Dag.DagNode<JobExecutionPlan> dagNode) throws IOException;
 
+  /**
+   * Returns true if lease can be acquired on entity provided in leaseParams.
+   * @param leaseParams   uniquely identifies the flow, the present action upon it, the time the action
+   *                      was triggered, and if the dag action event we're checking on is a reminder event
+   */
+  boolean canAcquireLeaseOnEntity(DagActionStore.LeaseParams leaseParams) throws IOException;
+
   /**
    * Returns the requested {@link  org.apache.gobblin.service.modules.flowgraph.Dag.DagNode} and its {@link JobStatus}.
    * Both params are returned as optional and are empty if not present in the store.

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

@@ -90,6 +90,11 @@ public LeaseAttemptStatus tryAcquireLease(DagActionStore.LeaseParams leaseParams
     throw new RuntimeException(String.format("Unexpected LeaseAttemptStatus (%s) for %s", leaseAttemptStatus.getClass().getName(), leaseParams));
   }
 
+  @Override
+  public boolean canAcquireLeaseOnEntity(DagActionStore.LeaseParams leaseParams) throws IOException {
+    return decoratedMultiActiveLeaseArbiter.canAcquireLeaseOnEntity(leaseParams);
+  }
+
   @Override
   public boolean recordLeaseSuccess(LeaseAttemptStatus.LeaseObtainedStatus status)
       throws IOException {

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

@@ -61,6 +61,16 @@ public interface MultiActiveLeaseArbiter {
   LeaseAttemptStatus tryAcquireLease(DagActionStore.LeaseParams leaseParams, boolean adoptConsensusFlowExecutionId)
       throws IOException;
 
+  /**
+   * This method checks if lease can be acquired on provided flow in lease params
+   * returns true if entry for the same flow does not exists within epsilon time
+   * in leaseArbiterStore
+   * @param leaseParams   uniquely identifies the flow, the present action upon it, the time the action
+   *                      was triggered, and if the dag action event we're checking on is a reminder event
+   */
+  boolean canAcquireLeaseOnEntity(DagActionStore.LeaseParams leaseParams)
+      throws IOException;
+
   /**
    * This method is used to indicate the owner of the lease has successfully completed required actions while holding
    * the lease of the dag action event. It marks the lease as "no longer leasing", if the eventTimeMillis and

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

@@ -65,6 +65,7 @@ public class MySqlDagManagementStateStore implements DagManagementStateStore {
   // todo - these two stores should merge
   private DagStateStoreWithDagNodes dagStateStore;
   private DagStateStoreWithDagNodes failedDagStateStore;
+  private MultiActiveLeaseArbiter multiActiveLeaseArbiter;
   private final JobStatusRetriever jobStatusRetriever;
   private boolean dagStoresInitialized = false;
   private final UserQuotaManager quotaManager;
@@ -79,13 +80,14 @@ public class MySqlDagManagementStateStore implements DagManagementStateStore {
 
   @Inject
   public MySqlDagManagementStateStore(Config config, FlowCatalog flowCatalog, UserQuotaManager userQuotaManager,
-      JobStatusRetriever jobStatusRetriever, DagActionStore dagActionStore) {
+      JobStatusRetriever jobStatusRetriever, DagActionStore dagActionStore, MultiActiveLeaseArbiter multiActiveLeaseArbiter) {
     this.quotaManager = userQuotaManager;
     this.config = config;
     this.flowCatalog = flowCatalog;
     this.jobStatusRetriever = jobStatusRetriever;
     this.dagManagerMetrics.activate();
     this.dagActionStore = dagActionStore;
+    this.multiActiveLeaseArbiter = multiActiveLeaseArbiter;
    }
 
   // It should be called after topology spec map is set
@@ -168,6 +170,11 @@ public synchronized void updateDagNode(Dag.DagNode<JobExecutionPlan> dagNode)
     this.dagStateStore.updateDagNode(dagNode);
   }
 
+  @Override
+  public boolean canAcquireLeaseOnEntity(DagActionStore.LeaseParams leaseParams) throws IOException {
+    return multiActiveLeaseArbiter.canAcquireLeaseOnEntity(leaseParams);
+  }
+
   @Override
   public Optional<Dag<JobExecutionPlan>> getDag(Dag.DagId dagId) throws IOException {
     return Optional.ofNullable(this.dagStateStore.getDag(dagId));

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

@@ -362,6 +362,12 @@ else if (leaseValidityStatus == 2) {
     }
   }
 
+  @Override
+  public boolean canAcquireLeaseOnEntity(DagActionStore.LeaseParams leaseParams) throws IOException {
+    Optional<GetEventInfoResult> infoResult = getExistingEventInfo(leaseParams);
+    return infoResult.isPresent() ? !infoResult.get().isWithinEpsilon() : true;
+  }
+
   /**
    * Checks leaseArbiterTable for an existing entry for this dag action and event time
    */

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

@@ -17,13 +17,16 @@
 
 package org.apache.gobblin.service.modules.orchestration;
 
+import com.linkedin.restli.common.HttpStatus;
+import com.linkedin.restli.server.RestLiServiceException;
 import java.io.IOException;
 import java.net.URI;
 import java.util.Collections;
 import java.util.List;
 import java.util.Properties;
 import java.util.concurrent.TimeUnit;
 
+import org.apache.gobblin.runtime.api.LeaseUnavailableException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -78,6 +81,7 @@ public class Orchestrator implements SpecCatalogListener, Instrumentable {
   protected final SpecCompiler specCompiler;
   protected final TopologyCatalog topologyCatalog;
   private final JobStatusRetriever jobStatusRetriever;
+  private final DagManagementStateStore dagManagementStateStore;
 
   protected final MetricContext metricContext;
 
@@ -100,6 +104,7 @@ public Orchestrator(Config config, TopologyCatalog topologyCatalog, Optional<Log
     this.topologyCatalog = topologyCatalog;
     this.flowLaunchHandler = flowLaunchHandler;
     this.sharedFlowMetricsSingleton = sharedFlowMetricsSingleton;
+    this.dagManagementStateStore = dagManagementStateStore;
     this.jobStatusRetriever = jobStatusRetriever;
     this.specCompiler = flowCompilationValidationHelper.getSpecCompiler();
     // todo remove the need to set topology factory outside of constructor GOBBLIN-2056
@@ -125,6 +130,7 @@ public AddSpecResponse onAddSpec(Spec addedSpec) {
       _log.info("Orchestrator - onAdd[Topology]Spec: " + addedSpec);
       this.specCompiler.onAddSpec(addedSpec);
     } else if (addedSpec instanceof FlowSpec) {
+      validateAdhocFlowLeasability((FlowSpec) addedSpec);
       _log.info("Orchestrator - onAdd[Flow]Spec: " + addedSpec);
       return this.specCompiler.onAddSpec(addedSpec);
     } else {
@@ -133,6 +139,25 @@ public AddSpecResponse onAddSpec(Spec addedSpec) {
     return new AddSpecResponse<>(null);
   }
 
+  private void validateAdhocFlowLeasability(FlowSpec flowSpec) {
+    if (!flowSpec.isScheduled()) {
+      Config flowConfig = flowSpec.getConfig();
+      String flowGroup = flowConfig.getString(ConfigurationKeys.FLOW_GROUP_KEY);
+      String flowName = flowConfig.getString(ConfigurationKeys.FLOW_NAME_KEY);
+      DagActionStore.DagAction dagAction = DagActionStore.DagAction.forFlow(flowGroup, flowName,
+          FlowUtils.getOrCreateFlowExecutionId(flowSpec), DagActionStore.DagActionType.LAUNCH);
+      DagActionStore.LeaseParams leaseParams = new DagActionStore.LeaseParams(dagAction, System.currentTimeMillis());
+      try {
+        if (!dagManagementStateStore.canAcquireLeaseOnEntity(leaseParams)) {
+          throw new LeaseUnavailableException("Lease already occupied by another execution of this flow");
+        }
+      } catch (IOException exception) {
+        _log.error(String.format("Failed to query leaseArbiterTable for existing flow details: %s", flowSpec), exception);
+        throw new RuntimeException("Error querying leaseArbiterTable", exception);
+      }
+    }
+  }
+
   public void onDeleteSpec(URI deletedSpecURI, String deletedSpecVersion) {
     onDeleteSpec(deletedSpecURI, deletedSpecVersion, new Properties());
   }

gobblin-service/src/test/java/org/apache/gobblin/service/modules/orchestration/MySqlDagManagementStateStoreTest.java

@@ -24,6 +24,7 @@
 import java.util.Map;
 import java.util.Set;
 
+import org.mockito.Mockito;
 import org.testng.Assert;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
@@ -59,6 +60,7 @@
 public class MySqlDagManagementStateStoreTest {
 
   private ITestMetastoreDatabase testDb;
+  private static MultiActiveLeaseArbiter leaseArbiter;
   private MySqlDagManagementStateStore dagManagementStateStore;
   private static final String TEST_USER = "testUser";
   public static final String TEST_PASSWORD = "testPassword";
@@ -68,6 +70,7 @@ public class MySqlDagManagementStateStoreTest {
   @BeforeClass
   public void setUp() throws Exception {
     // Setting up mock DB
+    this.leaseArbiter = mock(MultiActiveLeaseArbiter.class);
     this.testDb = TestMetastoreDatabaseFactory.get();
     this.dagManagementStateStore = getDummyDMSS(this.testDb);
   }
@@ -92,6 +95,16 @@ public static <T> boolean compareLists(List<T> list1, List<T> list2) {
     return true;
   }
 
+  @Test
+  public void testcanAcquireLeaseOnEntity() throws Exception{
+    Mockito.when(leaseArbiter.canAcquireLeaseOnEntity(Mockito.any(DagActionStore.LeaseParams.class))).thenReturn(true);
+    String flowName = "testFlow";
+    String flowGroup = "testGroup";
+    DagActionStore.DagAction dagAction = new DagActionStore.DagAction(flowName, flowGroup, System.currentTimeMillis(), "testJob", DagActionStore.DagActionType.LAUNCH);
+    DagActionStore.LeaseParams leaseParams = new DagActionStore.LeaseParams(dagAction);
+    Assert.assertTrue(dagManagementStateStore.canAcquireLeaseOnEntity(leaseParams));
+  }
+
   @Test
   public void testAddDag() throws Exception {
     Dag<JobExecutionPlan> dag = DagTestUtils.buildDag("test", 12345L);
@@ -150,9 +163,11 @@ public static MySqlDagManagementStateStore getDummyDMSS(ITestMetastoreDatabase t
     TopologySpec topologySpec = LaunchDagProcTest.buildNaiveTopologySpec(TEST_SPEC_EXECUTOR_URI);
     URI specExecURI = new URI(TEST_SPEC_EXECUTOR_URI);
     topologySpecMap.put(specExecURI, topologySpec);
+    MultiActiveLeaseArbiter multiActiveLeaseArbiter = Mockito.mock(MultiActiveLeaseArbiter.class);
+    leaseArbiter = multiActiveLeaseArbiter;
     MySqlDagManagementStateStore dagManagementStateStore =
         new MySqlDagManagementStateStore(config, null, null, jobStatusRetriever,
-            MysqlDagActionStoreTest.getTestDagActionStore(testMetastoreDatabase));
+            MysqlDagActionStoreTest.getTestDagActionStore(testMetastoreDatabase), multiActiveLeaseArbiter);
     dagManagementStateStore.setTopologySpecMap(topologySpecMap);
     return dagManagementStateStore;
   }

gobblin-service/src/test/java/org/apache/gobblin/service/modules/orchestration/MysqlMultiActiveLeaseArbiterTest.java

@@ -44,6 +44,7 @@
 public class MysqlMultiActiveLeaseArbiterTest {
   private static final long EPSILON = 10000L;
   private static final long MORE_THAN_EPSILON = (long) (EPSILON * 1.1);
+  private static final long LESS_THAN_EPSILON = (long) (EPSILON * 0.90);
   // NOTE: `sleep`ing this long SIGNIFICANTLY slows tests, but we need a large enough value that exec. variability won't cause spurious failure
   private static final long LINGER = 20000L;
   private static final long MORE_THAN_LINGER = (long) (LINGER * 1.1);
@@ -53,6 +54,8 @@ public class MysqlMultiActiveLeaseArbiterTest {
   private static final String CONSTANTS_TABLE = "constants_store";
   private static final String flowGroup = "testFlowGroup";
   private static final String flowGroup2 = "testFlowGroup2";
+  private static final String flowGroup3 = "testFlowGroup3";
+  private static final String flowGroup4 = "testFlowGroup4";
   private static final String flowName = "testFlowName";
   private static final String jobName = "testJobName";
   private static final long flowExecutionId = 12345677L;
@@ -70,6 +73,14 @@ public class MysqlMultiActiveLeaseArbiterTest {
       new DagActionStore.DagAction(flowGroup2, flowName, flowExecutionId, jobName, DagActionStore.DagActionType.LAUNCH);
   private static final DagActionStore.LeaseParams
       launchLeaseParams2 = new DagActionStore.LeaseParams(launchDagAction2, false, eventTimeMillis);
+  private static final DagActionStore.DagAction launchDagAction3 =
+      new DagActionStore.DagAction(flowGroup3, flowName, flowExecutionId, jobName, DagActionStore.DagActionType.LAUNCH);
+  private static final DagActionStore.LeaseParams
+      launchLeaseParams3 = new DagActionStore.LeaseParams(launchDagAction3, false, eventTimeMillis);
+  private static final DagActionStore.DagAction launchDagAction4 =
+      new DagActionStore.DagAction(flowGroup4, flowName, flowExecutionId, jobName, DagActionStore.DagActionType.LAUNCH);
+  private static final DagActionStore.LeaseParams
+      launchLeaseParams4 = new DagActionStore.LeaseParams(launchDagAction4, false, eventTimeMillis);
   private static final Timestamp dummyTimestamp = new Timestamp(99999);
   private ITestMetastoreDatabase testDb;
   private MysqlMultiActiveLeaseArbiter mysqlMultiActiveLeaseArbiter;
@@ -201,6 +212,33 @@ public void testAcquireLeaseSingleParticipant() throws Exception {
         <= sixthObtainedStatus.getLeaseAcquisitionTimestamp());
   }
 
+  /*
+   test to verify if leasable entity is unavailable before epsilon time
+   to account for clock drift
+  */
+  @Test
+  public void testWhenLeasableEntityUnavailable() throws Exception{
+    LeaseAttemptStatus firstLaunchStatus =
+        mysqlMultiActiveLeaseArbiter.tryAcquireLease(launchLeaseParams3, true);
+    Assert.assertTrue(firstLaunchStatus instanceof LeaseAttemptStatus.LeaseObtainedStatus);
+    completeLeaseHelper(launchLeaseParams3);
+    Thread.sleep(LESS_THAN_EPSILON);
+    Assert.assertFalse(mysqlMultiActiveLeaseArbiter.canAcquireLeaseOnEntity(launchLeaseParams3));
+  }
+
+  /*
+     test to verify if leasable entity exists post epsilon time
+   */
+  @Test
+  public void testWhenLeasableEntityAvailable() throws Exception{
+    LeaseAttemptStatus firstLaunchStatus =
+        mysqlMultiActiveLeaseArbiter.tryAcquireLease(launchLeaseParams4, true);
+    Assert.assertTrue(firstLaunchStatus instanceof LeaseAttemptStatus.LeaseObtainedStatus);
+    completeLeaseHelper(launchLeaseParams4);
+    Thread.sleep(MORE_THAN_EPSILON);
+    Assert.assertTrue(mysqlMultiActiveLeaseArbiter.canAcquireLeaseOnEntity(launchLeaseParams4));
+  }
+
   /*
      Tests attemptLeaseIfNewRow() method to ensure a new row is inserted if no row matches the primary key in the table.
      If such a row does exist, the method should disregard the resulting SQL error and return 0 rows updated, indicating