Merge pull request #51 from embulk/fix_stalling_issue

Fix stalling issue when closing `remoteFile`

Author: tvhung83

src/main/java/org/embulk/output/sftp/SftpUtils.java

@@ -20,6 +20,7 @@
 import org.slf4j.Logger;
 
 import java.io.BufferedOutputStream;
+import java.io.Closeable;
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
@@ -39,7 +40,7 @@
 public class SftpUtils
 {
     private final Logger logger = Exec.getLogger(SftpUtils.class);
-    private final DefaultFileSystemManager manager;
+    private DefaultFileSystemManager manager;
     private final FileSystemOptions fsOptions;
     private final String userInfo;
     private final String user;
@@ -159,15 +160,24 @@ public Void call() throws Exception
                 final FileObject remoteFile = newSftpFile(getSftpFileUri(remotePath));
                 final BufferedOutputStream outputStream = openStream(remoteFile);
                 // When channel is broken, closing resource may hang, hence the time-out wrapper
-                // Note: closing FileObject will also close OutputStream
-                try (TimeoutCloser ignored = new TimeoutCloser(outputStream)) {
+                try (final TimeoutCloser ignored = new TimeoutCloser(outputStream)) {
                     appendFile(localTempFile, remoteFile, outputStream);
                     return null;
                 }
                 finally {
-                    remoteFile.close();
+                    // closing sequentially
+                    new TimeoutCloser(remoteFile).close();
                 }
             }
+
+            @Override
+            public void onRetry(Exception exception, int retryCount, int retryLimit, int retryWait)
+            {
+                super.onRetry(exception, retryCount, retryLimit, retryWait);
+                // re-connect
+                manager.close();
+                manager = initializeStandardFileSystemManager();
+            }
         });
     }
 
@@ -258,18 +268,9 @@ FileObject resolve(final String remoteFilePath) throws FileSystemException
         return manager.resolveFile(getSftpFileUri(remoteFilePath).toString(), fsOptions);
     }
 
-    BufferedOutputStream openStream(final FileObject remoteFile)
+    BufferedOutputStream openStream(final FileObject remoteFile) throws FileSystemException
     {
-        // output stream is already a BufferedOutputStream, no need to wrap
-        final String taskName = "SFTP open stream";
-        return withRetry(new DefaultRetry<BufferedOutputStream>(taskName)
-        {
-            @Override
-            public BufferedOutputStream call() throws Exception
-            {
-                return new BufferedOutputStream(remoteFile.getContent().getOutputStream());
-            }
-        });
+        return new BufferedOutputStream(remoteFile.getContent().getOutputStream());
     }
 
     URI getSftpFileUri(String remoteFilePath)

src/main/java/org/embulk/output/sftp/utils/TimeoutCloser.java

@@ -10,8 +10,7 @@
 
 public class TimeoutCloser implements Closeable
 {
-    @VisibleForTesting
-    int timeout = 300; // 5 minutes
+    private static int timeout = 300; // 5 minutes
     private Closeable wrapped;
 
     public TimeoutCloser(Closeable wrapped)
@@ -39,4 +38,10 @@ public Void call() throws Exception
             throw Throwables.propagate(e);
         }
     }
+
+    @VisibleForTesting
+    public static void setTimeout(int timeout)
+    {
+        TimeoutCloser.timeout = timeout;
+    }
 }

src/test/java/org/embulk/output/sftp/TestSftpFileOutputPlugin.java

@@ -768,7 +768,7 @@ public void testResolveWithoutRetry()
     }
 
     @Test
-    public void testOpenStreamWithRetry() throws FileSystemException
+    public void testOpenStreamWithoutRetry() throws FileSystemException
     {
         SftpFileOutputPlugin.PluginTask task = defaultTask();
         SftpUtils utils = new SftpUtils(task);
@@ -778,9 +778,13 @@ public void testOpenStreamWithRetry() throws FileSystemException
                 .doCallRealMethod()
                 .when(mock).getContent();
 
-        OutputStream stream = utils.openStream(mock);
-        assertNotNull(stream);
-        Mockito.verify(mock, Mockito.times(2)).getContent();
+        try {
+            utils.openStream(mock);
+            fail("Should not reach here");
+        }
+        catch (FileSystemException e) {
+            Mockito.verify(mock, Mockito.times(1)).getContent();
+        }
     }
 
     @Test

src/test/java/org/embulk/output/sftp/utils/TestTimeoutCloser.java

@@ -32,7 +32,7 @@ public void close()
                 }
             }
         });
-        closer.timeout = 1;
+        TimeoutCloser.setTimeout(1);
         try {
             closer.close();
             fail("Should not finish");