Error management poc

cdap-api-common/pom.xml

@@ -43,6 +43,12 @@
       <groupId>junit</groupId>
       <artifactId>junit</artifactId>
     </dependency>
+    <dependency>
+      <groupId>io.cdap.cdap</groupId>
+      <artifactId>cdap-error-api</artifactId>
+      <version>6.11.0-SNAPSHOT</version>
+      <scope>compile</scope>
+    </dependency>
   </dependencies>
 
 </project>

cdap-api-common/src/main/java/io/cdap/cdap/api/exception/WrappedException.java

@@ -0,0 +1,58 @@
+/*
+ * Copyright © 2024 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.api.exception;
+
+import io.cdap.cdap.error.api.ErrorTagProvider;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Set;
+
+/**
+ * Sets the stage name in exception.
+ */
+public class WrappedException extends RuntimeException implements ErrorTagProvider {
+
+  private final String stageName;
+  private final Set<ErrorTag> errorTags = new HashSet<>();
+
+  public WrappedException(Throwable cause, String stageName) {
+    super(cause);
+    this.stageName = stageName;
+    this.errorTags.add(ErrorTag.PLUGIN);
+  }
+
+  public WrappedException(String message, String stageName) {
+    super(message);
+    this.stageName = stageName;
+    this.errorTags.add(ErrorTag.PLUGIN);
+  }
+
+  public WrappedException(Throwable cause, String message, String stageName) {
+    super(message, cause);
+    this.stageName = stageName;
+    this.errorTags.add(ErrorTag.PLUGIN);
+  }
+
+  public String getStageName() {
+    return stageName;
+  }
+
+  @Override
+  public Set<ErrorTag> getErrorTags() {
+    return Collections.unmodifiableSet(errorTags);
+  }
+}

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

@@ -15,15 +15,19 @@
  */
 package io.cdap.cdap.internal.app.runtime.workflow;
 
+import com.google.common.base.Throwables;
 import com.google.common.util.concurrent.Service;
+import io.cdap.cdap.api.exception.WrappedException;
 import io.cdap.cdap.internal.app.runtime.AbstractProgramController;
 import io.cdap.cdap.proto.id.ProgramId;
 import io.cdap.cdap.proto.id.ProgramRunId;
+import java.util.List;
 import org.apache.twill.api.RunId;
 import org.apache.twill.common.Threads;
 import org.apache.twill.internal.ServiceListenerAdapter;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.slf4j.MDC;
 
 /**
  *
@@ -87,6 +91,15 @@ public void terminated(Service.State from) {
 
       @Override
       public void failed(Service.State from, Throwable failure) {
+        List<Throwable> causalChain = Throwables.getCausalChain(failure);
+        for(Throwable cause : causalChain) {
+          if (cause instanceof WrappedException) {
+            String stageName = ((WrappedException) cause).getStageName();
+            LOG.error("Stage: {}", stageName);
+            MDC.put("Failed_Stage", stageName);
+            break;
+          }
+        }
         LOG.error("Workflow service '{}' failed.", serviceName, failure);
         error(failure);
       }

cdap-app-templates/cdap-etl/cdap-etl-core/src/main/java/io/cdap/cdap/etl/batch/WrappedInputFormat.java

@@ -0,0 +1,86 @@
+/*
+ * Copyright © 2024 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.etl.batch;
+
+import io.cdap.cdap.api.exception.WrappedException;
+import java.io.IOException;
+import java.util.List;
+import org.apache.hadoop.mapreduce.InputFormat;
+import org.apache.hadoop.mapreduce.InputSplit;
+import org.apache.hadoop.mapreduce.JobContext;
+import org.apache.hadoop.mapreduce.RecordReader;
+import org.apache.hadoop.mapreduce.TaskAttemptContext;
+
+public class WrappedInputFormat<K, V> extends InputFormat<K, V> {
+  private final InputFormat<K, V> inputFormat;
+  private final String stageName;
+
+  @Override
+  public List<InputSplit> getSplits(JobContext jobContext)
+      throws IOException, InterruptedException {
+    try {
+      return inputFormat.getSplits(jobContext);
+    } catch (Exception e) {
+      if (stageName != null) {
+        throw new WrappedException(e, stageName);
+      }
+      throw e;
+    }
+  }
+
+  @Override
+  public RecordReader<K, V> createRecordReader(InputSplit inputSplit,
+      TaskAttemptContext taskAttemptContext) throws IOException, InterruptedException {
+    try {
+      return new WrappedRecordReader<>(inputFormat.createRecordReader(inputSplit,
+          taskAttemptContext), stageName);
+    } catch (Exception e) {
+      if (stageName != null) {
+        throw new WrappedException(e, stageName);
+      }
+      throw e;
+    }
+  }
+
+  /**
+   * Returns the delegating {@link InputFormat} based on the current configuration.
+   *
+   * @param classLoader the {@link ClassLoader} for loading input format
+   * @param inputFormatClassName the name of {@link InputFormat} class
+   * @throws IOException if failed to instantiate the input format class
+   */
+  public WrappedInputFormat(ClassLoader classLoader, String inputFormatClassName,
+      String stageName) throws IOException {
+    this.stageName = stageName;
+    if (inputFormatClassName == null) {
+      throw new IllegalArgumentException("Missing configuration for the InputFormat to use");
+    }
+    if (inputFormatClassName.equals(getClass().getName())) {
+      throw new IllegalArgumentException("Cannot delegate InputFormat to the same class");
+    }
+    try {
+      //noinspection unchecked
+      @SuppressWarnings("unchecked")
+      Class<InputFormat<K, V>> inputFormatClass = (Class<InputFormat<K, V>>) classLoader.loadClass(
+          inputFormatClassName);
+      this.inputFormat = inputFormatClass.newInstance();
+    } catch (InstantiationException | IllegalAccessException | ClassNotFoundException e) {
+      throw new IOException(
+          String.format("Unable to instantiate delegate input format %s", inputFormatClassName), e);
+    }
+  }
+}

