Skip to content

Commit

Permalink
Validate workflow token on program status trigger.
Browse files Browse the repository at this point in the history
  • Loading branch information
ritwiksahani committed Jan 23, 2025
1 parent 5730f51 commit ac0b41b
Show file tree
Hide file tree
Showing 11 changed files with 305 additions and 25 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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
Expand All @@ -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>
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
}
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,71 @@
/*
* 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.conf.CConfiguration;
import io.cdap.cdap.common.service.RetryStrategies;
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());
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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
Expand All @@ -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);
Expand All @@ -59,13 +71,46 @@ 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 -> {
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;
});
}

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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down Expand Up @@ -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, "{}");
}
}

Expand Down
Loading

0 comments on commit ac0b41b

Please sign in to comment.