UNOMI-939/940: Scheduler dispatch and cache refresh race fixes

services-common/src/main/java/org/apache/unomi/services/common/cache/AbstractMultiTypeCachingService.java

@@ -43,6 +43,9 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
@@ -68,6 +71,25 @@ public abstract class AbstractMultiTypeCachingService extends AbstractContextAwa
     protected TenantService tenantService;
     protected AuditService auditService;
 
+    /**
+     * Per-item-type locks ensuring saveItem()/removeItem() (direct writes) and refreshTypeCache()'s
+     * persistence-read-then-cache-write cycle can never interleave for the same item type. Without this,
+     * refreshTypeCache() could read a snapshot from persistence just before a concurrent direct write
+     * commits, then overwrite the cache with that stale snapshot after the direct write's own cache.put()
+     * already applied the fresh value. Direct writes only need to exclude a concurrent refresh of the
+     * same type, not each other, so they take the read lock (shared, concurrent) while refreshTypeCache()
+     * takes the write lock (exclusive) for its whole read-then-cache cycle, including the "remove items no
+     * longer in persistence" step, which reads from the same persistence snapshot.
+     *
+     * Cardinality is bounded by the number of distinct item types a service registers (getTypeConfigs()),
+     * typically one to a handful per service, so this map needs no eviction.
+     */
+    private final Map<String, ReadWriteLock> typeRefreshLocks = new ConcurrentHashMap<>();
+
+    private ReadWriteLock typeRefreshLock(String itemType) {
+        return typeRefreshLocks.computeIfAbsent(itemType, k -> new ReentrantReadWriteLock());
+    }
+
     /**
      * Map tracking which plugin/bundle contributed which items.
      * Key is the bundle ID, value is the list of items contributed by that bundle.
@@ -261,12 +283,25 @@ protected void shutdownTimers() {
         scheduledRefreshTasks.clear();
     }
 
-    @SuppressWarnings("unchecked")
     protected <T extends Serializable> void refreshTypeCache(CacheableTypeConfig<T> config) {
         if (!config.isRequiresRefresh()) {
             return;
         }
 
+        // Excludes any concurrent saveItem()/removeItem() for this item type for the duration of this
+        // refresh, so its persistence read can never be overwritten by - or overwrite - a direct write.
+        // See typeRefreshLocks for the full rationale.
+        Lock refreshLock = typeRefreshLock(config.getItemType()).writeLock();
+        refreshLock.lock();
+        try {
+            refreshTypeCacheLocked(config);
+        } finally {
+            refreshLock.unlock();
+        }
+    }
+
+    @SuppressWarnings("unchecked")
+    private <T extends Serializable> void refreshTypeCacheLocked(CacheableTypeConfig<T> config) {
         // Only create the global state maps if we need them
         Map<String, Map<String, T>> oldGlobalState = null;
         Map<String, Map<String, T>> newGlobalState = null;
@@ -858,8 +893,15 @@ protected <T extends Item & Serializable> void saveItem(T item, Function<T, Stri
             }
         }
 
-        persistenceService.save(item);
-        cacheService.put(itemType, itemId, currentTenant, item);
+        // Excludes a concurrent refreshTypeCache() for this item type: see typeRefreshLocks.
+        Lock writeGuard = typeRefreshLock(itemType).readLock();
+        writeGuard.lock();
+        try {
+            persistenceService.save(item);
+            cacheService.put(itemType, itemId, currentTenant, item);
+        } finally {
+            writeGuard.unlock();
+        }
     }
 
     /**
@@ -872,7 +914,14 @@ protected <T extends Item & Serializable> void saveItem(T item, Function<T, Stri
      */
     protected <T extends Item & Serializable> void removeItem(String id, Class<T> itemClass, String itemType) {
         String currentTenant = contextManager.getCurrentContext().getTenantId();
-        persistenceService.remove(id, itemClass);
-        cacheService.remove(itemType, id, currentTenant, itemClass);
+        // Excludes a concurrent refreshTypeCache() for this item type: see typeRefreshLocks.
+        Lock writeGuard = typeRefreshLock(itemType).readLock();
+        writeGuard.lock();
+        try {
+            persistenceService.remove(id, itemClass);
+            cacheService.remove(itemType, id, currentTenant, itemClass);
+        } finally {
+            writeGuard.unlock();
+        }
     }
 }

services-common/src/test/java/org/apache/unomi/services/common/cache/AbstractMultiTypeCachingServiceTest.java

@@ -16,6 +16,7 @@
  */
 package org.apache.unomi.services.common.cache;
 
+import org.apache.unomi.api.ExecutionContext;
 import org.apache.unomi.api.Item;
 import org.apache.unomi.api.services.ExecutionContextManager;
 import org.apache.unomi.api.services.cache.CacheableTypeConfig;