cdap-app-templates/cdap-etl/cdap-etl-core/src/main/java/io/cdap/cdap/etl/batch/WrappedRecordReader.java

@@ -0,0 +1,106 @@
+/*
+ * Copyright © 2024 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.etl.batch;
+
+import io.cdap.cdap.api.exception.WrappedException;
+import java.io.IOException;
+import org.apache.hadoop.mapreduce.InputSplit;
+import org.apache.hadoop.mapreduce.RecordReader;
+import org.apache.hadoop.mapreduce.TaskAttemptContext;
+
+public class WrappedRecordReader<K, V> extends RecordReader<K, V> {
+  private  final RecordReader<K, V> recordReader;
+  private final String stageName;
+
+  public WrappedRecordReader(RecordReader<K, V> recordReader, String stageName) {
+    this.recordReader = recordReader;
+    this.stageName = stageName;
+  }
+
+  @Override
+  public void initialize(InputSplit inputSplit, TaskAttemptContext taskAttemptContext)
+      throws IOException, InterruptedException {
+    try {
+      recordReader.initialize(inputSplit, taskAttemptContext);
+    } catch (Exception e) {
+      if (stageName != null) {
+        throw new WrappedException(e, stageName);
+      }
+      throw e;
+    }
+  }
+
+  @Override
+  public boolean nextKeyValue() throws IOException, InterruptedException {
+    try {
+      return recordReader.nextKeyValue();
+    } catch (Exception e) {
+      if (stageName != null) {
+        throw new WrappedException(e, stageName);
+      }
+      throw e;
+    }
+  }
+
+  @Override
+  public K getCurrentKey() throws IOException, InterruptedException {
+    try {
+      return recordReader.getCurrentKey();
+    } catch (Exception e) {
+      if (stageName != null) {
+        throw new WrappedException(e, stageName);
+      }
+      throw e;
+    }
+  }
+
+  @Override
+  public V getCurrentValue() throws IOException, InterruptedException {
+    try {
+      return recordReader.getCurrentValue();
+    } catch (Exception e) {
+      if (stageName != null) {
+        throw new WrappedException(e, stageName);
+      }
+      throw e;
+    }
+  }
+
+  @Override
+  public float getProgress() throws IOException, InterruptedException {
+    try {
+      return recordReader.getProgress();
+    } catch (Exception e) {
+      if (stageName != null) {
+        throw new WrappedException(e, stageName);
+      }
+      throw e;
+    }
+  }
+
+  @Override
+  public void close() throws IOException {
+    try {
+      recordReader.close();
+    } catch (Exception e) {
+      if (stageName != null) {
+        throw new WrappedException(e, stageName);
+      }
+      throw e;
+    }
+  }
+}

cdap-app-templates/cdap-etl/cdap-etl-core/src/main/java/io/cdap/cdap/etl/common/plugin/StageLoggingCaller.java

@@ -33,6 +33,10 @@ private StageLoggingCaller(Caller delegate, String stageName) {
     this.stageName = stageName;
   }
 
+  public String getStageName() {
+    return stageName;
+  }
+
   @Override
   public <T> T call(Callable<T> callable) throws Exception {
     MDC.put(Constants.MDC_STAGE_KEY, stageName);

cdap-app-templates/cdap-etl/cdap-etl-core/src/main/java/io/cdap/cdap/etl/common/plugin/WrappedAction.java

@@ -16,10 +16,14 @@
 
 package io.cdap.cdap.etl.common.plugin;
 
+import io.cdap.cdap.api.exception.WrappedException;
 import io.cdap.cdap.etl.api.PipelineConfigurer;
 import io.cdap.cdap.etl.api.action.Action;
 import io.cdap.cdap.etl.api.action.ActionContext;
 import java.util.concurrent.Callable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.slf4j.MDC;
 
 /**
  * Wrapper around {@link Action} that makes sure logging, classloading, and other pipeline
@@ -29,6 +33,7 @@ public class WrappedAction extends Action implements PluginWrapper<Action> {
 
   private final Action action;
   private final Caller caller;
+  private static final Logger LOG = LoggerFactory.getLogger(WrappedAction.class);
 
   public WrappedAction(Action action, Caller caller) {
     this.action = action;
@@ -45,10 +50,20 @@ public void configurePipeline(final PipelineConfigurer pipelineConfigurer) {
 
   @Override
   public void run(final ActionContext context) throws Exception {
-    caller.call((Callable<Void>) () -> {
-      action.run(context);
-      return null;
-    });
+    try {
+      caller.call((Callable<Void>) () -> {
+        action.run(context);
+        return null;
+      });
+    } catch(Exception e) {
+      if (caller instanceof StageLoggingCaller) {
+        String stageName = ((StageLoggingCaller) caller).getStageName();
+        MDC.put("Failed_Stage", stageName);
+        LOG.error("Stage: {}", stageName);
+        throw new WrappedException(e, stageName);
+      }
+      throw e;
+    }
   }
 
   @Override