Skip to content
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import com.openframe.data.document.delivery.DeliveryType;
import com.openframe.data.document.rmm.schedule.ScheduleScript;
import com.openframe.data.document.rmm.schedule.ScheduledScriptCustomParams;
import com.openframe.data.document.rmm.script.ExecutionSource;
import com.openframe.data.document.rmm.script.ExecutionStatus;
import com.openframe.data.document.rmm.script.Script;
import com.openframe.data.document.rmm.script.ScriptExecution;
Expand All @@ -23,7 +24,9 @@
import org.springframework.stereotype.Service;

import java.time.Instant;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
Expand All @@ -36,14 +39,16 @@
import static java.util.stream.Collectors.toSet;

// a batch waiting on the engine for a machine carries what the user sees now, not what was there when the schedule
// fired: an edit of a script or of the schedule's parameters replaces the waiting bodies, a schedule or a device that
// is gone withdraws them. A leaf still QUEUED is the sign the machine has not received the batch
// fired: an edit of a script or of the schedule replaces the waiting bodies and the leaves follow (a script added to
// the schedule gets a leaf, one removed closes), a schedule or a device that is gone withdraws them. A leaf still
// QUEUED is the sign the machine has not received the batch
@Slf4j
@Service
@RequiredArgsConstructor
public class ScheduleDeliveryRefresher {

private static final String NOTHING_LEFT_TO_RUN = "no runnable script left in the batch";
private static final String REMOVED_FROM_SCHEDULE = "Script was removed from the schedule before the agent received it";

private final ScriptExecutionRepository scriptExecutionRepository;
private final ScriptScheduleRepository scheduleRepository;
Expand Down Expand Up @@ -109,7 +114,7 @@
return;
}
ScheduleScript current = schedule.get();
List<Script> scripts = runnableScriptsOf(tenantId, executionId, current);
List<Script> scripts = runnableScriptsOf(tenantId, current);
if (scripts.isEmpty()) {
cancel(queuedLeaves, NOTHING_LEFT_TO_RUN);
return;
Expand All @@ -126,27 +131,80 @@
ScheduleOfflineRule rule = ScheduleOfflineRule.of(current.getTrigger(), current.getOfflineBehavior(), current.getReconnectWindowSeconds());
DeliveryOfflineBehavior offlineBehavior = rule.behavior();
Long reconnectWindowSeconds = rule.reconnectWindowSeconds();
long replaced = waitingMachines.stream()
Set<String> replaced = waitingMachines.stream()
.filter(machineId -> deliveryRefresher.refresh(DeliveryType.SCRIPT_SCHEDULE, executionId, machineId, message,
offlineBehavior, reconnectWindowSeconds))
.count();
followTheScripts(queuedLeaves, scripts);
log.info("Waiting batches refreshed: executionId={} machines={} replaced={}", executionId, waitingMachines.size(), replaced);
.collect(toSet());
List<ScriptExecution> leavesOfTheFire = scriptExecutionRepository.findByTenantIdAndExecutionId(tenantId, executionId);
alignLeaves(leavesOfTheFire, replaced, scripts, first);

Check failure on line 139 in openframe-api-lib/src/main/java/com/openframe/api/service/rmm/schedule/ScheduleDeliveryRefresher.java

View workflow job for this annotation

GitHub Actions / Flamingo Code Review

alignLeaves closes leaves of machines not in replacedMachines, orphaning their rows

alignLeaves closes leaves of machines not in replacedMachines, orphaning their rows refreshExecution computes `replaced` as only the machines for which deliveryRefresher.refresh(...) returned true (line 134-137). alignLeaves is then called only with that `replaced` set (line 139), and inside alignLeaves the for-loop over `replacedMachines` (line 151) is the only place that looks at `leavesOfTheFire`. A machine that was NOT replaced (deliveryRefresher.refresh returned false, e.g. because the device is offline and the offline-behavior keeps the old row, or the machine simply isn't in `waitingMachines` for this executionId) never has its leaves touched by alignLeaves. That by itself is correct per the new test 'refreshForSchedule_rowNotReplaced_leavesUntouched'. However, leavesOfTheFire is fetched via `findByTenantIdAndExecutionId(tenantId, executionId)`, i.e. it contains leaves for ALL machines under this executionId, not just the waiting/queued ones. Machines whose leaf status is already a terminal state (e.g. RUNNING or already FAILED from a prior refresh) for scripts removed from the schedule are never marked FAILED with REMOVED_FROM_SCHEDULE, because the forEach at line 155-160 only considers machines in `replacedMachines`. So a schedule edit that removes a script will leave stale QUEUED leaves for any machine whose delivery-refresh call returned false (e.g. transient refresh failure, or a machine that is not currently in the `waitingMachines` set from a stale read) indefinitely in QUEUED state for a script that no longer exists in the schedule — contradicting the class-level invariant comment that 'a script removed from the schedule closes' its leaf for every machine, not only the replaced ones. **What it risks:** A leaf for a machine whose delivery-refresh happened to return false (network hiccup, offline machine not covered by the offline-behavior path, or any other false from DeliveryRefresher.refresh) keeps a QUEUED leaf referencing a script-id removed from the schedule; the watchdog/agent-side reconciliation for that leaf has nothing to resolve it since runnableScriptsOf no longer returns that script, leaving an orphaned QUEUED row that never transitions to a terminal state. **How to fix:** Iterate leaf alignment over all machines present in leavesOfTheFire (or at minimum over waitingMachines, not just the subset that returned true from refresh), so a script removed from the schedule closes the leaf regardless of whether the delivery-refresh call itself succeeded for that machine.
Comment thread
semen-flamingo marked this conversation as resolved.
Resolved
log.info("Waiting batches refreshed: executionId={} machines={} replaced={}", executionId, waitingMachines.size(), replaced.size());
}

// the leaf keeps what the agent will run with: the script's privilege and timeout as they are now (the result
// watchdog reads the timeout from the leaf)
private void followTheScripts(List<ScriptExecution> queuedLeaves, List<Script> scripts) {
// the leaves of a machine whose waiting body was replaced are the batch it will get: a script that left the batch
// closes, a script that joined it gets a leaf (or its closed leaf back, the fire holds one leaf per script and
// machine), the rest keep the script's privilege and timeout as they are now (the result watchdog reads the leaf)
private void alignLeaves(List<ScriptExecution> leavesOfTheFire, Set<String> replacedMachines, List<Script> scripts,
ScriptExecution sample) {
Instant now = Instant.now();
Map<String, Script> byId = scripts.stream().collect(toMap(Script::getId, Function.identity(), (first, second) -> first));
List<ScriptExecution> changed = queuedLeaves.stream()
.filter(leaf -> byId.containsKey(leaf.getScriptId()))
.filter(leaf -> follow(leaf, byId.get(leaf.getScriptId())))
.toList();
List<ScriptExecution> changed = new ArrayList<>();
for (String machineId : replacedMachines) {
Map<String, ScriptExecution> byScript = leavesOfTheFire.stream()
.filter(leaf -> machineId.equals(leaf.getMachineId()))
.collect(toMap(ScriptExecution::getScriptId, Function.identity(), (first, second) -> first));
byScript.values().stream()
.filter(leaf -> isQueued(leaf) && !byId.containsKey(leaf.getScriptId()))
.forEach(leaf -> {
markFailed(leaf, now, REMOVED_FROM_SCHEDULE);
changed.add(leaf);
});
for (Script script : scripts) {
ScriptExecution leaf = byScript.get(script.getId());
if (leaf == null) {
changed.add(newLeaf(sample, machineId, script, now));
} else if (!isQueued(leaf)) {
reopen(leaf, script, now);
changed.add(leaf);
} else if (follow(leaf, script)) {
changed.add(leaf);
}
}
}
if (!changed.isEmpty()) {
scriptExecutionRepository.saveAll(changed);
}
}

private static ScriptExecution newLeaf(ScriptExecution sample, String machineId, Script script, Instant now) {
return ScriptExecution.builder()
.tenantId(sample.getTenantId())
.executionId(sample.getExecutionId())
.scheduleId(sample.getScheduleId())
.machineId(machineId)
.scriptId(script.getId())
.privilegeLevel(script.getPrivilegeLevel())
.timeoutSeconds(script.getDefaultTimeoutSeconds())
.initiatedBy(sample.getInitiatedBy())
.source(ExecutionSource.SCHEDULED)
.status(ExecutionStatus.QUEUED)
.dispatchedAt(sample.getDispatchedAt())
.statusChangedAt(now)
.build();
}

Check warning on line 193 in openframe-api-lib/src/main/java/com/openframe/api/service/rmm/schedule/ScheduleDeliveryRefresher.java

View workflow job for this annotation

GitHub Actions / Flamingo Code Review

newLeaf ignores per-machine existing timestamps and always uses sample.getDispatchedAt()

newLeaf ignores per-machine existing timestamps and always uses sample.getDispatchedAt() In alignLeaves, when a machine gets a brand-new leaf for a script that joined the schedule (line 163-164, newLeaf), the leaf's dispatchedAt is copied from `sample` -- the first element of the original queuedLeaves list for the whole execution (set at line 110, `ScriptExecution first = queuedLeaves.getFirst()`). That `first` leaf belongs to whichever machine happened to be first in the queued list, not to the `machineId` the new leaf is being created for. If different machines in the same executionId were dispatched at different times (e.g. staggered delivery, or one machine joined the schedule set later), the new leaf for machineId gets a dispatchedAt timestamp that belongs to a different machine's leaf, which can mislead the result-watchdog's timeout computation (timeoutSeconds is measured from dispatchedAt). **What it risks:** A script newly added to the schedule gets a leaf whose dispatchedAt is borrowed from an unrelated machine's leaf; if that borrowed dispatchedAt is earlier than when this leaf is actually created, the watchdog could compute the timeout window as already elapsed and prematurely mark the leaf as timed out. **How to fix:** Use the dispatch time relevant to the machine the leaf is created for (e.g. `now`, since the leaf is newly queued at this moment) rather than `sample.getDispatchedAt()`.
Comment thread
semen-flamingo marked this conversation as resolved.
Resolved

private static void reopen(ScriptExecution leaf, Script script, Instant now) {
leaf.setStatus(ExecutionStatus.QUEUED);
leaf.setStatusChangedAt(now);
leaf.setFinishedAt(null);
leaf.setTimedOut(null);
leaf.setError(null);
follow(leaf, script);
}

private static boolean isQueued(ScriptExecution leaf) {
return leaf.getStatus() == ExecutionStatus.QUEUED;
}

private static boolean follow(ScriptExecution leaf, Script script) {
boolean changed = !Objects.equals(leaf.getPrivilegeLevel(), script.getPrivilegeLevel())
|| !Objects.equals(leaf.getTimeoutSeconds(), script.getDefaultTimeoutSeconds());
Expand All @@ -164,21 +222,14 @@
.filter(ScheduleDeliveryRefresher::isActive);
}

// the fire's own scripts, in the schedule's order, as they are now: a script the user removed from the schedule
// since stays in a batch that was built with it, the leaves of the fire are what the user sees of it
private List<Script> runnableScriptsOf(String tenantId, String executionId, ScheduleScript schedule) {
Set<String> fired = scriptExecutionRepository.findByTenantIdAndExecutionId(tenantId, executionId).stream()
.map(ScriptExecution::getScriptId)
.collect(toSet());
Map<String, Script> byId = scriptRepository.findByTenantIdAndIdIn(tenantId, fired).stream()
// the schedule's scripts as they are now, in its order, active only: what the machine gets is the schedule the
// user sees, not the one that fired
private List<Script> runnableScriptsOf(String tenantId, ScheduleScript schedule) {
List<String> ids = schedule.getScriptIds() == null ? List.of() : schedule.getScriptIds().stream().distinct().toList();
Map<String, Script> byId = scriptRepository.findByTenantIdAndIdIn(tenantId, new HashSet<>(ids)).stream()
.filter(ScheduleDeliveryRefresher::isActive)
.collect(toMap(Script::getId, Function.identity(), (first, second) -> first));
List<String> orderedIds = schedule.getScriptIds() == null ? List.of() : schedule.getScriptIds();
List<Script> ordered = orderedIds.stream().distinct().map(byId::get).filter(Objects::nonNull).toList();
if (!ordered.isEmpty()) {
return ordered;
}
return fired.stream().map(byId::get).filter(Objects::nonNull).toList();
return ids.stream().map(byId::get).filter(Objects::nonNull).toList();
}

private void cancel(List<ScriptExecution> queuedLeaves, String reason) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -187,7 +187,6 @@ void retireScript_otherScriptsLeft_batchRebuiltWithoutItAndItsLeavesClosed() {
void retireScript_nothingLeft_rowsWithdrawnAndLeavesClosed() {
ScriptExecution leaf = queued(EXECUTION_ID, MACHINE_1, SCRIPT_A);
when(scriptExecutionRepository.findByTenantIdAndScriptIdAndStatus(TENANT_ID, SCRIPT_A, ExecutionStatus.QUEUED)).thenReturn(List.of(leaf));
when(scriptExecutionRepository.findByTenantIdAndExecutionId(TENANT_ID, EXECUTION_ID)).thenReturn(List.of(leaf));
when(scheduleRepository.findByTenantIdAndId(TENANT_ID, SCHEDULE_ID)).thenReturn(Optional.of(schedule(ScriptStatus.ACTIVE, List.of(SCRIPT_A))));
Script deleted = script(SCRIPT_A, "echo a");
deleted.setStatus(ScriptStatus.DELETED);
Expand Down Expand Up @@ -223,6 +222,72 @@ void refreshForSchedule_waitingSchedule_windowTravelsAndLeavesFollowTheScript()
verify(scriptExecutionRepository).saveAll(List.of(leaf));
}

@Test
@DisplayName("refreshForSchedule: a script added to the schedule joins the waiting batch: the body carries it and the machine gets a QUEUED leaf for it under the same fire")
void refreshForSchedule_scriptAdded_bodyAndLeafFollow() {
ScriptExecution leafA = queued(EXECUTION_ID, MACHINE_1, SCRIPT_A);
when(scriptExecutionRepository.findByTenantIdAndScheduleIdAndStatus(TENANT_ID, SCHEDULE_ID, ExecutionStatus.QUEUED)).thenReturn(List.of(leafA));
when(scriptExecutionRepository.findByTenantIdAndExecutionId(TENANT_ID, EXECUTION_ID)).thenReturn(List.of(leafA));
when(scheduleRepository.findByTenantIdAndId(TENANT_ID, SCHEDULE_ID)).thenReturn(Optional.of(schedule(ScriptStatus.ACTIVE, List.of(SCRIPT_A, SCRIPT_B))));
when(scriptRepository.findByTenantIdAndIdIn(eq(TENANT_ID), any())).thenReturn(List.of(script(SCRIPT_A, "echo a"), script(SCRIPT_B, "echo b")));
when(deliveryRefresher.refresh(eq(DeliveryType.SCRIPT_SCHEDULE), eq(EXECUTION_ID), eq(MACHINE_1), messageCaptor.capture(), any(), any())).thenReturn(true);

refresher.refreshForSchedule(SCHEDULE_ID);

assertThat(messageCaptor.getValue().getScripts()).extracting(ScriptScheduleExecutionItem::getScriptId).containsExactly(SCRIPT_A, SCRIPT_B);
ArgumentCaptor<List<ScriptExecution>> saved = leavesCaptor();
verify(scriptExecutionRepository).saveAll(saved.capture());
assertThat(saved.getValue()).extracting(ScriptExecution::getScriptId).containsExactly(SCRIPT_A, SCRIPT_B);
ScriptExecution leafB = saved.getValue().getLast();
assertThat(leafB.getExecutionId()).isEqualTo(EXECUTION_ID);
assertThat(leafB.getMachineId()).isEqualTo(MACHINE_1);
assertThat(leafB.getStatus()).isEqualTo(ExecutionStatus.QUEUED);
assertThat(leafB.getInitiatedBy()).isEqualTo(OWNER);
assertThat(leafB.getTimeoutSeconds()).isEqualTo(60);
}

@Test
@DisplayName("refreshForSchedule: a script removed from the schedule leaves the waiting batch: its leaf closes with the reason, the body goes on without it")
void refreshForSchedule_scriptRemoved_leafClosesBodyGoesOn() {
ScriptExecution leafA = queued(EXECUTION_ID, MACHINE_1, SCRIPT_A);
ScriptExecution leafB = queued(EXECUTION_ID, MACHINE_1, SCRIPT_B);
when(scriptExecutionRepository.findByTenantIdAndScheduleIdAndStatus(TENANT_ID, SCHEDULE_ID, ExecutionStatus.QUEUED)).thenReturn(List.of(leafA, leafB));
when(scriptExecutionRepository.findByTenantIdAndExecutionId(TENANT_ID, EXECUTION_ID)).thenReturn(List.of(leafA, leafB));
when(scheduleRepository.findByTenantIdAndId(TENANT_ID, SCHEDULE_ID)).thenReturn(Optional.of(schedule(ScriptStatus.ACTIVE, List.of(SCRIPT_B))));
when(scriptRepository.findByTenantIdAndIdIn(eq(TENANT_ID), any())).thenReturn(List.of(script(SCRIPT_B, "echo b")));
when(deliveryRefresher.refresh(eq(DeliveryType.SCRIPT_SCHEDULE), eq(EXECUTION_ID), eq(MACHINE_1), messageCaptor.capture(), any(), any())).thenReturn(true);

refresher.refreshForSchedule(SCHEDULE_ID);

assertThat(messageCaptor.getValue().getScripts()).extracting(ScriptScheduleExecutionItem::getScriptId).containsExactly(SCRIPT_B);
assertThat(leafA.getStatus()).isEqualTo(ExecutionStatus.FAILED);
assertThat(leafA.getError()).isEqualTo("Script was removed from the schedule before the agent received it");
assertThat(leafB.getStatus()).isEqualTo(ExecutionStatus.QUEUED);
assertThat(leafB.getTimeoutSeconds()).isEqualTo(60);
verify(scriptExecutionRepository).saveAll(List.of(leafA, leafB));
}

@Test
@DisplayName("refreshForSchedule: the machine already has the batch (its row was not replaced): its leaves are left as they are")
void refreshForSchedule_rowNotReplaced_leavesUntouched() {
ScriptExecution leafA = queued(EXECUTION_ID, MACHINE_1, SCRIPT_A);
when(scriptExecutionRepository.findByTenantIdAndScheduleIdAndStatus(TENANT_ID, SCHEDULE_ID, ExecutionStatus.QUEUED)).thenReturn(List.of(leafA));
when(scriptExecutionRepository.findByTenantIdAndExecutionId(TENANT_ID, EXECUTION_ID)).thenReturn(List.of(leafA));
when(scheduleRepository.findByTenantIdAndId(TENANT_ID, SCHEDULE_ID)).thenReturn(Optional.of(schedule(ScriptStatus.ACTIVE, List.of(SCRIPT_B))));
when(scriptRepository.findByTenantIdAndIdIn(eq(TENANT_ID), any())).thenReturn(List.of(script(SCRIPT_B, "echo b")));
when(deliveryRefresher.refresh(eq(DeliveryType.SCRIPT_SCHEDULE), eq(EXECUTION_ID), eq(MACHINE_1), any(), any(), any())).thenReturn(false);

refresher.refreshForSchedule(SCHEDULE_ID);

assertThat(leafA.getStatus()).isEqualTo(ExecutionStatus.QUEUED);
verify(scriptExecutionRepository, never()).saveAll(any());
}

@SuppressWarnings("unchecked")
private static ArgumentCaptor<List<ScriptExecution>> leavesCaptor() {
return ArgumentCaptor.forClass(List.class);
}

private static ScriptExecution queued(String executionId, String machineId, String scriptId) {
return ScriptExecution.builder()
.tenantId(TENANT_ID)
Expand Down
Loading