embulk 0.2.1 → 0.3.0

data/embulk-core/src/main/java/org/embulk/exec/LocalExecutor.java

@@ -1,6 +1,7 @@
 package org.embulk.exec;
 
 import java.util.List;
+import java.util.Arrays;
 import java.util.ArrayList;
 import java.util.concurrent.Callable;
 import java.util.concurrent.Future;
@@ -13,9 +14,12 @@ import com.google.common.collect.ImmutableList;
 import com.google.inject.Inject;
 import com.google.inject.Injector;
 import com.google.common.base.Throwables;
+import com.google.common.base.Predicates;
+import com.google.common.collect.Iterables;
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import org.embulk.config.Task;
 import org.embulk.config.Config;
+import org.embulk.config.ConfigDefault;
 import org.embulk.config.ConfigSource;
 import org.embulk.config.TaskSource;
 import org.embulk.config.NextConfig;
@@ -26,8 +30,11 @@ import org.embulk.spi.Exec;
 import org.embulk.spi.ExecSession;
 import org.embulk.spi.ExecAction;
 import org.embulk.spi.InputPlugin;
+import org.embulk.spi.FilterPlugin;
 import org.embulk.spi.OutputPlugin;
+import org.embulk.spi.PageOutput;
 import org.embulk.spi.TransactionalPageOutput;
+import org.embulk.spi.util.Filters;
 import org.slf4j.Logger;
 
 public class LocalExecutor
@@ -37,22 +44,25 @@ public class LocalExecutor
     private final int maxThreads;
     private final ExecutorService executor;
 
-    private Logger log;
-    private final AtomicInteger runningTaskCount;
-    private final AtomicInteger completedTaskCount;
-
     public interface ExecutorTask
             extends Task
     {
         @Config("in")
         public ConfigSource getInputConfig();
 
+        @Config("filters")
+        @ConfigDefault("[]")
+        public List<ConfigSource> getFilterConfigs();
+
         @Config("out")
         public ConfigSource getOutputConfig();
 
         public TaskSource getInputTask();
         public void setInputTask(TaskSource taskSource);
 
+        public List<TaskSource> getFilterTasks();
+        public void setFilterTasks(List<TaskSource> taskSources);
+
         public TaskSource getOutputTask();
         public void setOutputTask(TaskSource taskSource);
     }
@@ -71,70 +81,251 @@ public class LocalExecutor
                         .setNameFormat("embulk-executor-%d")
                         .setDaemon(true)
                         .build());
-
-        this.runningTaskCount = new AtomicInteger(0);
-        this.completedTaskCount = new AtomicInteger(0);
     }
 