@@ -32,6 +33,9 @@
 
 import java.io.Serializable;
 import java.util.*;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
 
 import static org.junit.Assert.*;
@@ -110,6 +114,18 @@ public String toString() {
         }
     }
 
+    // Minimal Item subclass: saveItem()/removeItem() require <T extends Item & Serializable>,
+    // which the plain-Serializable TestSerializable above does not satisfy.
+    // Public (not just the ITEM_TYPE field): Item's reflective ITEM_TYPE lookup needs the
+    // declaring class itself to be accessible, not just the field's own modifiers.
+    public static class TestItem extends Item {
+        public static final String ITEM_TYPE = TEST_ITEM_TYPE;
+
+        public TestItem(String itemId) {
+            super(itemId);
+        }
+    }
+
     private static class TestCachingServiceImpl extends AbstractMultiTypeCachingService {
         private final Set<CacheableTypeConfig<?>> typeConfigs = new HashSet<>();
 
@@ -136,6 +152,11 @@ protected Set<CacheableTypeConfig<?>> getTypeConfigs() {
             return typeConfigs;
         }
 
+        // Exposes the protected saveItem() for the concurrency test below.
+        void callSaveItem(TestItem item) {
+            saveItem(item, Item::getItemId, TEST_ITEM_TYPE);
+        }
+
         // Helper method to set a config as persistable for testing
         void makeConfigPersistable() {
             try {
@@ -194,6 +215,9 @@ public void setUp() {
             runnable.run();
             return null;
         }).when(contextManager).executeAsSystem(any(Runnable.class));
+
+        // saveItem() reads the current tenant from contextManager.getCurrentContext().getTenantId().
+        when(contextManager.getCurrentContext()).thenReturn(new ExecutionContext(TEST_TENANT, null, null));
     }
 
     @Test
@@ -377,4 +401,70 @@ public void testRefreshCacheHandlesMultipleTenants() {
         // Verify items were removed from system tenant
         verify(cacheService).remove(eq(TEST_ITEM_TYPE), eq("item3"), eq(SYSTEM_TENANT), eq(TestSerializable.class));
     }
+
+    /**
+     * Validates the ReadWriteLock fix in AbstractMultiTypeCachingService: saveItem() must block while
+     * a refreshTypeCache() for the same item type holds the write lock, so a refresh's persistence read
+     * can never be followed by it overwriting a cache entry that a concurrent direct write just applied.
+     * Without the fix, saveItem() and refreshTypeCache() could interleave freely and the two race
+     * scenarios this test would otherwise need to get lucky on (read-before-write, write-after-read)
+     * would only manifest occasionally - here we force the interleaving deterministically instead.
+     */
+    @Test
+    public void testSaveItemExcludedByConcurrentRefresh() throws Exception {
+        when(cacheService.getTenantCache(eq(TEST_TENANT), eq(TestSerializable.class))).thenReturn(new HashMap<>());
+        when(cacheService.getTenantCache(eq(SYSTEM_TENANT), eq(TestSerializable.class))).thenReturn(new HashMap<>());
+        doReturn(Collections.singleton(TEST_TENANT)).when(testCachingService).getTenants();
+
+        CacheableTypeConfig<TestSerializable> config = null;
+        for (CacheableTypeConfig<?> typeConfig : testCachingService.getTypeConfigs()) {
+            if (typeConfig.getType().equals(TestSerializable.class)) {
+                @SuppressWarnings("unchecked")
+                CacheableTypeConfig<TestSerializable> typedConfig = (CacheableTypeConfig<TestSerializable>) typeConfig;
+                config = typedConfig;
+                break;
+            }
+        }
+        assertNotNull("Should find config for TestSerializable", config);
+        final CacheableTypeConfig<TestSerializable> finalConfig = config;
+
+        CountDownLatch refreshHoldingLock = new CountDownLatch(1);
+        CountDownLatch releaseRefresh = new CountDownLatch(1);
+        AtomicBoolean refreshStillRunningWhenSaveReturned = new AtomicBoolean(true);
+
+        // Stand-in for the persistence read inside refreshTypeCache(): signal that the refresh has
+        // entered its critical section (and therefore holds the write lock), then block until the
+        // test releases it.
+        doAnswer(invocation -> {
+            refreshHoldingLock.countDown();
+            assertTrue("Test setup: refresh should be released within timeout",
+                releaseRefresh.await(5, TimeUnit.SECONDS));
+            return Collections.emptyList();
+        }).when(testCachingService).loadItemsForTenant(eq(TEST_TENANT), eq(finalConfig));
+
+        Thread refreshThread = new Thread(() -> testCachingService.refreshTypeCache(finalConfig));
+        refreshThread.start();
+
+        assertTrue("Refresh should reach loadItemsForTenant while holding the write lock",
+            refreshHoldingLock.await(5, TimeUnit.SECONDS));
+
+        Thread saveThread = new Thread(() -> {
+            testCachingService.callSaveItem(new TestItem("item1"));
+            refreshStillRunningWhenSaveReturned.set(refreshThread.isAlive());
+        });
+        saveThread.start();
+
+        // saveItem()'s read lock should be blocked behind the refresh's write lock: give it a moment,
+        // then confirm it has NOT returned yet.
+        saveThread.join(500);
+        assertTrue("saveItem() should still be blocked while refresh holds the write lock", saveThread.isAlive());
+
+        releaseRefresh.countDown();
+        refreshThread.join(5000);
+        saveThread.join(5000);
+
+        assertFalse("saveThread should have finished", saveThread.isAlive());
+        assertFalse("saveItem() must not proceed until the concurrent refresh released its write lock",
+            refreshStillRunningWhenSaveReturned.get());
+    }
 }

