TEZ-4737: Session DAG submission fails with FileAlreadyExistsException when the session outlives the TezClient

tez-api/src/main/java/org/apache/tez/client/TezClient.java

@@ -688,11 +688,11 @@ private DAGClient submitDAGSession(DAG dag) throws TezException, IOException {
     SubmitDAGRequestProto request = requestBuilder.build();
     if (request.getSerializedSize() > maxSubmitDAGRequestSizeThroughIPC) {
       Path dagPlanPath = new Path(TezCommonUtils.getTezSystemStagingPath(amConfig.getTezConfiguration(),
-          sessionAppId.toString()), TezConstants.TEZ_PB_PLAN_BINARY_NAME +
-          serializedSubmitDAGPlanRequestCounter.incrementAndGet());
+          sessionAppId.toString()), TezConstants.TEZ_PB_PLAN_BINARY_NAME_FORMAT.formatted(
+          serializedSubmitDAGPlanRequestCounter.incrementAndGet()));
 
       FileSystem fs = dagPlanPath.getFileSystem(stagingFs.getConf());
-      try (FSDataOutputStream fsDataOutputStream = fs.create(dagPlanPath, false)) {
+      try (FSDataOutputStream fsDataOutputStream = fs.create(dagPlanPath, true)) {
         LOG.info("Send dag plan using YARN local resources since it's too large"
             + ", dag plan size=" + request.getSerializedSize()
             + ", max dag plan size through IPC=" + maxSubmitDAGRequestSizeThroughIPC

tez-api/src/main/java/org/apache/tez/dag/api/TezConstants.java

@@ -48,6 +48,7 @@ public final class TezConstants {
   public static final String SERVICE_PLUGINS_DESCRIPTOR_JSON = "service_plugins_descriptor.json";
   public static final String TEZ_PB_BINARY_CONF_NAME = "tez-conf.pb";
   public static final String TEZ_PB_PLAN_BINARY_NAME = "tez-dag.pb";
+  public static final String TEZ_PB_PLAN_BINARY_NAME_FORMAT = "tez-dag-%d.pb";
   public static final String TEZ_PB_PLAN_TEXT_NAME = "tez-dag.pb.txt";
 
   /*

tez-api/src/test/java/org/apache/tez/client/TestTezClient.java

@@ -70,6 +70,7 @@
 import org.apache.hadoop.yarn.client.api.YarnClient;
 import org.apache.hadoop.yarn.exceptions.ApplicationNotFoundException;
 import org.apache.hadoop.yarn.exceptions.YarnException;
+import org.apache.tez.common.TezCommonUtils;
 import org.apache.tez.common.counters.LimitExceededException;
 import org.apache.tez.common.counters.Limits;
 import org.apache.tez.common.counters.TezCounters;
@@ -269,20 +270,12 @@ private void _testTezClientSessionLargeDAGPlan(int maxIPCMsgSize, int payloadSiz
     localResourceMap.put(lrName, LocalResource.newInstance(URL.newInstance("file", "localhost", 0, "/test"),
         LocalResourceType.FILE, LocalResourceVisibility.PUBLIC, 1, 1));
 
-    ProcessorDescriptor processorDescriptor = ProcessorDescriptor.create("P");
-    processorDescriptor.setUserPayload(UserPayload.create(ByteBuffer.allocate(payloadSize)));
-    Vertex vertex = Vertex.create("Vertex", processorDescriptor, 1, Resource.newInstance(1, 1));
-    DAG dag = DAG.create("DAG-testTezClientSessionLargeDAGPlan").addVertex(vertex);
-
     client.start();
     client.addAppMasterLocalFiles(localResourceMap);
-    client.submitDAG(dag);
+    SubmitDAGRequestProto request =
+        submitDAGAndCaptureRequest(client, largeDAG("DAG-testTezClientSessionLargeDAGPlan", payloadSize));
     client.stop();
 
-    ArgumentCaptor<SubmitDAGRequestProto> captor = ArgumentCaptor.forClass(SubmitDAGRequestProto.class);
-    verify(client.sessionAmProxy).submitDAG(any(), captor.capture());
-    SubmitDAGRequestProto request = captor.getValue();
-
     if (shouldSerialize) {
       /* we need manually delete the serialized dagplan since staging path here won't be destroyed */
       Path dagPlanPath = new Path(request.getSerializedRequestPath());
@@ -300,6 +293,55 @@ private void _testTezClientSessionLargeDAGPlan(int maxIPCMsgSize, int payloadSiz
     }
   }
 
+  @Test
+  @Timeout(value = 10000, unit = TimeUnit.MILLISECONDS)
+  public void testSessionLargeDAGPlanWithLeftoverPlanFile() throws Exception {
+    // The session outlives its TezClient instances: a plan file left behind by a previous
+    // client must not fail subsequent submissions.
+    TezConfiguration conf = new TezConfiguration();
+    conf.setInt(CommonConfigurationKeys.IPC_MAXIMUM_DATA_LENGTH, 1024 * 1024);
+    conf.set(TezConfiguration.TEZ_AM_STAGING_DIR, STAGING_DIR.getAbsolutePath());
+
+    Path baseStagingPath = TezCommonUtils.getTezBaseStagingPath(conf);
+    FileSystem localFs = FileSystem.getLocal(conf);
+    localFs.delete(baseStagingPath, true);
+
+    // client1 writes tez-dag-1.pb and is abandoned without stop()
+    TezClientForTest client1 = configureAndCreateTezClient(null, true, conf);
+    client1.start();
+    SubmitDAGRequestProto request1 = submitDAGAndCaptureRequest(client1, largeDAG("DAG-client1", 2 * 1024 * 1024));
+    assertTrue(request1.hasSerializedRequestPath());
+    // the mocked AM never consumes the plan file, leaving it behind in the staging dir
+    Path leftoverPlanPath = new Path(request1.getSerializedRequestPath());
+    assertTrue(localFs.exists(leftoverPlanPath));
+
+    // client2 reuses the same session and computes the same tez-dag-1.pb path
+    TezClientForTest client2 = configureAndCreateTezClient(null, true, conf);
+    client2.start();
+    SubmitDAGRequestProto request2 = submitDAGAndCaptureRequest(client2, largeDAG("DAG-client2", 2 * 1024 * 1024));
+    assertTrue(request2.hasSerializedRequestPath());
+    // path collision: client2 overwrote client1's leftover instead of failing
+    assertEquals(leftoverPlanPath, new Path(request2.getSerializedRequestPath()));
+    assertTrue(localFs.exists(leftoverPlanPath));
+    client2.stop();
+
+    localFs.delete(baseStagingPath, true);
+  }
+
+  private DAG largeDAG(String name, int payloadSize) {
+    ProcessorDescriptor processorDescriptor = ProcessorDescriptor.create("P");
+    processorDescriptor.setUserPayload(UserPayload.create(ByteBuffer.allocate(payloadSize)));
+    Vertex vertex = Vertex.create("Vertex", processorDescriptor, 1, Resource.newInstance(1, 1));
+    return DAG.create(name).addVertex(vertex);
+  }
+
+  private SubmitDAGRequestProto submitDAGAndCaptureRequest(TezClientForTest client, DAG dag) throws Exception {
+    client.submitDAG(dag);
+    ArgumentCaptor<SubmitDAGRequestProto> captor = ArgumentCaptor.forClass(SubmitDAGRequestProto.class);
+    verify(client.sessionAmProxy).submitDAG(any(), captor.capture());
+    return captor.getValue();
+  }
+
   @Test
   @Timeout(value = 5000, unit = TimeUnit.MILLISECONDS)
   public void testGetClient() throws Exception {

tez-dag/src/main/java/org/apache/tez/dag/api/client/rpc/DAGClientAMProtocolBlockingPBServerImpl.java

@@ -58,8 +58,13 @@
 import com.google.protobuf.RpcController;
 import com.google.protobuf.ServiceException;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 public class DAGClientAMProtocolBlockingPBServerImpl implements DAGClientAMProtocolBlockingPB {
 
+  private static final Logger LOG = LoggerFactory.getLogger(DAGClientAMProtocolBlockingPBServerImpl.class);
+
   DAGClientHandler real;
   final FileSystem stagingFs;
 
@@ -176,6 +181,12 @@ public SubmitDAGResponseProto submitDAG(RpcController controller,
           request = SubmitDAGRequestProto.parseFrom(in);
         } catch (IOException e) {
           throw wrapException(e);
+        } finally {
+          try {
+            fs.delete(requestPath, false);
+          } catch (IOException e) {
+            LOG.warn("Failed to delete the serialized DAG plan file {}", requestPath, e);
+          }
         }
       }
       DAGPlan dagPlan = request.getDAGPlan();

tez-dag/src/test/java/org/apache/tez/dag/api/client/rpc/TestDAGClientAMProtocolBlockingPBServerImpl.java

@@ -20,6 +20,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.mock;
@@ -118,6 +119,9 @@ public void testSubmitDagInSessionWithLargeDagPlan() throws Exception {
     requestBuilder.clear().setSerializedRequestPath(requestFile.getAbsolutePath());
     serverImpl.submitDAG(null, requestBuilder.build());
 
+    assertFalse(requestFile.exists(),
+        "The serialized DAG plan file should be deleted once the AM has consumed it");
+
     ArgumentCaptor<DAGPlan> dagPlanCaptor = ArgumentCaptor.forClass(DAGPlan.class);
     verify(dagClientHandler).submitDAG(dagPlanCaptor.capture(), localResourcesCaptor.capture());
     dagPlan = dagPlanCaptor.getValue();