-    private static class ExecuteResultBuilder
+    private static class ProcessState
     {
-        private NextConfig inputNextConfig;
-        private NextConfig outputNextConfig;
+        private final Logger logger;
+        private volatile boolean[] started;
+        private volatile boolean[] finished;
+        private volatile Schema inputSchema;
+        private volatile Schema outputSchema;
+        private volatile Throwable[] exceptions;
+        private volatile CommitReport[] inputCommitReports;
+        private volatile CommitReport[] outputCommitReports;
+        private volatile NextConfig inputNextConfig;
+        private volatile NextConfig outputNextConfig;
+        private int processorCount;
+
+        public ProcessState(Logger logger)
+        {
+            this.logger = logger;
+        }
 
-        public void setInputNextConfig(NextConfig inputNextConfig)
+        public Logger getLogger()
         {
-            this.inputNextConfig = inputNextConfig;
+            return logger;
         }
 
-        public void setOutputNextConfig(NextConfig outputNextConfig)
+        public void initialize(int count)
         {
-            this.outputNextConfig = outputNextConfig;
+            this.started = new boolean[count];
+            this.finished = new boolean[count];
+            this.exceptions = new Throwable[count];
+            this.inputCommitReports = new CommitReport[count];
+            this.outputCommitReports = new CommitReport[count];
+            this.processorCount = count;
         }
 
-        public NextConfig getInputNextConfig()
+        public void setInputSchema(Schema inputSchema)
         {
-            return inputNextConfig;
+            this.inputSchema = inputSchema;
         }
 
-        public NextConfig getOutputNextConfig()
+        public void setOutputSchema(Schema outputSchema)
         {
-            return outputNextConfig;
+            this.outputSchema = outputSchema;
         }
 
-        public ExecuteResult build()
+        public Schema getInputSchema()
         {
-            if (inputNextConfig == null) {
-                inputNextConfig = Exec.newNextConfig();
+            return inputSchema;
+        }
+
+        public Schema getOutputSchema()
+        {
+            return outputSchema;
+        }
+
+        public boolean isAnyStarted()
+        {
+            if (started == null) {
+                return false;
             }
+            for (boolean b : started) {
+                if (b) { return true; }
+            }
+            return false;
+        }
+
+        public void start(int i)
+        {
+            started[i] = true;
+        }
+
+        public void finish(int i)
+        {
+            finished[i] = true;
+        }
+
+        public int getProcessorCount()
+        {
+            return processorCount;
+        }
+
+        public int getStartedCount()
+        {
+            int count = 0;
+            for (int i=0; i < started.length; i++) {
+                if (started[i]) { count++; }
+            }
+            return count;
+        }
+
+        public int getFinishedCount()
+        {
+            int count = 0;
+            for (int i=0; i < finished.length; i++) {
+                if (finished[i]) { count++; }
+            }
+            return count;
+        }
+
+        public void setInputCommitReport(int i, CommitReport inputCommitReport)
+        {
+            if (inputCommitReport == null) {
+                inputCommitReport = Exec.newCommitReport();
+            }
+            this.inputCommitReports[i] = inputCommitReport;
+        }
+
+        public void setOutputCommitReport(int i, CommitReport outputCommitReport)
+        {
+            if (outputCommitReport == null) {
+                outputCommitReport = Exec.newCommitReport();
+            }
+            this.outputCommitReports[i] = outputCommitReport;
+        }
+
+        public boolean isOutputCommitted(int i)
+        {
+            return outputCommitReports[i] != null;
+        }
+
+        public void setException(int i, Throwable exception)
+        {
+            this.exceptions[i] = exception;
+        }
+
+        public boolean isAllCommitted()
+        {
+            if (processorCount <= 0) {
+                // not initialized
+                return false;
+            }
+            for (int i=0; i < processorCount; i++) {
+                if (!isOutputCommitted(i)) {
+                    return false;
+                }
+            }
+            return true;
+        }
+
+        public boolean isAnyCommitted()
+        {
+            for (int i=0; i < processorCount; i++) {
+                if (isOutputCommitted(i)) {
+                    return true;
+                }
+            }
+            return false;
+        }
+
+        public void setOutputNextConfig(NextConfig outputNextConfig)
+        {
             if (outputNextConfig == null) {
                 outputNextConfig = Exec.newNextConfig();
             }
-            NextConfig nextConfig = Exec.newNextConfig();
-            nextConfig.getNestedOrSetEmpty("in").merge(inputNextConfig);
-            nextConfig.getNestedOrSetEmpty("out").merge(outputNextConfig);
-            return new ExecuteResult(nextConfig);
+            this.outputNextConfig = outputNextConfig;
         }
-    }
 
-    private static class ProcessResult
-    {
-        private final CommitReport inputCommitReport;
-        private final CommitReport outputCommitReport;
+        public void setInputNextConfig(NextConfig inputNextConfig)
+        {
+            if (inputNextConfig == null) {
+                inputNextConfig = Exec.newNextConfig();
+            }
+            this.inputNextConfig = inputNextConfig;
+        }
 
-        public ProcessResult(CommitReport inputCommitReport, CommitReport outputCommitReport)
+        public List<CommitReport> getInputCommitReports()
         {
-            this.inputCommitReport = inputCommitReport;
-            this.outputCommitReport = outputCommitReport;
+            return ImmutableList.copyOf(inputCommitReports);
         }
 
-        public CommitReport getInputCommitReport()
+        public List<CommitReport> getOutputCommitReports()
         {
-            return inputCommitReport;
+            return ImmutableList.copyOf(outputCommitReports);
         }
 
-        public CommitReport getOutputCommitReport()
+        public RuntimeException getRepresentativeException()
         {
-            return outputCommitReport;
+            RuntimeException top = null;
+            for (Throwable ex : exceptions) {
+                if (ex != null) {
+                    if (top != null) {
+                        top.addSuppressed(ex);
+                    } else {
+                        if (ex instanceof RuntimeException) {
+                            top = (RuntimeException) ex;
+                        } else {
+                            top = new RuntimeException(ex);
+                        }
+                    }
+                }
+            }
+            if (top == null) {
+                top = new RuntimeException("Some transactions are not committed");
+            }
+            return top;
+        }
+
+        public int getCommittedUnclosedCount()
+        {
+            int count = 0;
+            for (int i=0; i < exceptions.length; i++) {
+                if (exceptions[i] != null && isOutputCommitted(i)) {
+                    count++;
+                }
+            }
+            return count;
+        }
+
+        public ExecutionResult buildExecuteResult()
+        {
+            return buildExecuteResultWithWarningException(null);
+        }
+
+        public ExecutionResult buildExecuteResultWithWarningException(Throwable ex)
+        {
+            NextConfig nextConfig = Exec.newNextConfig();
+            if (inputNextConfig != null) {
+                nextConfig.getNestedOrSetEmpty("in").merge(inputNextConfig);
+            }
+            if (outputNextConfig != null) {
+                nextConfig.getNestedOrSetEmpty("out").merge(outputNextConfig);
+            }
+
+            ImmutableList.Builder<Throwable> ignoredExceptions = ImmutableList.builder();
+            for (Throwable e : exceptions) {
+                if (e != null) {
+                    ignoredExceptions.add(e);
+                }
+            }
+            if (ex != null) {
+                ignoredExceptions.add(ex);
+            }
+
+            return new ExecutionResult(nextConfig, ignoredExceptions.build());
+        }
+
+        public PartialExecutionException buildPartialExecuteException(Throwable cause,
+                ExecutorTask task, ExecSession exec)
+        {
+            return new PartialExecutionException(cause, new ResumeState(
+                        exec.getSessionConfigSource(),
+                        task.getInputTask(), task.getOutputTask(),
+                        inputSchema, outputSchema,
+                        Arrays.asList(inputCommitReports), Arrays.asList(outputCommitReports)));
         }
     }
 
@@ -143,17 +334,21 @@ public class LocalExecutor
         return Exec.newPlugin(InputPlugin.class, task.getInputConfig().get(PluginType.class, "type"));
     }
 
