embulk-executor-remoteserver 0.1.0

data/src/main/java/org/embulk/executor/remoteserver/RemoteServerExecutor.java

@@ -0,0 +1,131 @@
+package org.embulk.executor.remoteserver;
+
+import com.google.inject.Inject;
+import org.embulk.config.Config;
+import org.embulk.config.ConfigDefault;
+import org.embulk.config.ConfigInject;
+import org.embulk.config.ConfigSource;
+import org.embulk.config.ModelManager;
+import org.embulk.config.Task;
+import org.embulk.exec.ForSystemConfig;
+import org.embulk.spi.Exec;
+import org.embulk.spi.ExecutorPlugin;
+import org.embulk.spi.ProcessState;
+import org.embulk.spi.ProcessTask;
+import org.embulk.spi.Schema;
+import org.jruby.embed.ScriptingContainer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.nio.file.Files;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.TimeoutException;
+
+public class RemoteServerExecutor implements ExecutorPlugin {
+    private static final Logger log = LoggerFactory.getLogger(RemoteServerExecutor.class);
+    private static final Host DEFAULT_HOST = new Host("localhost", 30000);
+    private final ConfigSource systemConfig;
+    private final ScriptingContainer jruby;
+
+    interface PluginTask extends Task {
+        @Config("hosts")
+        @ConfigDefault("[]")
+        List<Host> getHosts();
+
+        @Config("timeout_seconds")
+        @ConfigDefault("3600")
+        int getTimeoutSeconds();
+
+        @ConfigInject
+        ModelManager getModelManager();
+    }
+
+    @Inject
+    public RemoteServerExecutor(@ForSystemConfig ConfigSource systemConfig, ScriptingContainer jruby) {
+        this.systemConfig = systemConfig;
+        this.jruby = jruby;
+    }
+
+    @Override
+    public void transaction(ConfigSource config, Schema outputSchema, int inputTaskCount, Control control) {
+        PluginTask task = config.loadConfig(PluginTask.class);
+        if (task.getHosts().isEmpty()) {
+            log.info("Hosts is empty. Run with a local server.");
+            try (EmbulkServer _autoclosed = EmbulkServer.start(DEFAULT_HOST.getName(), DEFAULT_HOST.getPort(), 1)) {
+                control.transaction(outputSchema, inputTaskCount, new ExecutorImpl(inputTaskCount, task, Collections.singletonList(DEFAULT_HOST)));
+            } catch (IOException e) {
+                throw new UncheckedIOException(e);
+            }
+        } else {
+            control.transaction(outputSchema, inputTaskCount, new ExecutorImpl(inputTaskCount, task, task.getHosts()));
+        }
+    }
+
+    private class ExecutorImpl implements ExecutorPlugin.Executor {
+        private final PluginTask pluginTask;
+        private final int inputTaskCount;
+        private final List<Host> hosts;
+
+        ExecutorImpl(int inputTaskCount, PluginTask pluginTask, List<Host> hosts) {
+            this.inputTaskCount = inputTaskCount;
+            this.pluginTask = pluginTask;
+            this.hosts = hosts;
+        }
+
+        @Override
+        public void execute(ProcessTask processTask, ProcessState state) {
+            byte[] pluginArchiveBytes;
+            List<PluginArchive.GemSpec> gemSpecs;
+            try {
+                File tempFile = Exec.getTempFileSpace().createTempFile("gems", ".zip");
+                gemSpecs = archivePlugins(tempFile);
+                pluginArchiveBytes = Files.readAllBytes(tempFile.toPath());
+            } catch (IOException e) {
+                throw new UncheckedIOException(e);
+            }
+            // Remove 'jruby_global_bundler_plugin_source_directory' (--bundle option)
+            // because all gems will be loaded via PluginArchive on server
+            ConfigSource systemConfigToSend = systemConfig.deepCopy().remove("jruby_global_bundler_plugin_source_directory");
+
+            ModelManager modelManager = pluginTask.getModelManager();
+            String systemConfigJson = modelManager.writeObject(systemConfigToSend);
+            String pluginTaskJson = modelManager.writeObject(pluginTask);
+            String processTaskJson = modelManager.writeObject(processTask);
+
+            SessionState sessionState = new SessionState(
+                    systemConfigJson, pluginTaskJson, processTaskJson, gemSpecs, pluginArchiveBytes, state, inputTaskCount, modelManager);
+            try (EmbulkClient client = EmbulkClient.open(sessionState, hosts)) {
+                client.createSession();
+
+                state.initialize(inputTaskCount, inputTaskCount);
+                for (int i = 0; i < inputTaskCount; i++) {
+                    if (state.getOutputTaskState(i).isCommitted()) {
+                        log.warn("Skipped resumed task {}", i);
+                        continue;
+                    }
+                    client.startTask(i);
+                }
+                sessionState.waitUntilCompleted(pluginTask.getTimeoutSeconds() + 1); // Add 1 sec to consider network latency
+            } catch (InterruptedException | TimeoutException e) {
+                throw new IllegalStateException(e);
+            } catch (IOException e) {
+                throw new UncheckedIOException(e);
+            }
+        }
+
+        private List<PluginArchive.GemSpec> archivePlugins(File tempFile) throws IOException {
+            // archive plugins
+            PluginArchive archive = new PluginArchive.Builder()
+                    .addLoadedRubyGems(jruby)
+                    .build();
+            try (FileOutputStream fos = new FileOutputStream(tempFile)) {
+                return archive.dump(fos);
+            }
+        }
+    }
+}

