[CDAP-21064] Validate workflow token on program status trigger.

cdap-app-fabric/src/main/java/io/cdap/cdap/internal/app/runtime/schedule/queue/JobQueueTable.java

@@ -29,10 +29,13 @@
 import io.cdap.cdap.api.schedule.Trigger;
 import io.cdap.cdap.common.conf.CConfiguration;
 import io.cdap.cdap.common.conf.Constants;
+import io.cdap.cdap.common.service.RetryStrategies;
+import io.cdap.cdap.common.service.RetryStrategy;
 import io.cdap.cdap.internal.app.runtime.schedule.ProgramSchedule;
 import io.cdap.cdap.internal.app.runtime.schedule.ProgramScheduleRecord;
 import io.cdap.cdap.internal.app.runtime.schedule.ProgramScheduleStatus;
 import io.cdap.cdap.internal.app.runtime.schedule.constraint.ConstraintCodec;
+import io.cdap.cdap.internal.app.runtime.schedule.trigger.NotificationContext;
 import io.cdap.cdap.internal.app.runtime.schedule.trigger.SatisfiableTrigger;
 import io.cdap.cdap.internal.app.runtime.schedule.trigger.TriggerCodec;
 import io.cdap.cdap.internal.app.store.AppMetadataStore;
@@ -57,8 +60,8 @@
 /**
  * Dataset that stores {@link Job}s, which correspond to schedules that have been triggered, but not
  * yet executed. The queue can have only one {@link Job.State#PENDING_TRIGGER} job per schedule. It
- * can have other jobs for the same schedule in various other states - marked for deletion, {@link
- * Job.State#PENDING_LAUNCH}, etc. <p/>
+ * can have other jobs for the same schedule in various other states - marked for deletion,
+ * {@link Job.State#PENDING_LAUNCH}, etc. <p/>
  *
  * The queue is designed to avoid conflicts when different services use it concurrently. For
  * instance, when notification processors are adding notifications to a job, the job can be
@@ -73,7 +76,8 @@
  * deletion, then the generation id is incremented by one and a new job is created, again creating
  * conflicts on concurrent creation of jobs. <p/>
  *
- * Row Key is in the following format for a job: <p/> &lt;partition_id>&lt;scheduleId>&lt;generationI>&lt;rowType
+ * Row Key is in the following format for a job: <p/>
+ * &lt;partition_id>&lt;scheduleId>&lt;generationI>&lt;rowType
  * <ul>
  *   <li>The &lt;partition_id> is a hash based upon the scheduleId</li>
  *   <li>The &lt;generationId> is used to distinguish jobs for the same schedule in the queue</li>
@@ -93,18 +97,20 @@ public class JobQueueTable implements JobQueue {
   private final StructuredTable jobQueueTable;
   private final AppMetadataStore appMetadataStore;
   private final int numPartitions;
+  private final RetryStrategy notificationRetryStrategy;
 
   JobQueueTable(StructuredTable jobQueueTable, AppMetadataStore appMetadataStore,
-      int numPartitions) {
+      CConfiguration cConf) {
     this.jobQueueTable = jobQueueTable;
     this.appMetadataStore = appMetadataStore;
-    this.numPartitions = numPartitions;
+    this.numPartitions = cConf.getInt(Constants.Scheduler.JOB_QUEUE_NUM_PARTITIONS);
+    this.notificationRetryStrategy = RetryStrategies.fromConfiguration(cConf,
+        "system.notification.");
   }
 
   public static JobQueueTable getJobQueue(StructuredTableContext context, CConfiguration cConf) {
     StructuredTable jobQueueTable = context.getTable(StoreDefinition.JobQueueStore.JOB_QUEUE_TABLE);
-    return new JobQueueTable(jobQueueTable, AppMetadataStore.create(context),
-        cConf.getInt(Constants.Scheduler.JOB_QUEUE_NUM_PARTITIONS));
+    return new JobQueueTable(jobQueueTable, AppMetadataStore.create(context), cConf);
   }
 
   @Override
@@ -209,7 +215,8 @@ private void addNotification(Job job, Notification notification) throws IOExcept
   }
 
   private boolean isTriggerSatisfied(ProgramSchedule schedule, List<Notification> notifications) {
-    return ((SatisfiableTrigger) schedule.getTrigger()).isSatisfied(schedule, notifications);
+    return ((SatisfiableTrigger) schedule.getTrigger()).isSatisfied(schedule,
+        new NotificationContext(notifications, appMetadataStore, notificationRetryStrategy));
   }
 
   @Override

cdap-app-fabric/src/main/java/io/cdap/cdap/internal/app/runtime/schedule/trigger/AndTrigger.java

@@ -19,7 +19,6 @@
 import io.cdap.cdap.api.schedule.Trigger;
 import io.cdap.cdap.api.schedule.TriggerInfo;
 import io.cdap.cdap.internal.app.runtime.schedule.ProgramSchedule;
-import io.cdap.cdap.proto.Notification;
 import io.cdap.cdap.proto.id.ProgramId;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -41,9 +40,9 @@ public AndTrigger(List<SatisfiableTrigger> triggers) {
 
 
   @Override
-  public boolean isSatisfied(ProgramSchedule schedule, List<Notification> notifications) {
+  public boolean isSatisfied(ProgramSchedule schedule, NotificationContext notificationContext) {
     for (Trigger trigger : getTriggers()) {
-      if (!((SatisfiableTrigger) trigger).isSatisfied(schedule, notifications)) {
+      if (!((SatisfiableTrigger) trigger).isSatisfied(schedule, notificationContext)) {
         return false;
       }
     }

cdap-app-fabric/src/main/java/io/cdap/cdap/internal/app/runtime/schedule/trigger/NotificationContext.java

@@ -0,0 +1,68 @@
+/*
+ * Copyright © 2025 Cask Data, Inc.
+ *
+ * Licensed 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 io.cdap.cdap.internal.app.runtime.schedule.trigger;
+
+import io.cdap.cdap.api.workflow.WorkflowToken;
+import io.cdap.cdap.common.service.RetryStrategy;
+import io.cdap.cdap.internal.app.store.AppMetadataStore;
+import io.cdap.cdap.proto.Notification;
+import io.cdap.cdap.proto.ProgramType;
+import io.cdap.cdap.proto.id.ProgramId;
+import io.cdap.cdap.proto.id.ProgramRunId;
+import io.cdap.cdap.proto.id.WorkflowId;
+import java.io.IOException;
+import java.util.List;
+
+/**
+ * Context object, exposing information that may be useful processing
+ * {@link io.cdap.cdap.proto.Notification} for a trigger.
+ */
+public class NotificationContext {
+
+  private final List<Notification> notifications;
+  private final AppMetadataStore appMetadataStore;
+  private final RetryStrategy retryStrategy;
+
+  public NotificationContext(List<Notification> notifications, AppMetadataStore appMetadataStore,
+      RetryStrategy retryStrategy) {
+    this.notifications = notifications;
+    this.appMetadataStore = appMetadataStore;
+    this.retryStrategy = retryStrategy;
+  }
+
+  public List<Notification> getNotifications() {
+    return notifications;
+  }
+
+  public RetryStrategy getRetryStrategy() {
+    return retryStrategy;
+  }
+
+  /**
+   * Fetches the {@link WorkflowToken} for the provided {@link ProgramRunId}.
+   *
+   * @return The workflow token if the program is a workflow, {@code null} otherwise.
+   */
+  public WorkflowToken getWorkflowToken(ProgramRunId programRunId) throws IOException {
+    ProgramId programId = programRunId.getParent();
+    if (!programId.getType().equals(ProgramType.WORKFLOW)) {
+      return null;
+    }
+    return appMetadataStore.getWorkflowToken(
+        new WorkflowId(programId.getParent(), programId.getProgram()), programRunId.getRun());
+  }
+}