+    protected List<FilterPlugin> newFilterPlugins(ExecutorTask task)
+    {
+        return Filters.newFilterPlugins(Exec.session(), task.getFilterConfigs());
+    }
+
     protected OutputPlugin newOutputPlugin(ExecutorTask task)
     {
         return Exec.newPlugin(OutputPlugin.class, task.getOutputConfig().get(PluginType.class, "type"));
     }
 
-    public ExecuteResult run(ExecSession exec, final ConfigSource config)
+    public ExecutionResult run(ExecSession exec, final ConfigSource config)
     {
-        log = exec.getLogger(LocalExecutor.class);
         try {
-            return Exec.doWith(exec, new ExecAction<ExecuteResult>() {
-                public ExecuteResult run()
+            return Exec.doWith(exec, new ExecAction<ExecutionResult>() {
+                public ExecutionResult run()
                 {
                     return doRun(config);
                 }
@@ -163,114 +358,270 @@ public class LocalExecutor
         }
     }
 
-    private ExecuteResult doRun(ConfigSource config)
+    public ExecutionResult resume(final ConfigSource config, final ResumeState resume)
+    {
+        try {
+            ExecSession exec = new ExecSession(injector, resume.getExecSessionConfigSource());
+            return Exec.doWith(exec, new ExecAction<ExecutionResult>() {
+                public ExecutionResult run()
+                {
+                    return doResume(config, resume);
+                }
+            });
+        } catch (Exception ex) {
+            throw Throwables.propagate(ex);
+        }
+    }
+
+    public void cleanup(final ConfigSource config, final ResumeState resume)
+    {
+        try {
+            ExecSession exec = new ExecSession(injector, resume.getExecSessionConfigSource());
+            Exec.doWith(exec, new ExecAction<Void>() {
+                public Void run()
+                {
+                    doCleanup(config, resume);
+                    return null;
+                }
+            });
+        } catch (Exception ex) {
+            throw Throwables.propagate(ex);
+        }
+    }
+
+    public void doCleanup(ConfigSource config, ResumeState resume)
+    {
+        ExecutorTask task = config.loadConfig(ExecutorTask.class);
+        InputPlugin in = newInputPlugin(task);
+        OutputPlugin out = newOutputPlugin(task);
+
+        List<CommitReport> successInputCommitReports = ImmutableList.copyOf(
+                Iterables.filter(resume.getInputCommitReports(), Predicates.notNull()));
+        List<CommitReport> successOutputCommitReports = ImmutableList.copyOf(
+                Iterables.filter(resume.getOutputCommitReports(), Predicates.notNull()));
+
+        in.cleanup(resume.getInputTaskSource(), resume.getInputSchema(),
+                resume.getInputCommitReports().size(), successInputCommitReports);
+
+        out.cleanup(resume.getOutputTaskSource(), resume.getOutputSchema(),
+                resume.getOutputCommitReports().size(), successOutputCommitReports);
+    }
+
+    private ExecutionResult doRun(ConfigSource config)
     {
         final ExecutorTask task = config.loadConfig(ExecutorTask.class);
 
         final InputPlugin in = newInputPlugin(task);
+        final List<FilterPlugin> filterPlugins = newFilterPlugins(task);
         final OutputPlugin out = newOutputPlugin(task);
 
-        final ExecuteResultBuilder execResult = new ExecuteResultBuilder();
-
-        NextConfig inputNextConfig = in.transaction(task.getInputConfig(), new InputPlugin.Control() {
-            public List<CommitReport> run(final TaskSource inputTask, final Schema schema, final int processorCount)
-            {
-                final ImmutableList.Builder<CommitReport> inputCommitReports = ImmutableList.builder();
-                NextConfig outputNextConfig = out.transaction(task.getOutputConfig(), schema, processorCount, new OutputPlugin.Control() {
-                    public List<CommitReport> run(final TaskSource outputTask)
-                    {
-                        final ImmutableList.Builder<CommitReport> outputCommitReports = ImmutableList.builder();
-                        task.setInputTask(inputTask);
-                        task.setOutputTask(outputTask);
-
-                        //log.debug("input: %s", task.getInputTask());
-                        //log.debug("output: %s", task.getOutputTask());
-
-                        List<ProcessResult> results = process(task.dump(), schema, processorCount);
-                        for (ProcessResult result : results) {
-                            inputCommitReports.add(result.getInputCommitReport());
-                            outputCommitReports.add(result.getOutputCommitReport());
+        final ProcessState state = new ProcessState(Exec.getLogger(LocalExecutor.class));
+        try {
+            NextConfig inputNextConfig = in.transaction(task.getInputConfig(), new InputPlugin.Control() {
+                public List<CommitReport> run(final TaskSource inputTask, final Schema inputSchema, final int processorCount)
+                {
+                    state.initialize(processorCount);
+                    state.setInputSchema(inputSchema);
+                    Filters.transaction(filterPlugins, task.getFilterConfigs(), inputSchema, new Filters.Control() {
+                        public void run(final List<TaskSource> filterTasks, final List<Schema> filterSchemas)
+                        {
+                            Schema outputSchema = last(filterSchemas);
+                            state.setOutputSchema(outputSchema);
+                            NextConfig outputNextConfig = out.transaction(task.getOutputConfig(), outputSchema, processorCount, new OutputPlugin.Control() {
+                                public List<CommitReport> run(final TaskSource outputTask)
+                                {
+                                    task.setInputTask(inputTask);
+                                    task.setFilterTasks(filterTasks);
+                                    task.setOutputTask(outputTask);
+
+                                    process(task.dump(), filterSchemas, processorCount, state);
+                                    if (!state.isAllCommitted()) {
+                                        throw state.getRepresentativeException();
+                                    }
+                                    return state.getOutputCommitReports();
+                                }
+                            });
+                            state.setOutputNextConfig(outputNextConfig);
                         }
+                    });
+                    return state.getInputCommitReports();
+                }
+            });
+            state.setInputNextConfig(inputNextConfig);
 
-                        return outputCommitReports.build();
-                    }
-                });
-                execResult.setOutputNextConfig(outputNextConfig);
-                return inputCommitReports.build();
+            return state.buildExecuteResult();
+
+        } catch (Throwable ex) {
+            if (state.isAllCommitted()) {
+                // ignore the exception
+                return state.buildExecuteResultWithWarningException(ex);
             }
-        });
-        execResult.setInputNextConfig(inputNextConfig);
+            if (!state.isAnyStarted()) {
+                throw ex;
+            }
+            throw state.buildPartialExecuteException(ex, task, Exec.session());
+        }
+    }
+
+    private ExecutionResult doResume(ConfigSource config, final ResumeState resume)
+    {
+        final ExecutorTask task = config.loadConfig(ExecutorTask.class);
+
+        final InputPlugin in = newInputPlugin(task);
+        final List<FilterPlugin> filterPlugins = newFilterPlugins(task);
+        final OutputPlugin out = newOutputPlugin(task);
+
+        final ProcessState state = new ProcessState(Exec.getLogger(LocalExecutor.class));
+        try {
+            NextConfig inputNextConfig = in.resume(resume.getInputTaskSource(), resume.getInputSchema(), resume.getInputCommitReports().size(), new InputPlugin.Control() {
+                public List<CommitReport> run(final TaskSource inputTask, final Schema inputSchema, final int processorCount)
+                {
+                    // TODO validate inputTask?
+                    // TODO validate inputSchema
+                    // TODO validate processorCount
+                    state.initialize(processorCount);
+                    Filters.transaction(filterPlugins, task.getFilterConfigs(), inputSchema, new Filters.Control() {
+                        public void run(final List<TaskSource> filterTasks, final List<Schema> filterSchemas)
+                        {
+                            Schema outputSchema = last(filterSchemas);
+                            state.setOutputSchema(outputSchema);
+                            NextConfig outputNextConfig = out.resume(resume.getOutputTaskSource(), outputSchema, processorCount, new OutputPlugin.Control() {
+                                public List<CommitReport> run(final TaskSource outputTask)
+                                {
+                                    // TODO validate outputTask?
+                                    task.setInputTask(inputTask);
+                                    task.setFilterTasks(filterTasks);
+                                    task.setOutputTask(outputTask);
+
+                                    for (int i=0; i < resume.getOutputCommitReports().size(); i++) {
+                                        if (resume.getOutputCommitReports().get(i) != null) {
+                                            state.start(i);
+                                            state.setInputCommitReport(i, resume.getInputCommitReports().get(i));
+                                            state.setOutputCommitReport(i, resume.getOutputCommitReports().get(i));
+                                            state.finish(i);
+                                        }
+                                    }
+
+                                    process(task.dump(), filterSchemas, processorCount, state);
+                                    if (!state.isAllCommitted()) {
+                                        throw state.getRepresentativeException();
+                                    }
+                                    return state.getOutputCommitReports();
+                                }
+                            });
+                            state.setOutputNextConfig(outputNextConfig);
+                        }
+                    });
+                    return state.getInputCommitReports();
+                }
+            });
+            state.setInputNextConfig(inputNextConfig);
 
-        return execResult.build();
+            return state.buildExecuteResult();
+
+        } catch (Throwable ex) {
+            if (state.isAllCommitted()) {
+                // ignore the exception
+                return state.buildExecuteResultWithWarningException(ex);
+            }
+            if (!state.isAnyStarted()) {
+                throw ex;
+            }
+            throw state.buildPartialExecuteException(ex, task, Exec.session());
+        }
     }
 
-    private List<ProcessResult> process(TaskSource taskSource, Schema schema, int processorCount)
+    private void process(TaskSource taskSource, List<Schema> filterSchemas, int processorCount,
+            ProcessState state)
     {
-        List<Future<ProcessResult>> futures = new ArrayList<>();
-        List<ProcessResult> joined = new ArrayList<>();
+        List<Future<Throwable>> futures = new ArrayList<>(processorCount);
         try {
-            log.info("Running {} tasks using {} local threads", processorCount, maxThreads);
-            showProgress(processorCount);
             for (int i=0; i < processorCount; i++) {
-                futures.add(startProcessor(taskSource, schema, i));
+                if (state.isOutputCommitted(i)) {
+                    state.getLogger().warn("Skipped resumed task {}", i);
+                    futures.add(null);  // resumed
+                } else {
+                    futures.add(startProcessor(taskSource, filterSchemas, i, state));
+                }
             }
+            showProgress(state);
 
             for (int i=0; i < processorCount; i++) {
+                if (futures.get(i) == null) {
+                    continue;
+                }
                 try {
-                    joined.add(futures.get(i).get());
-                    showProgress(processorCount);
-
+                    state.setException(i, futures.get(i).get());
                 } catch (ExecutionException ex) {
-                    throw Throwables.propagate(ex.getCause());
+                    state.setException(i, ex.getCause());
+                    //Throwables.propagate(ex.getCause());
                 } catch (InterruptedException ex) {
-                    throw new ExecuteInterruptedException(ex);
+                    state.setException(i, new ExecutionInterruptedException(ex));
                 }
+                showProgress(state);
             }
-            return joined;
         } finally {
-            for (int i=joined.size(); i < futures.size(); i++) {
-                futures.get(i).cancel(true);
-                // TODO join?
+            for (Future<Throwable> future : futures) {
+                if (future != null && !future.isDone()) {
+                    future.cancel(true);
+                    // TODO join?
+                }
             }
         }
     }
 
-    private void showProgress(int total)
+    private void showProgress(ProcessState state)
     {
-        int running = runningTaskCount.get();
-        int done = completedTaskCount.get();
-        log.info(String.format("{done:%3d / %d, running: %d}", done, total, running));
+        int total = state.getProcessorCount();
+        int finished = state.getFinishedCount();
+        int started = state.getStartedCount();
+        state.getLogger().info(String.format("{done:%3d / %d, running: %d}", finished, total, started - finished));
     }
 
-    private Future<ProcessResult> startProcessor(final TaskSource taskSource, final Schema schema, final int index)
+    private Future<Throwable> startProcessor(final TaskSource taskSource,
+            final List<Schema> filterSchemas, final int index,
+            final ProcessState state)
     {
-        return executor.submit(new Callable<ProcessResult>() {
-            public ProcessResult call()
+        return executor.submit(new Callable<Throwable>() {
+            public Throwable call()
             {
+                final ExecutorTask task = taskSource.loadTask(ExecutorTask.class);
+                final InputPlugin in = newInputPlugin(task);
+                final List<FilterPlugin> filterPlugins = newFilterPlugins(task);
+                final OutputPlugin out = newOutputPlugin(task);
+
+                TransactionalPageOutput tran = out.open(task.getOutputTask(), last(filterSchemas), index);
+                PageOutput closeThis = tran;
+                state.start(index);
                 try {
-                    runningTaskCount.getAndIncrement();
-                    final ExecutorTask task = taskSource.loadTask(ExecutorTask.class);
-                    final InputPlugin in = newInputPlugin(task);
-                    final OutputPlugin out = newOutputPlugin(task);
-
-                    TransactionalPageOutput tran = out.open(task.getOutputTask(), schema, index);
-                    boolean committed = false;
+                    PageOutput filtered = closeThis = Filters.open(filterPlugins, task.getFilterTasks(), filterSchemas, tran);
+                    state.setInputCommitReport(index, in.run(task.getInputTask(), first(filterSchemas), index, filtered));
+                    state.setOutputCommitReport(index, tran.commit());  // TODO check output.finish() is called. wrap or abstract
+                    return null;
+                } finally {
                     try {
-                        CommitReport inReport = in.run(task.getInputTask(), schema, index, tran);
-                        CommitReport outReport = tran.commit();  // TODO check output.finish() is called. wrap or abstract
-                        committed = true;
-                        return new ProcessResult(inReport, outReport);
-                    } finally {
-                        if (!committed) {
-                            tran.abort();
+                        try {
+                            if (!state.isOutputCommitted(index)) {
+                                tran.abort();
+                            }
+                        } finally {
+                            closeThis.close();
                         }
-                        tran.close();
+                    } finally {
+                        state.finish(index);
                     }
-                } finally {
-                    runningTaskCount.getAndDecrement();
-                    completedTaskCount.getAndIncrement();
                 }
             }
         });
     }
+
+    private static Schema first(List<Schema> filterSchemas)
+    {
+        return filterSchemas.get(0);
+    }
+
+    private static Schema last(List<Schema> filterSchemas)
+    {
+        return filterSchemas.get(filterSchemas.size() - 1);
+    }
 }