From e67c2e35ad8f0e37c98a6e1e012d6f7b95dc1fe1 Mon Sep 17 00:00:00 2001 From: bvolpato Date: Tue, 4 Aug 2026 22:21:20 -0400 Subject: [PATCH] [examples] Atomically publish subprocess executables --- .../examples/subprocess/utils/FileUtils.java | 30 ++++--- .../subprocess/utils/FileUtilsTest.java | 80 +++++++++++++++++++ 2 files changed, 99 insertions(+), 11 deletions(-) create mode 100644 examples/java/src/test/java/org/apache/beam/examples/subprocess/utils/FileUtilsTest.java diff --git a/examples/java/src/main/java/org/apache/beam/examples/subprocess/utils/FileUtils.java b/examples/java/src/main/java/org/apache/beam/examples/subprocess/utils/FileUtils.java index d0244d233722..eff913b9f4fb 100644 --- a/examples/java/src/main/java/org/apache/beam/examples/subprocess/utils/FileUtils.java +++ b/examples/java/src/main/java/org/apache/beam/examples/subprocess/utils/FileUtils.java @@ -28,6 +28,8 @@ import java.nio.file.Files; import java.nio.file.Path; import java.nio.file.Paths; +import java.nio.file.StandardCopyOption; +import java.util.UUID; import org.apache.beam.examples.subprocess.configuration.SubProcessConfiguration; import org.apache.beam.sdk.io.FileSystems; import org.apache.beam.sdk.io.fs.ResolveOptions.StandardResolveOptions; @@ -77,30 +79,36 @@ public static String copyFileFromWorkerToGCS( } } - public static String copyFileFromGCSToWorker(ExecutableFile execuableFile) throws Exception { + public static String copyFileFromGCSToWorker(ExecutableFile executableFile) throws Exception { ResourceId sourceFile = - FileSystems.matchNewResource(execuableFile.getSourceGCSLocation(), false); - ResourceId destinationFile = - FileSystems.matchNewResource(execuableFile.getDestinationLocation(), false); + FileSystems.matchNewResource(executableFile.getSourceGCSLocation(), false); try { LOG.info( "Moving File {} to {} ", - execuableFile.getSourceGCSLocation(), - execuableFile.getDestinationLocation()); - Path path = Paths.get(execuableFile.getDestinationLocation()); + executableFile.getSourceGCSLocation(), + executableFile.getDestinationLocation()); + Path path = Paths.get(executableFile.getDestinationLocation()); if (path.toFile().exists()) { LOG.warn( "Overwriting file {}, should only see this once per worker.", - execuableFile.getDestinationLocation()); + executableFile.getDestinationLocation()); + } + Path stagedFile = path.resolveSibling(".beam-executable-" + UUID.randomUUID() + ".tmp"); + try { + ResourceId stagedResource = FileSystems.matchNewResource(stagedFile.toString(), false); + copyFile(sourceFile, stagedResource); + stagedFile.toFile().setExecutable(true); + Files.move( + stagedFile, path, StandardCopyOption.REPLACE_EXISTING, StandardCopyOption.ATOMIC_MOVE); + } finally { + Files.deleteIfExists(stagedFile); } - copyFile(sourceFile, destinationFile); - path.toFile().setExecutable(true); return path.toString(); } catch (Exception ex) { - LOG.error("Error moving file : {} ", execuableFile.fileName, ex); + LOG.error("Error moving file : {} ", executableFile.fileName, ex); throw ex; } } diff --git a/examples/java/src/test/java/org/apache/beam/examples/subprocess/utils/FileUtilsTest.java b/examples/java/src/test/java/org/apache/beam/examples/subprocess/utils/FileUtilsTest.java new file mode 100644 index 000000000000..02c23d705c27 --- /dev/null +++ b/examples/java/src/test/java/org/apache/beam/examples/subprocess/utils/FileUtilsTest.java @@ -0,0 +1,80 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.beam.examples.subprocess.utils; + +import static java.nio.charset.StandardCharsets.UTF_8; +import static java.nio.file.StandardOpenOption.WRITE; +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; +import static org.junit.Assume.assumeTrue; + +import java.io.File; +import java.nio.channels.FileChannel; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.attribute.PosixFilePermission; +import java.util.Set; +import java.util.stream.Stream; +import org.apache.beam.examples.subprocess.configuration.SubProcessConfiguration; +import org.apache.commons.lang3.SystemUtils; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class FileUtilsTest { + @Rule public TemporaryFolder temporaryFolder = new TemporaryFolder(); + + @Test + public void copyFileFromGCSToWorkerAtomicallyReplacesExecutable() throws Exception { + assumeTrue(SystemUtils.IS_OS_LINUX); + + File sourceDirectory = temporaryFolder.newFolder("source"); + File workerDirectory = temporaryFolder.newFolder("worker"); + String fileName = "echo.sh"; + Path source = sourceDirectory.toPath().resolve(fileName); + Path destination = workerDirectory.toPath().resolve(fileName); + Files.write(source, "#!/bin/sh\nexit 0\n".getBytes(UTF_8)); + Files.write(destination, "#!/bin/sh\nexit 1\n".getBytes(UTF_8)); + assertTrue(destination.toFile().setExecutable(true)); + Set destinationPermissions = Files.getPosixFilePermissions(destination); + + SubProcessConfiguration configuration = new SubProcessConfiguration(); + configuration.setSourcePath(sourceDirectory.getAbsolutePath()); + configuration.setWorkerPath(workerDirectory.getAbsolutePath()); + + try (FileChannel ignored = FileChannel.open(destination, WRITE)) { + String copiedFile = + FileUtils.copyFileFromGCSToWorker(new ExecutableFile(configuration, fileName)); + + assertEquals(destination.toString(), copiedFile); + assertArrayEquals(Files.readAllBytes(source), Files.readAllBytes(destination)); + assertEquals(destinationPermissions, Files.getPosixFilePermissions(destination)); + assertTrue(Files.isExecutable(destination)); + try (Stream files = Files.list(workerDirectory.toPath())) { + assertEquals(1, files.count()); + } + + Process process = new ProcessBuilder(destination.toString()).start(); + assertEquals(0, process.waitFor()); + } + } +}