[BACKLOG-24582] - Updated FileStreamMap logic

pentaho · Oct 16, 2018 · d9e45ab · d9e45ab
1 parent df49e05
commit d9e45ab
Showing 1 changed file with 49 additions and 42 deletions.
diff --git a/engine/src/main/java/org/pentaho/di/trans/steps/textfileoutput/TextFileOutputData.java b/engine/src/main/java/org/pentaho/di/trans/steps/textfileoutput/TextFileOutputData.java
@@ -59,13 +59,15 @@ interface IFileStreamsCollection {
     void add( String filename, FileStream fileStreams );
   }
 
-  public class FileStreamsKey implements Comparable<FileStreamsKey> {
+  public class FileStreamsCollectionEntry {
-    String fileName;
+    private String fileName;
-    long index = 0;
+    private long index = 0;
+    private FileStream fileStream;
 
-    public FileStreamsKey( String fileName, long index ) {
+    public FileStreamsCollectionEntry( String fileName, long index, FileStream fileStream ) {
       this.fileName = fileName;
       this.index = index;
+      this.fileStream = fileStream;
     }
 
     public String getFileName() {
@@ -84,17 +86,12 @@ public void setIndex( int index ) {
       this.index = index;
     }
 
-    @Override
+    public FileStream getFileStream() {
-    public int compareTo( FileStreamsKey o ) {
+      return fileStream;
-      if ( o.fileName.equals( fileName ) ) {
-        return 0;
-      }
-      return Long.compare( index, o.index );
     }
 
-    @Override
+    public void setFileStream( FileStream fileStream ) {
-    public boolean equals( Object obj ) {
+      this.fileStream = fileStream;
-      return fileName.equals( obj );
     }
   }
 
@@ -267,16 +264,25 @@ public void add( String filename, FileStream fileStreams ) {
     }
   }
 
-  public class FileStreamsMap extends TreeMap<FileStreamsKey, FileStream> implements IFileStreamsCollection  {
+  public class FileStreamsMap implements IFileStreamsCollection  {
-    int numOpenFiles = 0;
+    private int numOpenFiles = 0;
+    private TreeMap<String, FileStreamsCollectionEntry> fileNameMap = new TreeMap<>();
+    private TreeMap<Long, FileStreamsCollectionEntry> indexMap = new TreeMap<>();
 
     @Override
-    public void add( String filename, FileStream fileWriterOutputStream ) {
+    public int size() {
+      return fileNameMap.size();
+    }
+
+    @Override
+    public void add( String fileName, FileStream fileWriterOutputStream ) {
       long index = 0;
       if ( size() > 0 ) {
-        index = lastKey().index + 1;
+        index = indexMap.lastKey() + 1;
       }
-      put( new FileStreamsKey( filename, index ), fileWriterOutputStream );
+      FileStreamsCollectionEntry newEntry = new FileStreamsCollectionEntry( fileName, index, fileWriterOutputStream );
+      fileNameMap.put( fileName, newEntry );
+      indexMap.put( index, newEntry );
       if ( fileWriterOutputStream.isOpen() ) {
         numOpenFiles++;
       }
@@ -285,29 +291,27 @@ public void add( String filename, FileStream fileWriterOutputStream ) {
 
     @Override
     public FileStream getStream( String filename ) {
-      return get( new FileStreamsKey( filename, -1 ) );
+      if ( fileNameMap.containsKey( filename ) ) {
-    }
+        return fileNameMap.get( filename ).getFileStream();
-
+      } else {
-    private FileStream remove( String filename ) {
+        return null;
-      return remove( new FileStreamsKey( filename, -1 ) );
+      }
     }
 
     @Override
     public String getLastFileName() {
       String filename = null;
-      Map.Entry<TextFileOutputData.FileStreamsKey, FileStream> lastEntry = lastEntry();
+      if ( indexMap.size() > 0 ) {
-      if ( lastEntry != null ) {
+        filename = indexMap.lastEntry().getValue().getFileName();
-        filename = lastEntry.getKey( ).getFileName( );
       }
       return filename;
     }
 
     @Override
     public FileStream getLastStream() {
       FileStream lastStream = null;
-      Map.Entry<TextFileOutputData.FileStreamsKey, FileStream> lastEntry = lastEntry();
+      if ( indexMap.size() > 0 ) {
-      if ( lastEntry != null ) {
+        lastStream = indexMap.lastEntry().getValue().getFileStream();
-        lastStream = lastEntry.getValue( );
       }
       return lastStream;
     }
@@ -321,11 +325,13 @@ public int getNumOpenFiles() {
     public void closeOldestOpenFile( boolean removeFileFromCollection ) throws IOException {
       FileStream oldestOpenStream = null;
       String oldestOpenFileName = null;
-      for ( Map.Entry<TextFileOutputData.FileStreamsKey, FileStream> mapEntry : entrySet() ) {
+      Long oldestOpenFileIndex = null;
-        FileStream existingStream = mapEntry.getValue();
+      for ( Map.Entry<Long, FileStreamsCollectionEntry> mapEntry : indexMap.entrySet() ) {
-        if ( existingStream.isOpen() ) {
+        FileStreamsCollectionEntry existingStream = mapEntry.getValue();
-          oldestOpenStream = existingStream;
+        if ( existingStream.getFileStream().isOpen() ) {
-          oldestOpenFileName = mapEntry.getKey().getFileName();
+          oldestOpenStream = existingStream.getFileStream();
+          oldestOpenFileName = existingStream.getFileName();
+          oldestOpenFileIndex = existingStream.getIndex();
           break;
         }
       }
@@ -334,19 +340,20 @@ public void closeOldestOpenFile( boolean removeFileFromCollection ) throws IOExc
         oldestOpenStream.close();
         numOpenFiles--;
         if ( removeFileFromCollection ) {
-          remove( oldestOpenFileName );
+          fileNameMap.remove( oldestOpenFileName );
+          indexMap.remove( oldestOpenFileIndex );
         }
       }
     }
 
     @Override
     public void flushOpenFiles( boolean closeAfterFlush ) {
-      for ( FileStream outputStream : values() ) {
+      for ( FileStreamsCollectionEntry collectionEntry : indexMap.values() ) {
-        if ( outputStream.isDirty() ) {
+        if ( collectionEntry.getFileStream().isDirty() ) {
           try {
-            outputStream.flush();
+            collectionEntry.getFileStream().flush();
             if ( closeAfterFlush ) {
-              outputStream.close();
+              collectionEntry.getFileStream().close();
             }
           } catch ( IOException e ) {
             e.printStackTrace();
@@ -367,10 +374,10 @@ public void closeFile( String filename ) throws IOException {
 
     @Override
     public void closeStream( OutputStream outputStream ) throws IOException {
-      for ( Map.Entry<TextFileOutputData.FileStreamsKey, FileStream> mapEntry : entrySet() ) {
+      for ( Map.Entry<Long, FileStreamsCollectionEntry> mapEntry : indexMap.entrySet() ) {
-        FileStream fileStream = mapEntry.getValue();
+        FileStream fileStream = mapEntry.getValue().getFileStream();
         if ( ( fileStream.getBufferedOutputStream() == outputStream ) || ( fileStream.getCompressedOutputStream() == outputStream ) || ( fileStream.getFileOutputStream() == outputStream ) ) {
-          closeFile( mapEntry.getKey().getFileName() );
+          closeFile( mapEntry.getValue().getFileName() );
         }
       }
     }