Support atomic writes with FileWriter

3 files changed, 47 insertions, 8 deletions

diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/file/FileSync.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/file/FileSync.java
index 78f720074dc..dc13ea1c9ab 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/file/FileSync.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/file/FileSync.java
@@ -4,11 +4,15 @@ package com.yahoo.vespa.hosted.node.admin.task.util.file;
 
 import com.yahoo.vespa.hosted.node.admin.component.TaskContext;
 
+import java.nio.file.FileSystems;
+import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Arrays;
 import java.util.Optional;
 import java.util.logging.Logger;
 
+import static com.yahoo.yolean.Exceptions.uncheck;
+
 /**
  * Class to minimize resource usage with repetitive and mostly identical, idempotent, and
  * mutating file operations, e.g. setting file content, setting owner, etc.
@@ -29,15 +33,22 @@ public class FileSync {
         this.contentCache = new FileContentCache(this.path);
     }
 
+    public boolean convergeTo(TaskContext taskContext, PartialFileData partialFileData) {
+        return convergeTo(taskContext, partialFileData, false);
+    }
+
     /**
      * CPU, I/O, and memory usage is optimized for repeated calls with the same arguments.
+     *
+     * @param atomicWrite Whether to write updates to a temporary file in the same directory, and atomically move it
+     *                    to path. Ensures the file cannot be read while in the middle of writing it.
      * @return true if the system was modified: content was written, or owner was set, etc.
      *         system is only modified if necessary (different).
      */
-    public boolean convergeTo(TaskContext taskContext, PartialFileData partialFileData) {
+    public boolean convergeTo(TaskContext taskContext, PartialFileData partialFileData, boolean atomicWrite) {
         FileAttributesCache currentAttributes = new FileAttributesCache(path);
 
-        boolean modifiedSystem = maybeUpdateContent(taskContext, partialFileData.getContent(), currentAttributes);
+        boolean modifiedSystem = maybeUpdateContent(taskContext, partialFileData.getContent(), currentAttributes, atomicWrite);
 
         AttributeSync attributeSync = new AttributeSync(path.toPath()).with(partialFileData);
         modifiedSystem |= attributeSync.converge(taskContext, currentAttributes);
@@ -47,7 +58,8 @@ public class FileSync {
 
     private boolean maybeUpdateContent(TaskContext taskContext,
                                        Optional<byte[]> content,
-                                       FileAttributesCache currentAttributes) {
+                                       FileAttributesCache currentAttributes,
+                                       boolean atomicWrite) {
         if (!content.isPresent()) {
             return false;
         }
@@ -55,7 +67,7 @@ public class FileSync {
         if (!currentAttributes.exists()) {
             taskContext.recordSystemModification(logger, "Creating file " + path);
             path.createParents();
-            path.writeBytes(content.get());
+            writeBytes(content.get(), atomicWrite);
             contentCache.updateWith(content.get(), currentAttributes.forceGet().lastModifiedTime());
             return true;
         }
@@ -64,9 +76,20 @@ public class FileSync {
             return false;
         } else {
             taskContext.recordSystemModification(logger, "Patching file " + path);
-            path.writeBytes(content.get());
+            writeBytes(content.get(), atomicWrite);
             contentCache.updateWith(content.get(), currentAttributes.forceGet().lastModifiedTime());
             return true;
         }
     }
+
+    private void writeBytes(byte[] content, boolean atomic) {
+        if (atomic) {
+            String tmpPath = path.toPath().toString() + ".FileSyncTmp";
+            new UnixPath(path.toPath().getFileSystem().getPath(tmpPath))
+                    .writeBytes(content)
+                    .atomicMove(path.toPath());
+        } else {
+            path.writeBytes(content);
+        }
+    }
 }
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/file/FileWriter.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/file/FileWriter.java
index afc0e7b5c22..57f3417b789 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/file/FileWriter.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/file/FileWriter.java
@@ -3,7 +3,6 @@ package com.yahoo.vespa.hosted.node.admin.task.util.file;
 
 import com.yahoo.vespa.hosted.node.admin.component.TaskContext;
 
-import java.io.File;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -21,6 +20,7 @@ public class FileWriter {
     private final PartialFileData.Builder fileDataBuilder = PartialFileData.builder();
     private final Optional<ByteArraySupplier> contentProducer;
 
+    private boolean atomicWrite = false;
     private boolean overwriteExistingFile = true;
 
     public FileWriter(Path path) {
@@ -58,6 +58,11 @@ public class FileWriter {
         return this;
     }
 
+    public FileWriter atomicWrite(boolean atomicWrite) {
+        this.atomicWrite = atomicWrite;
+        return this;
+    }
+
     public FileWriter onlyIfFileDoesNotAlreadyExist() {
         overwriteExistingFile = false;
         return this;
@@ -78,7 +83,7 @@ public class FileWriter {
 
         fileDataBuilder.withContent(content);
         PartialFileData fileData = fileDataBuilder.create();
-        return fileSync.convergeTo(context, fileData);
+        return fileSync.convergeTo(context, fileData, atomicWrite);
     }
 
     @FunctionalInterface
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/task/util/file/FileWriterTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/task/util/file/FileWriterTest.java
index 2bc64a3fdb3..e1a0ed5d972 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/task/util/file/FileWriterTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/task/util/file/FileWriterTest.java
@@ -21,6 +21,7 @@ import static org.mockito.Mockito.verify;
 
 public class FileWriterTest {
     private final FileSystem fileSystem = TestFileSystem.create();
+    private final TaskContext context = mock(TaskContext.class);
 
     @Test
     public void testWrite() {
@@ -35,7 +36,6 @@ public class FileWriterTest {
                 .withOwner(owner)
                 .withGroup(group)
                 .onlyIfFileDoesNotAlreadyExist();
-        TaskContext context = mock(TaskContext.class);
         assertTrue(writer.converge(context));
         verify(context, times(1)).recordSystemModification(any(), eq("Creating file " + path));
 
@@ -50,4 +50,15 @@ public class FileWriterTest {
         assertFalse(writer.converge(context));
         assertEquals(fileTime, unixPath.getLastModifiedTime());
     }
+
+    @Test
+    public void testAtomicWrite() {
+        FileWriter writer = new FileWriter(fileSystem.getPath("/foo/bar"))
+                .atomicWrite(true);
+
+        assertTrue(writer.converge(context, "content"));
+
+        verify(context).recordSystemModification(any(), eq("Creating file /foo/bar"));
+        assertEquals("content", new UnixPath(writer.path()).readUtf8File());
+    }
 }