services/src/main/java/org/apache/unomi/services/impl/scheduler/SchedulerServiceImpl.java

@@ -1172,10 +1172,14 @@ private void cancelTaskInternal(String taskId) {
         }
         ScheduledTask task = getTask(taskId);
         if (task != null) {
-            // Only cancel if in a cancellable state
+            // Only cancel if in a cancellable state. COMPLETED is included because a periodic task
+            // can be observed in that transient state by a caller racing handleTaskCompletion()
+            // (status briefly goes RUNNING -> COMPLETED -> SCHEDULED); without it, a cancelTask()
+            // call landing in that window would silently no-op and leave the task enabled.
             if (task.getStatus() == ScheduledTask.TaskStatus.SCHEDULED ||
                 task.getStatus() == ScheduledTask.TaskStatus.WAITING ||
-                task.getStatus() == ScheduledTask.TaskStatus.RUNNING) {
+                task.getStatus() == ScheduledTask.TaskStatus.RUNNING ||
+                task.getStatus() == ScheduledTask.TaskStatus.COMPLETED) {
 
                 task.setEnabled(false);
                 stateManager.updateTaskState(task, ScheduledTask.TaskStatus.CANCELLED, null, nodeId);

services/src/main/java/org/apache/unomi/services/impl/scheduler/TaskExecutionManager.java

@@ -176,13 +176,28 @@ public void executeTask(ScheduledTask task, TaskExecutor executor) {
             // Ensure the executing set exists even under concurrent clears during shutdown
             Set<String> executingSet = executingTasksByType.computeIfAbsent(taskType, k -> ConcurrentHashMap.newKeySet());
 
+            // Atomically claim this task before dispatching. The persisted/cached status check above
+            // is only updated once prepareForExecution() actually runs inside the async taskWrapper below,
+            // so a second poll tick landing in that gap could otherwise see a stale non-RUNNING status and
+            // dispatch the same task again. Set.add() is atomic, so only one caller can win this race.
+            if (!executingSet.add(task.getItemId())) {
+                LOGGER.debug("Node {} : Task {} is already dispatched for execution, skipping duplicate dispatch", nodeId, task.getItemId());
+                return;
+            }
+
             TaskExecutor.TaskStatusCallback statusCallback = createStatusCallback(task);
             Runnable taskWrapper = createTaskWrapper(task, executor, statusCallback);
 
             // Execute task immediately using the scheduler
-            ScheduledFuture<?> future = scheduler.schedule(taskWrapper, 0, TimeUnit.MILLISECONDS);
-            scheduledTasks.put(task.getItemId(), future);
-            executingSet.add(task.getItemId());
+            try {
+                ScheduledFuture<?> future = scheduler.schedule(taskWrapper, 0, TimeUnit.MILLISECONDS);
+                scheduledTasks.put(task.getItemId(), future);
+            } catch (Exception e) {
+                // Scheduling failed (e.g. rejected during shutdown): release the claim so a later
+                // attempt isn't permanently blocked, since the wrapper that would normally do so never ran.
+                executingSet.remove(task.getItemId());
+                throw e;
+            }
         } catch (Exception e) {
             LOGGER.error("Node "+nodeId+", Error executing task: " + task.getItemId(), e);
             handleTaskError(task, e.getMessage(), System.currentTimeMillis());
@@ -281,35 +296,34 @@ private Runnable createTaskWrapper(ScheduledTask task, TaskExecutor executor,
                 return;
             }
 
-            // Check shutdown again before preparing for execution
-            if (schedulerService != null && schedulerService.isShutdownNow()) {
-                LOGGER.debug("Node {} : Skipping task {} execution as scheduler is shutting down", nodeId, taskId);
-                return;
-            }
-
-            // Prepare task for execution (both persistent and in-memory)
-            if (!prepareForExecution(task)) {
-                return;
-            }
-
-            // Final shutdown check before executing
-            if (schedulerService != null && schedulerService.isShutdownNow()) {
-                LOGGER.debug("Node {} : Skipping task {} execution as scheduler is shutting down", nodeId, taskId);
-                return;
-            }
-
+            // From here on, executeTask() has already atomically claimed taskId in executingTasksByType
+            // before scheduling this wrapper. Every exit path below - including the early returns for
+            // shutdown and prepareForExecution() failing (e.g. lock already held by the caller, as
+            // happens when TaskRecoveryManager.attemptTaskResumption()/attemptTaskRestart() acquire the
+            // lock themselves before calling executeTask()) - must release that claim in the outer
+            // finally below, or the task is permanently blocked from ever being dispatched again.
+            boolean executingNodeIdSet = false;
             try {
-                // Get or create the executing tasks set
-                Set<String> executingTasks = executingTasksByType.computeIfAbsent(taskType,
-                    k -> ConcurrentHashMap.newKeySet());
+                // Check shutdown again before preparing for execution
+                if (schedulerService != null && schedulerService.isShutdownNow()) {
+                    LOGGER.debug("Node {} : Skipping task {} execution as scheduler is shutting down", nodeId, taskId);
+                    return;
+                }
+
+                // Prepare task for execution (both persistent and in-memory)
+                if (!prepareForExecution(task)) {
+                    return;
+                }
 
-                // Only add to executing set if not already there
-                if (taskId != null) {
-                    executingTasks.add(taskId);
+                // Final shutdown check before executing
+                if (schedulerService != null && schedulerService.isShutdownNow()) {
+                    LOGGER.debug("Node {} : Skipping task {} execution as scheduler is shutting down", nodeId, taskId);
+                    return;
                 }
 
                 // Set the executing node ID
                 task.setExecutingNodeId(nodeId);
+                executingNodeIdSet = true;
                 schedulerService.saveTask(task);
 
                 long startTime = System.currentTimeMillis();
@@ -329,11 +343,15 @@ private Runnable createTaskWrapper(ScheduledTask task, TaskExecutor executor,
                 LOGGER.error("Unexpected error while executing task: " + taskId, e);
                 statusCallback.fail("Unexpected error: " + e.getMessage());
             } finally {
-                // Clear executing node ID
-                task.setExecutingNodeId(null);
-                schedulerService.saveTask(task);
+                // Only clear/save executingNodeId if we actually set it above; otherwise we never
+                // touched the task and a redundant save here could race a concurrent legitimate holder.
+                if (executingNodeIdSet) {
+                    task.setExecutingNodeId(null);
+                    schedulerService.saveTask(task);
+                }
 
-                // Remove task from executing set
+                // Always release the dispatch claim taken by executeTask(), regardless of which path
+                // above was taken.
                 try {
                     Set<String> executingTasks = executingTasksByType.get(taskType);
                     if (executingTasks != null && taskId != null) {
@@ -424,7 +442,11 @@ private void handleTaskError(ScheduledTask task, String error, long startTime) {
                                 executeTask(task, executor);
                             }
                         };
-                        long retryDelay = task.getNextScheduledExecution().getTime() - System.currentTimeMillis();
+                        // Use the configured retry delay directly rather than re-deriving it from
+                        // nextScheduledExecution: that target was computed before the state/history/metrics
+                        // bookkeeping above ran, so subtracting "now" here would silently erode the delay
+                        // by however long that bookkeeping took (worse under slower/contended runners).
+                        long retryDelay = task.getRetryDelay();
                         scheduler.schedule(retryTask, retryDelay, TimeUnit.MILLISECONDS);
                         LOGGER.debug("Scheduled retry #{} for task {} in {} ms",
                             task.getFailureCount(), task.getItemId(), retryDelay);

services/src/main/java/org/apache/unomi/services/impl/scheduler/TaskRecoveryManager.java

@@ -134,6 +134,15 @@ private void recoverCrashedTask(ScheduledTask task) {
             return;
         }
 
+        // Re-check shutdown right before writing: preDestroy() marks RUNNING tasks owned by this
+        // node as crashed with a more specific "Interrupted by scheduler shutdown" cause, and that
+        // write must win over the generic "Node failure detected" cause below if both race.
+        if (shutdownNow) {
+            LOGGER.debug("Node {} Skipping recovery of task {} : {} as scheduler is shutting down",
+                nodeId, task.getTaskType(), task.getItemId());
+            return;
+        }
+
         // First mark as crashed and release lock
         String previousOwner = task.getLockOwner();
         if (task.getStatus() != ScheduledTask.TaskStatus.CRASHED) {