Add releaseOnCompletion hook to OperationalJob

Author: andresbeckruiz

server/src/main/java/org/apache/cassandra/sidecar/job/OperationalJob.java

@@ -183,6 +183,24 @@ public boolean requiresCoordination()
         return false;
     }
 
+    /**
+     * Whether the manager should release the active operation lock (via
+     * {@link OperationalJobCoordinator#clearActive}) when this job completes locally. Only consulted for jobs that
+     * {@link #requiresCoordination() require coordination}.
+     * <p>
+     * Single-node coordinated jobs return {@code true} (the default): the Sidecar that acquires the lock also
+     * finishes the work, so releasing on local completion is correct. Distributed cluster-wide jobs whose work
+     * finishes on other nodes return {@code false} and rely on the orchestration layer to call
+     * {@link OperationalJobCoordinator#clearActive} once all nodes reach a terminal state.
+     *
+     * @return {@code true} if the manager should release the lock on local completion; {@code false} to defer
+     *         release to the orchestration layer
+     */
+    public boolean releasesOnCompletion()
+    {
+        return true;
+    }
+
     @Override
     public final Void result()
     {
@@ -240,11 +258,6 @@ public String name()
         return simpleName.isEmpty() ? this.getClass().getName() : simpleName;
     }
 
-    /**
-     * @return the type of operation this job performs
-     */
-    public abstract OperationType operationType();
-
     /**
      * {@inheritDoc}
      */

server/src/main/java/org/apache/cassandra/sidecar/job/OperationalJobManager.java

@@ -127,7 +127,12 @@ public void trySubmitJob(OperationalJob job,
         {
             if (ar.succeeded() && Boolean.TRUE.equals(ar.result()))
             {
-                job.asyncResult().onComplete(result -> releaseActiveOperationLock(job));
+                // Distributed jobs that finish on other nodes opt out of auto-release; the orchestration
+                // layer clears the lock once all nodes reach a terminal state.
+                if (job.releasesOnCompletion())
+                {
+                    job.asyncResult().onComplete(result -> releaseActiveOperationLock(job));
+                }
                 trackAndExecute(job, onComplete, serviceExecutorPool, waitTime);
             }
             else

server/src/test/java/org/apache/cassandra/sidecar/job/OperationalJobManagerTest.java

@@ -49,6 +49,7 @@
 import static org.apache.cassandra.testing.utils.AssertionUtils.loopAssert;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.after;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -198,7 +199,7 @@ void testJobRemovedFromTrackerWhenPersistenceFails()
         DurableOperationalJobTracker durableTracker = new DurableOperationalJobTracker(new ServiceConfigurationImpl(),
                                                                                        storageProvider,
                                                                                        executorPool.service());
-        OperationalJobManager manager = new OperationalJobManager(durableTracker, executorPool);
+        OperationalJobManager manager = new OperationalJobManager(durableTracker, new DisabledOperationalJobCoordinator(), executorPool);
 
         UUID jobId = UUIDs.timeBased();
         OperationalJob job = OperationalJobTest.createOperationalJob(jobId, MillisecondBoundConfiguration.parse("50ms"));
@@ -283,6 +284,29 @@ void testCoordinationFailsWhenCoordinationDisabled() throws InterruptedException
         assertThat(tracker.inflightJobsByOperation(job.name())).doesNotContain(job);
     }
 
+    @Test
+    void testLockNotReleasedWhenJobOptsOut() throws InterruptedException
+    {
+        OperationalJobTracker tracker = new InMemoryOperationalJobTracker(4);
+        OperationalJobCoordinator coordinator = mock(OperationalJobCoordinator.class);
+        when(coordinator.trySetActive(any(), any())).thenReturn(true);
+        OperationalJobManager manager = new OperationalJobManager(tracker, coordinator, executorPool);
+        CountDownLatch latch = new CountDownLatch(1);
+
+        // A distributed cluster-wide job that acquires the lock locally but relies on the orchestration
+        // layer to clear it once all nodes finish, so the manager must not auto-release on local completion.
+        OperationalJob job = createNonReleasingCoordinatedJob(UUIDs.timeBased());
+        BiConsumer<OperationalJob, OperationalJobConflictException> onComplete = (j, ex) -> {
+            assertThat(ex).isNull();
+            latch.countDown();
+        };
+
+        manager.trySubmitJob(job, onComplete, executorPool.service(), SecondBoundConfiguration.parse("5s"));
+        assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue();
+        verify(coordinator).trySetActive(OperationType.MOVE, job.jobId());
+        verify(coordinator, after(1000).never()).clearActive(any(), any());
+    }
+
     @Test
     void testCoordinatorNotCalledWhenJobDoesNotRequireCoordination() throws InterruptedException
     {
@@ -331,4 +355,40 @@ protected Future<Void> executeInternal()
             }
         };
     }
+
+    private static OperationalJob createNonReleasingCoordinatedJob(UUID jobId)
+    {
+        return new OperationalJob(jobId)
+        {
+            @Override
+            public boolean hasConflict(@NotNull List<OperationalJob> sameOperationJobs)
+            {
+                return false;
+            }
+
+            @Override
+            public OperationType operationType()
+            {
+                return OperationType.MOVE;
+            }
+
+            @Override
+            public boolean requiresCoordination()
+            {
+                return true;
+            }
+
+            @Override
+            public boolean releasesOnCompletion()
+            {
+                return false;
+            }
+
+            @Override
+            protected Future<Void> executeInternal()
+            {
+                return Future.succeededFuture();
+            }
+        };
+    }
 }

server/src/test/java/org/apache/cassandra/sidecar/job/OperationalJobTest.java

@@ -85,12 +85,6 @@ public OperationalJobStatus status()
                 return jobStatus;
             }
 
-            @Override
-            public OperationType operationType()
-            {
-                return OperationType.DRAIN;
-            }
-
             @Override
             public String name()
             {
@@ -127,12 +121,6 @@ public OperationalJobStatus status()
                 return jobStatus;
             }
 
-            @Override
-            public OperationType operationType()
-            {
-                return OperationType.DRAIN;
-            }
-
             @Override
             public String name()
             {
@@ -177,12 +165,6 @@ protected Future<Void> executeInternal() throws OperationalJobException
                 return Future.succeededFuture();
             }
 
-            @Override
-            public OperationType operationType()
-            {
-                return OperationType.DRAIN;
-            }
-
             @Override
             public String name()
             {