cdap-app-fabric/src/main/java/io/cdap/cdap/internal/app/runtime/schedule/trigger/OrTrigger.java

@@ -18,7 +18,6 @@
 
 import io.cdap.cdap.api.schedule.TriggerInfo;
 import io.cdap.cdap.internal.app.runtime.schedule.ProgramSchedule;
-import io.cdap.cdap.proto.Notification;
 import io.cdap.cdap.proto.id.ProgramId;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -40,9 +39,9 @@ public OrTrigger(List<SatisfiableTrigger> triggers) {
   }
 
   @Override
-  public boolean isSatisfied(ProgramSchedule schedule, List<Notification> notifications) {
+  public boolean isSatisfied(ProgramSchedule schedule, NotificationContext notificationContext) {
     for (SatisfiableTrigger trigger : getTriggers()) {
-      if (trigger.isSatisfied(schedule, notifications)) {
+      if (trigger.isSatisfied(schedule, notificationContext)) {
         return true;
       }
     }

cdap-app-fabric/src/main/java/io/cdap/cdap/internal/app/runtime/schedule/trigger/PartitionTrigger.java

@@ -39,8 +39,8 @@ public PartitionTrigger(DatasetId dataset, int numPartitions) {
   }
 
   @Override
-  public boolean isSatisfied(ProgramSchedule schedule, List<Notification> notifications) {
-    return getPartitionsCount(notifications) >= numPartitions;
+  public boolean isSatisfied(ProgramSchedule schedule, NotificationContext notificationContext) {
+    return getPartitionsCount(notificationContext.getNotifications()) >= numPartitions;
   }
 
   private int getPartitionsCount(List<Notification> notifications) {

cdap-app-fabric/src/main/java/io/cdap/cdap/internal/app/runtime/schedule/trigger/ProgramStatusTrigger.java

@@ -23,8 +23,12 @@
 import com.google.gson.Gson;
 import io.cdap.cdap.api.ProgramStatus;
 import io.cdap.cdap.api.app.ProgramType;
+import io.cdap.cdap.api.retry.RetryableException;
 import io.cdap.cdap.api.schedule.TriggerInfo;
+import io.cdap.cdap.api.workflow.WorkflowToken;
 import io.cdap.cdap.common.app.RunIds;
+import io.cdap.cdap.common.service.Retries;
+import io.cdap.cdap.common.service.RetryStrategy;
 import io.cdap.cdap.internal.app.runtime.ProgramOptionConstants;
 import io.cdap.cdap.internal.app.runtime.schedule.ProgramSchedule;
 import io.cdap.cdap.internal.app.runtime.schedule.store.Schedulers;
@@ -33,12 +37,15 @@
 import io.cdap.cdap.proto.ProtoTrigger;
 import io.cdap.cdap.proto.id.ProgramId;
 import io.cdap.cdap.proto.id.ProgramRunId;
+import java.io.IOException;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * A Trigger that schedules a ProgramSchedule, when a certain status of a program has been
@@ -47,7 +54,12 @@
 public class ProgramStatusTrigger extends ProtoTrigger.ProgramStatusTrigger implements
     SatisfiableTrigger {
 
+  private static final String RESOLVED_PLUGIN_PROPERTIES_MAP = "resolved.plugin.properties.map";
+
   private static final Gson GSON = new Gson();
+  private static final Logger LOG =
+      LoggerFactory.getLogger(
+          io.cdap.cdap.internal.app.runtime.schedule.trigger.ProgramStatusTrigger.class);
 
   public ProgramStatusTrigger(ProgramId programId, Set<ProgramStatus> programStatuses) {
     super(programId, programStatuses);
@@ -59,13 +71,55 @@ public ProgramStatusTrigger(ProgramId programId, ProgramStatus... programStatuse
   }
 
   @Override
-  public boolean isSatisfied(ProgramSchedule schedule, List<Notification> notifications) {
-    return getTriggerSatisfiedResult(notifications, false, new Function<ProgramRunInfo, Boolean>() {
-      @Override
-      public Boolean apply(ProgramRunInfo input) {
+  public boolean isSatisfied(ProgramSchedule schedule, NotificationContext notificationContext) {
+
+    return getTriggerSatisfiedResult(notificationContext.getNotifications(), false,
+        runInfo -> {
+          if (!io.cdap.cdap.proto.ProgramType.WORKFLOW.equals(
+              runInfo.getProgramRunId().getType())) {
+            LOG.info(
+                "Program {} is of type '{}' and not 'Workflow', skipping workflow token validation",
+                runInfo.getProgramRunId(), runInfo.getProgramRunId().getType());
+            return true;
+          }
+          RetryStrategy retryStrategy = notificationContext.getRetryStrategy();
+          try {
+            // Retries are added because there may be delays while the workflow token is published.
+            return Retries.callWithRetries(
+                () -> fetchWorkflowToken(runInfo.getProgramRunId(), notificationContext),
+                retryStrategy);
+          } catch (RetryableException e) {
+            LOG.error("Retries exhausted for program runId {} with exception: ",
+                runInfo.getProgramRunId(), e);
+          }
+          return false;
+        });
+  }
+  /**
+   * Exception thrown when the workflow token is not found.
+   */
+  public static class WorkflowTokenNotFoundException extends RetryableException {
+
+  }
+
+  private boolean fetchWorkflowToken(ProgramRunId programRunId, NotificationContext context) {
+    try {
+      WorkflowToken workflowToken = context.getWorkflowToken(
+          programRunId);
+      if (workflowToken != null
+          && workflowToken.get(RESOLVED_PLUGIN_PROPERTIES_MAP) != null) {
+        // Return true only if workflow token has been recorded and resolved properties have
+        // been added.
         return true;
       }
-    });
+      LOG.warn("Retrying invalid workflow token \"{}\" for program runId {}", workflowToken,
+          programRunId);
+    } catch (IOException e) {
+      LOG.warn("Retrying read of workflow token failed for program runId {} with error:",
+          programRunId,
+          e);
+    }
+    throw new WorkflowTokenNotFoundException();
   }
 
   @Override

cdap-app-fabric/src/main/java/io/cdap/cdap/internal/app/runtime/schedule/trigger/SatisfiableTrigger.java

@@ -35,10 +35,11 @@ public interface SatisfiableTrigger extends Trigger {
    * it will remain satisfied no matter what new notifications it receives.
    *
    * @param schedule the schedule that this trigger belongs to
-   * @param notifications the notifications used to check whether this trigger is satisfied
+   * @param notificationContext that provides necessary information related to notifications
+   *     received.
    * @return {@code true} if this trigger is satisfied, {@code false} otherwise
    */
-  boolean isSatisfied(ProgramSchedule schedule, List<Notification> notifications);
+  boolean isSatisfied(ProgramSchedule schedule, NotificationContext notificationContext);
 
   /**
    * Get all trigger keys which will be used to index the schedule containing this trigger, so that

cdap-app-fabric/src/main/java/io/cdap/cdap/internal/app/runtime/schedule/trigger/TimeTrigger.java

@@ -55,8 +55,8 @@ public void validate() {
   }
 
   @Override
-  public boolean isSatisfied(ProgramSchedule schedule, List<Notification> notifications) {
-    for (Notification notification : notifications) {
+  public boolean isSatisfied(ProgramSchedule schedule, NotificationContext notificationContext) {
+    for (Notification notification : notificationContext.getNotifications()) {
       if (isSatisfied(schedule, notification)) {
         return true;
       }

cdap-app-fabric/src/main/java/io/cdap/cdap/internal/app/runtime/workflow/BasicWorkflowToken.java

@@ -342,6 +342,18 @@ private void writeObject(ObjectOutputStream out) throws IOException {
     out.defaultWriteObject();
   }
 
+  @Override
+  public String toString() {
+    return "BasicWorkflowToken{"
+        + "tokenValueMap=" + tokenValueMap
+        + ", maxSizeBytes=" + maxSizeBytes
+        + ", mapReduceCounters=" + mapReduceCounters
+        + ", nodeName='" + nodeName + '\''
+        + ", putAllowed=" + putAllowed
+        + ", bytesLeft=" + bytesLeft
+        + '}';
+  }
+
   // Deserialize the WorkflowToken for using it inside the Spark executor. Set the putAllowed
   // flag to false so that we do not allow putting the values inside the Spark executor.
   private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {

cdap-app-fabric/src/main/java/io/cdap/cdap/internal/app/store/AppMetadataStore.java

@@ -3030,7 +3030,7 @@ private String getFlowControlStatus(ProgramRunId programRunId, ProgramRunStatus
     if (!CONTROL_FLOW_PROGRAM_TYPES.contains(programRunId.getType())) {
       return TYPE_FLOW_CONTROL_NONE;
     }
-    if (programRunId.getParent().getNamespace() == NamespaceId.SYSTEM.getNamespace()) {
+    if (NamespaceId.SYSTEM.getNamespace().equals(programRunId.getParent().getNamespace())) {
       return TYPE_FLOW_CONTROL_NONE;
     }
     if (systemArgs.containsKey(ProgramOptionConstants.WORKFLOW_NAME)) {

cdap-app-fabric/src/test/java/io/cdap/cdap/AppWithMultipleSchedules.java

@@ -57,6 +57,7 @@ public class AppWithMultipleSchedules extends AbstractApplication {
   public static final String TRIGGERED_RUNTIME_ARG_KEY = "TriggeredWorkflowRuntimeArgKey";
   public static final String TRIGGERED_TOKEN_KEY = "TriggeredWorkflowTokenKey";
   public static final String TRIGGERING_PROPERTIES_MAPPING = "triggering.properties.mapping";
+  public static final String RESOLVED_PLUGIN_PROPERTIES_MAP = "resolved.plugin.properties.map";
 
   @Override
   public void configure() {
@@ -135,6 +136,7 @@ public void configure() {
     public void initialize(WorkflowContext context) throws Exception {
       super.initialize(context);
       context.getToken().put(ANOTHER_TOKEN_KEY, ANOTHER_TOKEN_VALUE);
+      context.getToken().put(RESOLVED_PLUGIN_PROPERTIES_MAP, "{}");
     }
   }