data/src/main/java/org/embulk/executor/remoteserver/RemoveSessionCommand.java

@@ -0,0 +1,24 @@
+package org.embulk.executor.remoteserver;
+
+import com.github.kamatama41.nsocket.Connection;
+import com.github.kamatama41.nsocket.SyncCommand;
+
+public class RemoveSessionCommand implements SyncCommand<String, Void> {
+    static final String ID = "remove_session";
+    private final SessionManager sessionManager;
+
+    RemoveSessionCommand(SessionManager sessionManager) {
+        this.sessionManager = sessionManager;
+    }
+
+    @Override
+    public Void apply(String sessionId, Connection connection) {
+        sessionManager.removeSession(sessionId);
+        return null;
+    }
+
+    @Override
+    public String getId() {
+        return ID;
+    }
+}

data/src/main/java/org/embulk/executor/remoteserver/Session.java

@@ -0,0 +1,177 @@
+package org.embulk.executor.remoteserver;
+
+import com.github.kamatama41.nsocket.Connection;
+import org.embulk.EmbulkEmbed;
+import org.embulk.config.ConfigSource;
+import org.embulk.config.ModelManager;
+import org.embulk.config.TaskReport;
+import org.embulk.spi.Exec;
+import org.embulk.spi.ExecSession;
+import org.embulk.spi.ProcessTask;
+import org.embulk.spi.util.Executors;
+import org.jruby.embed.ScriptingContainer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.UncheckedIOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Queue;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.TimeUnit;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+class Session implements AutoCloseable {
+    private static final Logger log = LoggerFactory.getLogger(Session.class);
+    private final String id;
+    private final EmbulkEmbed embed;
+    private final ScriptingContainer jruby;
+    private final RemoteServerExecutor.PluginTask pluginTask;
+    private final ProcessTask processTask;
+    private final List<PluginArchive.GemSpec> gemSpecs;
+    private final byte[] pluginArchive;
+    private final ExecSession session;
+    private final ModelManager modelManager;
+    private final ConcurrentMap<Integer, Queue<UpdateTaskStateData>> bufferMap;
+    private final ExecutorService sessionRunner;
+    private volatile Connection connection;
+
+    Session(
+            String id,
+            String systemConfig,
+            String pluginTaskConfig,
+            String processTaskConfig,
+            List<PluginArchive.GemSpec> gemSpecs,
+            byte[] pluginArchive
+    ) {
+        this.id = id;
+        this.embed = newEmbulkBootstrap(systemConfig).initialize();
+        this.jruby = embed.getInjector().getInstance(ScriptingContainer.class);
+        this.modelManager = embed.getModelManager();
+        this.pluginTask = modelManager.readObject(RemoteServerExecutor.PluginTask.class, pluginTaskConfig);
+        this.processTask = modelManager.readObject(ProcessTask.class, processTaskConfig);
+        this.gemSpecs = gemSpecs;
+        this.pluginArchive = pluginArchive;
+        this.session = ExecSession.builder(embed.getInjector()).build();
+        loadPluginArchive();
+        this.bufferMap = new ConcurrentHashMap<>();
+        this.sessionRunner = java.util.concurrent.Executors.newCachedThreadPool(r -> {
+            Thread t = new Thread(r);
+            t.setName("session-runner-" + id);
+            t.setDaemon(true);
+            return t;
+        });
+    }
+
+    void runTaskAsynchronously(int taskIndex) {
+        sessionRunner.submit(() -> Exec.doWith(session, () -> {
+            runTask(taskIndex);
+            return null;
+        }));
+    }
+
+    private void runTask(int taskIndex) throws InterruptedException {
+        bufferMap.putIfAbsent(taskIndex, new LinkedList<>());
+        try {
+            Executors.process(session, processTask, taskIndex, new Executors.ProcessStateCallback() {
+                @Override
+                public void started() {
+                    sendCommand(taskIndex, new UpdateTaskStateData(id, taskIndex, TaskState.STARTED));
+                }
+
+                @Override
+                public void inputCommitted(TaskReport report) {
+                    UpdateTaskStateData data = new UpdateTaskStateData(id, taskIndex, TaskState.INPUT_COMMITTED);
+                    data.setTaskReport(modelManager.writeObject(report));
+                    sendCommand(taskIndex, data);
+                }
+
+                @Override
+                public void outputCommitted(TaskReport report) {
+                    UpdateTaskStateData data = new UpdateTaskStateData(id, taskIndex, TaskState.OUTPUT_COMMITTED);
+                    data.setTaskReport(modelManager.writeObject(report));
+                    sendCommand(taskIndex, data);
+                }
+            });
+            sendCommand(taskIndex, new UpdateTaskStateData(id, taskIndex, TaskState.FINISHED));
+        } catch (Exception e) {
+            log.warn(String.format("Failed to run task[%d]", taskIndex), e);
+            UpdateTaskStateData data = new UpdateTaskStateData(id, taskIndex, TaskState.FAILED);
+            data.setErrorMessage(e.getMessage());
+            sendCommand(taskIndex, data);
+        }
+
+        Queue<UpdateTaskStateData> buffer = bufferMap.get(taskIndex);
+        if (buffer.isEmpty()) {
+            return;
+        }
+
+        // Flush buffer if remaining
+        int waitSeconds = 10;
+        while (!buffer.isEmpty()) {
+            if (connection.isOpen()) {
+                flushBuffer(taskIndex, connection);
+                return;
+            }
+            log.warn("Connection is closed, wait {} seconds until reconnected.", waitSeconds);
+            TimeUnit.SECONDS.sleep(waitSeconds);
+        }
+    }
+
+    void updateConnection(Connection connection) {
+        this.connection = connection;
+    }
+
+    private void sendCommand(int taskIndex, UpdateTaskStateData data) {
+        bufferMap.get(taskIndex).offer(data);
+        if (!connection.isOpen()) {
+            log.warn("Connection is closed, add data to buffer.");
+            return;
+        }
+        flushBuffer(taskIndex, connection);
+    }
+
+    private void flushBuffer(int taskIndex, Connection connection) {
+        UpdateTaskStateData data;
+        Queue<UpdateTaskStateData> buffer = bufferMap.get(taskIndex);
+        while ((data = buffer.poll()) != null) {
+            connection.sendCommand(NotifyTaskStateCommand.ID, data);
+        }
+    }
+
+    private static EmbulkEmbed.Bootstrap newEmbulkBootstrap(String configJson) {
+        ConfigSource systemConfig = getSystemConfig(configJson);
+        return new EmbulkEmbed.Bootstrap().setSystemConfig(systemConfig);
+    }
+
+    private static ConfigSource getSystemConfig(String configJson) {
+        try (InputStream in = new ByteArrayInputStream(configJson.getBytes(UTF_8))) {
+            return EmbulkEmbed.newSystemConfigLoader().fromJson(in);
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+
+    private void loadPluginArchive() {
+        try (ByteArrayInputStream bis = new ByteArrayInputStream(pluginArchive)) {
+            Path gemsDir = Files.createTempDirectory("embulk_gems");
+            PluginArchive.load(gemsDir.toFile(), gemSpecs, bis).restoreLoadPathsTo(jruby);
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+
+    @Override
+    public void close() {
+        log.debug("Closing the session {}", id);
+        sessionRunner.shutdownNow();
+    }
+}

data/src/main/java/org/embulk/executor/remoteserver/SessionManager.java

@@ -0,0 +1,37 @@
+package org.embulk.executor.remoteserver;
+
+import com.github.kamatama41.nsocket.Connection;
+
+import java.util.List;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+
+class SessionManager {
+    private final ConcurrentMap<String, Session> sessionMap;
+
+    SessionManager() {
+        this.sessionMap = new ConcurrentHashMap<>();
+    }
+
+    void registerNewSession(String sessionId,
+                            String systemConfig,
+                            String pluginTaskConfig,
+                            String processTaskConfig,
+                            List<PluginArchive.GemSpec> gemSpecs,
+                            byte[] pluginArchive,
+                            Connection connection) {
+        Session session = sessionMap.computeIfAbsent(
+                sessionId, (k) -> new Session(
+                        sessionId, systemConfig, pluginTaskConfig, processTaskConfig, gemSpecs, pluginArchive));
+        session.updateConnection(connection);
+    }
+
+    Session getSession(String sessionId) {
+        return sessionMap.get(sessionId);
+    }
+
+    void removeSession(String sessionId) {
+        Session removed = sessionMap.remove(sessionId);
+        removed.close();
+    }
+}

data/src/main/java/org/embulk/executor/remoteserver/SessionState.java

@@ -0,0 +1,143 @@
+package org.embulk.executor.remoteserver;
+
+import org.embulk.config.ModelManager;
+import org.embulk.config.TaskReport;
+import org.embulk.spi.ProcessState;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.stream.Collectors;
+
+class SessionState {
+    private static final Logger log = LoggerFactory.getLogger(SessionState.class);
+
+    private String sessionId;
+    private final String systemConfigJson;
+    private final String pluginTaskJson;
+    private final String processTaskJson;
+    private final List<PluginArchive.GemSpec> gemSpecs;
+    private final byte[] pluginArchiveBytes;
+
+    private final ProcessState state;
+    private final CountDownLatch timer;
+    private final int inputTaskCount;
+    private final ModelManager modelManager;
+    private volatile boolean isFinished;
+    private final Map<Integer, String> errorMessages;
+
+    SessionState(
+            String systemConfigJson, String pluginTaskJson, String processTaskJson,
+            List<PluginArchive.GemSpec> gemSpecs, byte[] pluginArchiveBytes,
+            ProcessState state, int inputTaskCount, ModelManager modelManager) {
+        this.sessionId = UUID.randomUUID().toString();
+        this.systemConfigJson = systemConfigJson;
+        this.pluginTaskJson = pluginTaskJson;
+        this.processTaskJson = processTaskJson;
+        this.gemSpecs = gemSpecs;
+        this.pluginArchiveBytes = pluginArchiveBytes;
+        this.state = state;
+        this.timer = new CountDownLatch(inputTaskCount);
+        this.inputTaskCount = inputTaskCount;
+        this.modelManager = modelManager;
+        this.isFinished = false;
+        this.errorMessages = new ConcurrentHashMap<>();
+    }
+
+    String getSessionId() {
+        return sessionId;
+    }
+
+    String getSystemConfigJson() {
+        return systemConfigJson;
+    }
+
+    String getPluginTaskJson() {
+        return pluginTaskJson;
+    }
+
+    String getProcessTaskJson() {
+        return processTaskJson;
+    }
+
+    List<PluginArchive.GemSpec> getGemSpecs() {
+        return gemSpecs;
+    }
+
+    byte[] getPluginArchiveBytes() {
+        return pluginArchiveBytes;
+    }
+
+    ProcessState getState() {
+        return state;
+    }
+
+    boolean isFinished() {
+        return isFinished;
+    }
+
+    synchronized void update(UpdateTaskStateData data) {
+        switch (data.getTaskState()) {
+            case STARTED:
+                state.getInputTaskState(data.getTaskIndex()).start();
+                state.getOutputTaskState(data.getTaskIndex()).start();
+                break;
+            case INPUT_COMMITTED:
+                state.getInputTaskState(data.getTaskIndex()).setTaskReport(getTaskReport(data.getTaskReport()));
+                break;
+            case OUTPUT_COMMITTED:
+                state.getOutputTaskState(data.getTaskIndex()).setTaskReport(getTaskReport(data.getTaskReport()));
+                break;
+            case FAILED:
+                errorMessages.put(data.getTaskIndex(), data.getErrorMessage());
+                timer.countDown();
+                break;
+            case FINISHED:
+                state.getInputTaskState(data.getTaskIndex()).finish();
+                state.getOutputTaskState(data.getTaskIndex()).finish();
+                timer.countDown();
+                showProgress(state, inputTaskCount);
+                break;
+        }
+    }
+
+    void waitUntilCompleted(int timeoutSeconds) throws InterruptedException, TimeoutException {
+        try {
+            if (!timer.await(timeoutSeconds, TimeUnit.SECONDS)) {
+                throw new TimeoutException(String.format("The session (%s) was time-out.", sessionId));
+            }
+            if (!errorMessages.isEmpty()) {
+                String message = errorMessages.entrySet().stream()
+                        .map(e -> String.format("%d: %s", e.getKey(), e.getValue()))
+                        .collect(Collectors.joining(System.lineSeparator()));
+                throw new TaskExecutionException(message);
+            }
+        } finally {
+            isFinished = true;
+        }
+    }
+
+    private TaskReport getTaskReport(String json) {
+        return modelManager.readObject(TaskReport.class, json);
+    }
+
+    private static void showProgress(ProcessState state, int taskCount) {
+        int started = 0;
+        int finished = 0;
+        for (int i = 0; i < taskCount; i++) {
+            if (state.getOutputTaskState(i).isStarted()) {
+                started++;
+            }
+            if (state.getOutputTaskState(i).isFinished()) {
+                finished++;
+            }
+        }
+        log.info(String.format("{done:%3d / %d, running: %d}", finished, taskCount, started - finished));
+    }
+}