Write data to staging files before adding them to the output directory

Author: blootsvoets

src/main/java/org/radarcns/RestructureAvroRecords.java

@@ -28,11 +28,11 @@
 import org.apache.hadoop.fs.LocatedFileStatus;
 import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.fs.RemoteIterator;
-import org.radarcns.util.CsvAvroConverter;
-import org.radarcns.util.FileCacheStore;
-import org.radarcns.util.JsonAvroConverter;
+import org.radarcns.data.CsvAvroConverter;
+import org.radarcns.data.FileCacheStore;
+import org.radarcns.data.JsonAvroConverter;
 import org.radarcns.util.ProgressBar;
-import org.radarcns.util.RecordConverterFactory;
+import org.radarcns.data.RecordConverterFactory;
 import org.radarcns.util.commandline.CommandLineArgs;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -302,9 +302,9 @@ private void writeRecord(GenericRecord record, String topicName, FileCacheStore
         java.nio.file.Path outputPath = userTopicDir.resolve(outputFileName);
 
         // Write data
-        int response = cache.writeRecord(outputPath, record);
+        FileCacheStore.WriteStatus response = cache.writeRecord(outputPath, record);
 
-        if (response == FileCacheStore.CACHE_AND_NO_WRITE || response == FileCacheStore.NO_CACHE_AND_NO_WRITE) {
+        if (!response.isSuccessful()) {
             // Write was unsuccessful due to different number of columns,
             // try again with new file name
             writeRecord(record, topicName, cache, ++suffix);

src/main/java/org/radarcns/util/CsvAvroConverter.java renamed to src/main/java/org/radarcns/data/CsvAvroConverter.java

@@ -14,9 +14,8 @@
  * limitations under the License.
  */
 
-package org.radarcns.util;
+package org.radarcns.data;
 
-import com.fasterxml.jackson.databind.JsonMappingException;
 import com.fasterxml.jackson.databind.MappingIterator;
 import com.fasterxml.jackson.databind.ObjectReader;
 import com.fasterxml.jackson.databind.ObjectWriter;
@@ -34,7 +33,10 @@
 import java.io.Reader;
 import java.io.Writer;
 import java.nio.ByteBuffer;
-import java.util.*;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
 
 /**
  * Converts deep hierarchical Avro records into flat CSV format. It uses a simple dot syntax in the

src/main/java/org/radarcns/util/FileCache.java renamed to src/main/java/org/radarcns/data/FileCache.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package org.radarcns.util;
+package org.radarcns.data;
 
 import java.io.*;
 import java.nio.file.Files;
@@ -27,13 +27,17 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import static java.nio.file.StandardCopyOption.REPLACE_EXISTING;
+
 /** Keeps path handles of a path. */
 public class FileCache implements Closeable, Flushable, Comparable<FileCache> {
     private static final Logger logger = LoggerFactory.getLogger(FileCache.class);
+    private static final int BUFFER_SIZE = 8192;
 
     private final Writer writer;
     private final RecordConverter recordConverter;
     private final Path path;
+    private final Path tmpPath;
     private long lastUse;
 
     /**
@@ -45,19 +49,29 @@ public class FileCache implements Closeable, Flushable, Comparable<FileCache> {
      * @throws IOException
      */
     public FileCache(RecordConverterFactory converterFactory, Path path,
-            GenericRecord record, boolean gzip) throws IOException {
+            GenericRecord record, boolean gzip, Path tmpDir) throws IOException {
         this.path = path;
         boolean fileIsNew = !Files.exists(path) || Files.size(path) == 0;
+        this.tmpPath = Files.createTempFile(tmpDir, path.getFileName().toString(),
+                gzip ? ".tmp.gz" : ".tmp");
 
-        OutputStream outFile = Files.newOutputStream(path,
-                StandardOpenOption.APPEND, StandardOpenOption.CREATE);
-        InputStream inputStream = new BufferedInputStream(Files.newInputStream(path));
+        OutputStream outFile = Files.newOutputStream(tmpPath, StandardOpenOption.WRITE);
         OutputStream bufOut = new BufferedOutputStream(outFile);
         if (gzip) {
             bufOut = new GZIPOutputStream(bufOut);
-            if (!fileIsNew) {
+        }
+
+        InputStream inputStream;
+        if (fileIsNew) {
+            inputStream = new ByteArrayInputStream(new byte[0]);
+        } else {
+            inputStream = new BufferedInputStream(Files.newInputStream(path));
+            InputStream copyStream = Files.newInputStream(path);
+            if (gzip) {
+                copyStream = new GZIPInputStream(copyStream);
                 inputStream = new GZIPInputStream(inputStream);
             }
+            copy(copyStream, bufOut);
         }
 
         this.writer = new OutputStreamWriter(bufOut);
@@ -78,7 +92,7 @@ public FileCache(RecordConverterFactory converterFactory, Path path,
      * Write a record to the cache.
      * @param record AVRO record
      * @return true or false based on {@link RecordConverter} write result
-     * @throws IOException
+     * @throws IOException if the record cannot be used.
      */
     public boolean writeRecord(GenericRecord record) throws IOException {
         boolean result = this.recordConverter.writeRecord(record);
@@ -90,6 +104,7 @@ public boolean writeRecord(GenericRecord record) throws IOException {
     public void close() throws IOException {
         recordConverter.close();
         writer.close();
+        Files.move(tmpPath, path, REPLACE_EXISTING);
     }
 
     @Override
@@ -115,4 +130,16 @@ public int compareTo(@Nonnull FileCache other) {
     public Path getPath() {
         return path;
     }
+
+
+    /**
+     * Reads all bytes from an input stream and writes them to an output stream.
+     */
+    private static void copy(InputStream source, OutputStream sink) throws IOException {
+        byte[] buf = new byte[BUFFER_SIZE];
+        int n;
+        while ((n = source.read(buf)) > 0) {
+            sink.write(buf, 0, n);
+        }
+    }
 }

src/main/java/org/radarcns/util/FileCacheStore.java renamed to src/main/java/org/radarcns/data/FileCacheStore.java

@@ -14,9 +14,11 @@
  * limitations under the License.
  */
 
-package org.radarcns.util;
+package org.radarcns.data;
 
 import org.apache.avro.generic.GenericRecord;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.io.Closeable;
 import java.io.Flushable;
@@ -25,34 +27,34 @@
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.HashMap;
 import java.util.Map;
 
+import static org.radarcns.util.ThrowingConsumer.tryCatch;
+
 /**
  * Caches open file handles. If more than the limit is cached, the half of the files that were used
  * the longest ago cache are evicted from cache.
  */
 public class FileCacheStore implements Flushable, Closeable {
+    private static final Logger logger = LoggerFactory.getLogger(FileCacheStore.class);
+
     private final boolean gzip;
     private final boolean deduplicate;
+    private final Path tmpDir;
 
     private RecordConverterFactory converterFactory;
     private final int maxFiles;
     private final Map<Path, FileCache> caches;
 
-    // Response codes for each write record case
-    public static final int CACHE_AND_WRITE = 1; //used cache and write successful
-    public static final int NO_CACHE_AND_WRITE= 2;
-    public static final int CACHE_AND_NO_WRITE =3;
-    public static final int NO_CACHE_AND_NO_WRITE =4;
-
-
-    public FileCacheStore(RecordConverterFactory converterFactory, int maxFiles, boolean gzip, boolean deduplicate) {
+    public FileCacheStore(RecordConverterFactory converterFactory, int maxFiles, boolean gzip, boolean deduplicate) throws IOException {
         this.converterFactory = converterFactory;
         this.maxFiles = maxFiles;
         this.caches = new HashMap<>(maxFiles * 4 / 3 + 1);
         this.gzip = gzip;
         this.deduplicate = deduplicate;
+        this.tmpDir = Files.createTempDirectory("restructurehdfs");
     }
 
     /**
@@ -64,30 +66,30 @@ public FileCacheStore(RecordConverterFactory converterFactory, int maxFiles, boo
      * @return Integer value according to one of the response codes.
      * @throws IOException when failing to open a file or writing to it.
      */
-    public int writeRecord(Path path, GenericRecord record) throws IOException {
+    public WriteStatus writeRecord(Path path, GenericRecord record) throws IOException {
         FileCache cache = caches.get(path);
         if (cache != null) {
             if(cache.writeRecord(record)){
-                return CACHE_AND_WRITE;
+                return WriteStatus.CACHE_AND_WRITE;
             } else {
                 // This is the case when cache is used but write is unsuccessful
                 // because of different number columns in same topic
-                return CACHE_AND_NO_WRITE;
+                return WriteStatus.CACHE_AND_NO_WRITE;
             }
         } else {
             ensureCapacity();
 
             Path dir = path.getParent();
             Files.createDirectories(dir);
 
-            cache = new FileCache(converterFactory, path, record, gzip);
+            cache = new FileCache(converterFactory, path, record, gzip, tmpDir);
             caches.put(path, cache);
-            if(cache.writeRecord(record)) {
-                return NO_CACHE_AND_WRITE;
+            if (cache.writeRecord(record)) {
+                return WriteStatus.NO_CACHE_AND_WRITE;
             } else {
                 // The file path was not in cache but the file exists and this write is
                 // unsuccessful because of different number of columns
-                return NO_CACHE_AND_NO_WRITE;
+                return WriteStatus.NO_CACHE_AND_NO_WRITE;
             }
 
         }
@@ -127,9 +129,45 @@ public void close() throws IOException {
                     converterFactory.sortUnique(cache.getPath());
                 }
             }
+            Files.walk(tmpDir)
+                    .sorted(Comparator.reverseOrder())
+                    .forEach(tryCatch(Files::delete, (p, ex) -> logger.warn(
+                            "Failed to remove temporary file {}: {}", p, ex)));
         } finally {
             caches.clear();
         }
     }
 
+    // Response codes for each write record case
+    public enum WriteStatus {
+        /** Cache hit and write was successful. */
+        CACHE_AND_WRITE(true, true),
+        /** Cache hit and write was unsuccessful because of a mismatch in number of columns. */
+        CACHE_AND_NO_WRITE(true, false),
+        /** Cache miss and write was successful. */
+        NO_CACHE_AND_WRITE(false, true),
+        /** Cache miss and write was unsuccessful because of a mismatch in number of columns. */
+        NO_CACHE_AND_NO_WRITE(false, false);
+
+        private final boolean successful;
+        private final boolean cacheHit;
+
+        /**
+         * Write status.
+         * @param cacheHit whether the cache was used to write.
+         * @param successful whether the write was successful.
+         */
+        WriteStatus(boolean cacheHit, boolean successful) {
+            this.cacheHit = cacheHit;
+            this.successful = successful;
+        }
+
+        public boolean isSuccessful() {
+            return successful;
+        }
+
+        public boolean isCacheHit() {
+            return cacheHit;
+        }
+    }
 }

src/main/java/org/radarcns/util/JsonAvroConverter.java renamed to src/main/java/org/radarcns/data/JsonAvroConverter.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package org.radarcns.util;
+package org.radarcns.data;
 
 import com.fasterxml.jackson.core.JsonFactory;
 import com.fasterxml.jackson.core.JsonGenerator;

src/main/java/org/radarcns/util/RecordConverter.java renamed to src/main/java/org/radarcns/data/RecordConverter.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package org.radarcns.util;
+package org.radarcns.data;
 
 import java.io.Closeable;
 import java.io.Flushable;

src/main/java/org/radarcns/util/RecordConverterFactory.java renamed to src/main/java/org/radarcns/data/RecordConverterFactory.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package org.radarcns.util;
+package org.radarcns.data;
 
 import org.apache.avro.generic.GenericRecord;
 

src/main/java/org/radarcns/util/ThrowingConsumer.java

@@ -0,0 +1,20 @@
+package org.radarcns.util;
+
+import java.io.IOException;
+import java.util.function.BiConsumer;
+import java.util.function.Consumer;
+
+@FunctionalInterface
+public interface ThrowingConsumer<T> {
+    void accept(T t) throws IOException;
+
+    static <T> Consumer<T> tryCatch(ThrowingConsumer<T> consumer, BiConsumer<T, IOException> catchClause) {
+        return t -> {
+            try {
+                consumer.accept(t);
+            } catch (IOException ex) {
+                catchClause.accept(t, ex);
+            }
+        };
+    }
+}

src/test/java/org/radarcns/util/CsvAvroConverterTest.java

@@ -22,8 +22,6 @@
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
-import com.fasterxml.jackson.databind.JsonMappingException;
-
 import java.io.*;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -49,6 +47,9 @@
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
 import org.junit.rules.TemporaryFolder;
+import org.radarcns.data.CsvAvroConverter;
+import org.radarcns.data.RecordConverter;
+import org.radarcns.data.RecordConverterFactory;
 
 public class CsvAvroConverterTest {
     @Rule