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()
     {

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

@@ -45,6 +45,7 @@
 import static org.apache.cassandra.sidecar.common.data.OperationalJobStatus.SUCCEEDED;
 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.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.timeout;
@@ -253,6 +254,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
     {
@@ -301,4 +325,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();
+            }
+        };
+    }
 }