embulk-executor-remoteserver 0.1.0

data/gradlew.bat

@@ -0,0 +1,84 @@
+@if "%DEBUG%" == "" @echo off
+@rem ##########################################################################
+@rem
+@rem  Gradle startup script for Windows
+@rem
+@rem ##########################################################################
+
+@rem Set local scope for the variables with windows NT shell
+if "%OS%"=="Windows_NT" setlocal
+
+set DIRNAME=%~dp0
+if "%DIRNAME%" == "" set DIRNAME=.
+set APP_BASE_NAME=%~n0
+set APP_HOME=%DIRNAME%
+
+@rem Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script.
+set DEFAULT_JVM_OPTS="-Xmx64m"
+
+@rem Find java.exe
+if defined JAVA_HOME goto findJavaFromJavaHome
+
+set JAVA_EXE=java.exe
+%JAVA_EXE% -version >NUL 2>&1
+if "%ERRORLEVEL%" == "0" goto init
+
+echo.
+echo ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH.
+echo.
+echo Please set the JAVA_HOME variable in your environment to match the
+echo location of your Java installation.
+
+goto fail
+
+:findJavaFromJavaHome
+set JAVA_HOME=%JAVA_HOME:"=%
+set JAVA_EXE=%JAVA_HOME%/bin/java.exe
+
+if exist "%JAVA_EXE%" goto init
+
+echo.
+echo ERROR: JAVA_HOME is set to an invalid directory: %JAVA_HOME%
+echo.
+echo Please set the JAVA_HOME variable in your environment to match the
+echo location of your Java installation.
+
+goto fail
+
+:init
+@rem Get command-line arguments, handling Windows variants
+
+if not "%OS%" == "Windows_NT" goto win9xME_args
+
+:win9xME_args
+@rem Slurp the command line arguments.
+set CMD_LINE_ARGS=
+set _SKIP=2
+
+:win9xME_args_slurp
+if "x%~1" == "x" goto execute
+
+set CMD_LINE_ARGS=%*
+
+:execute
+@rem Setup the command line
+
+set CLASSPATH=%APP_HOME%\gradle\wrapper\gradle-wrapper.jar
+
+@rem Execute Gradle
+"%JAVA_EXE%" %DEFAULT_JVM_OPTS% %JAVA_OPTS% %GRADLE_OPTS% "-Dorg.gradle.appname=%APP_BASE_NAME%" -classpath "%CLASSPATH%" org.gradle.wrapper.GradleWrapperMain %CMD_LINE_ARGS%
+
+:end
+@rem End local scope for the variables with windows NT shell
+if "%ERRORLEVEL%"=="0" goto mainEnd
+
+:fail
+rem Set variable GRADLE_EXIT_CONSOLE if you need the _script_ return code instead of
+rem the _cmd.exe /c_ return code!
+if  not "" == "%GRADLE_EXIT_CONSOLE%" exit 1
+exit /b 1
+
+:mainEnd
+if "%OS%"=="Windows_NT" endlocal
+
+:omega

data/lib/embulk/executor/remoteserver.rb

@@ -0,0 +1,3 @@
+Embulk::JavaPlugin.register_executor(
+  "remoteserver", "org.embulk.executor.remoteserver.RemoteServerExecutor",
+  File.expand_path('../../../../classpath', __FILE__))

data/settings.gradle

@@ -0,0 +1 @@
+rootProject.name = 'embulk-executor-remoteserver'

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

@@ -0,0 +1,111 @@
+package org.embulk.executor.remoteserver;
+
+import com.github.kamatama41.nsocket.CommandListener;
+import com.github.kamatama41.nsocket.Connection;
+import com.github.kamatama41.nsocket.SocketClient;
+import org.embulk.config.ConfigException;
+
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.atomic.AtomicInteger;
+
+class EmbulkClient implements AutoCloseable {
+    private final SocketClient client;
+    private final List<Host> hosts;
+    private final SessionState sessionState;
+    private final AtomicInteger counter = new AtomicInteger(1);
+
+    private EmbulkClient(SocketClient client, List<Host> hosts, SessionState sessionState) {
+        this.client = client;
+        this.hosts = hosts;
+        this.sessionState = sessionState;
+    }
+
+    static EmbulkClient open(
+            SessionState sessionState,
+            List<Host> hosts) throws IOException {
+        SocketClient client = new SocketClient();
+        client.registerSyncCommand(new InitializeSessionCommand(null));
+        client.registerSyncCommand(new RemoveSessionCommand(null));
+        client.registerCommand(new NotifyTaskStateCommand(sessionState));
+        client.registerListener(new Reconnector(client, sessionState));
+        client.open();
+        for (Host host : hosts) {
+            client.addNode(host.toAddress());
+        }
+        return new EmbulkClient(client, hosts, sessionState);
+    }
+
+    void createSession() {
+        ExecutorService es = Executors.newFixedThreadPool(hosts.size());
+        List<Future> futures = new ArrayList<>();
+        for (Host host : hosts) {
+            futures.add(es.submit(() -> {
+                Connection connection = client.getConnection(host.toAddress());
+                connection.sendSyncCommand(InitializeSessionCommand.ID, toInitializeSessionData(sessionState));
+            }));
+        }
+        try {
+            for (Future future : futures) {
+                future.get();
+            }
+        } catch (ExecutionException e) {
+            throw new ConfigException(e.getCause());
+        } catch (Exception e) {
+            throw new ConfigException(e);
+        } finally {
+            es.shutdown();
+        }
+    }
+
+    void startTask(int taskIndex) {
+        // Round robin (more smart logic needed?)
+        InetSocketAddress target = hosts.get(counter.getAndIncrement() % hosts.size()).toAddress();
+        client.getConnection(target).sendCommand(
+                StartTaskCommand.ID, new StartTaskCommand.Data(sessionState.getSessionId(), taskIndex));
+    }
+
+    @Override
+    public void close() throws IOException {
+        for (Host host : hosts) {
+            Connection connection = client.getConnection(host.toAddress());
+            connection.sendSyncCommand(RemoveSessionCommand.ID, sessionState.getSessionId());
+        }
+        client.close();
+    }
+
+    private static class Reconnector implements CommandListener {
+        private final SocketClient client;
+        private final SessionState sessionState;
+
+        Reconnector(SocketClient client, SessionState sessionState) {
+            this.client = client;
+            this.sessionState = sessionState;
+        }
+
+        @Override
+        public void onDisconnected(Connection connection) {
+            if(!sessionState.isFinished()) {
+                Connection newConnection = client.getConnection((InetSocketAddress) connection.getRemoteSocketAddress());
+                newConnection.sendSyncCommand(InitializeSessionCommand.ID, toInitializeSessionData(sessionState));
+            }
+        }
+    }
+
+    private static InitializeSessionCommand.Data toInitializeSessionData(SessionState sessionState) {
+        return new InitializeSessionCommand.Data(
+                sessionState.getSessionId(),
+                sessionState.getSystemConfigJson(),
+                sessionState.getPluginTaskJson(),
+                sessionState.getProcessTaskJson(),
+                sessionState.getGemSpecs(),
+                sessionState.getPluginArchiveBytes()
+        );
+    }
+}

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

@@ -0,0 +1,32 @@
+package org.embulk.executor.remoteserver;
+
+import com.github.kamatama41.nsocket.SocketServer;
+
+import java.io.IOException;
+
+public class EmbulkServer implements AutoCloseable {
+    private SocketServer server;
+
+    EmbulkServer(SocketServer  server) {
+        this.server = server;
+    }
+
+    static EmbulkServer start(String host, int port, int numOfWorkers) throws IOException {
+        SocketServer server = new SocketServer();
+        SessionManager sessionManager = new SessionManager();
+        server.setHost(host);
+        server.setPort(port);
+        server.setDefaultContentBufferSize(4 * 1024 * 1024); // 4MB
+        server.setNumOfWorkers(numOfWorkers);
+        server.registerSyncCommand(new InitializeSessionCommand(sessionManager));
+        server.registerSyncCommand(new RemoveSessionCommand(sessionManager));
+        server.registerCommand(new StartTaskCommand(sessionManager));
+        server.start();
+        return new EmbulkServer(server);
+    }
+
+    @Override
+    public void close() throws IOException {
+        server.stop();
+    }
+}

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

@@ -0,0 +1,32 @@
+package org.embulk.executor.remoteserver;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.net.InetSocketAddress;
+
+class Host {
+    private String name;
+    private int port;
+
+    @JsonCreator
+    Host(@JsonProperty("name") String name,
+         @JsonProperty("port") int port) {
+        this.name = name;
+        this.port = port;
+    }
+
+    @JsonProperty
+    String getName() {
+        return name;
+    }
+
+    @JsonProperty
+    int getPort() {
+        return port;
+    }
+
+    InetSocketAddress toAddress() {
+        return new InetSocketAddress(name, port);
+    }
+}

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

@@ -0,0 +1,94 @@
+package org.embulk.executor.remoteserver;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.github.kamatama41.nsocket.Connection;
+import com.github.kamatama41.nsocket.SyncCommand;
+
+import java.util.List;
+
+class InitializeSessionCommand implements SyncCommand<InitializeSessionCommand.Data, Void> {
+    static final String ID = "initialize_session";
+    private final SessionManager sessionManager;
+
+    InitializeSessionCommand(SessionManager sessionManager) {
+        this.sessionManager = sessionManager;
+    }
+
+    @Override
+    public Void apply(Data data, Connection connection) {
+        sessionManager.registerNewSession(
+                data.getSessionId(),
+                data.getSystemConfigJson(),
+                data.getPluginTaskJson(),
+                data.getProcessTaskJson(),
+                data.getGemSpecs(),
+                data.getPluginArchive(),
+                connection);
+        return null;
+    }
+
+    @Override
+    public long getTimeoutMillis() {
+        return 60000L;
+    }
+
+    @Override
+    public String getId() {
+        return ID;
+    }
+
+    static class Data {
+        private String sessionId;
+        private String systemConfigJson;
+        private String pluginTaskJson;
+        private String processTaskJson;
+        private List<PluginArchive.GemSpec> gemSpecs;
+        private byte[] pluginArchive;
+
+        @JsonCreator
+        Data(@JsonProperty("sessionId") String sessionId,
+             @JsonProperty("systemConfigJson") String systemConfigJson,
+             @JsonProperty("pluginTaskJson") String pluginTaskJson,
+             @JsonProperty("processTaskJson") String processTaskJson,
+             @JsonProperty("gemSpecs") List<PluginArchive.GemSpec>  gemSpecs,
+             @JsonProperty("pluginArchive") byte[] pluginArchive) {
+            this.sessionId = sessionId;
+            this.systemConfigJson = systemConfigJson;
+            this.pluginTaskJson = pluginTaskJson;
+            this.processTaskJson = processTaskJson;
+            this.gemSpecs = gemSpecs;
+            this.pluginArchive = pluginArchive;
+        }
+
+        @JsonProperty
+        String getSessionId() {
+            return sessionId;
+        }
+
+        @JsonProperty
+        String getSystemConfigJson() {
+            return systemConfigJson;
+        }
+
+        @JsonProperty
+        String getPluginTaskJson() {
+            return pluginTaskJson;
+        }
+
+        @JsonProperty
+        String getProcessTaskJson() {
+            return processTaskJson;
+        }
+
+        @JsonProperty
+        List<PluginArchive.GemSpec> getGemSpecs() {
+            return gemSpecs;
+        }
+
+        @JsonProperty
+        byte[] getPluginArchive() {
+            return pluginArchive;
+        }
+    }
+}

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

@@ -0,0 +1,25 @@
+package org.embulk.executor.remoteserver;
+
+import ch.qos.logback.classic.Level;
+import ch.qos.logback.classic.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.util.Map;
+
+public class Launcher {
+    public static void main(String[] args) throws IOException {
+        Map<String, String> envVars = System.getenv();
+        String host = envVars.getOrDefault("BIND_ADDRESS", "0.0.0.0");
+        int port = Integer.parseInt(envVars.getOrDefault("PORT", "30001"));
+        int numOfWorkers = Integer.parseInt(envVars.getOrDefault("NUM_OF_WORKERS", "1"));
+        Level logLevel = Level.toLevel(envVars.getOrDefault("LOG_LEVEL", "info"));
+        configureLogLevel(logLevel);
+        EmbulkServer.start(host, port, numOfWorkers);
+    }
+
+    private static void configureLogLevel(Level logLevel) {
+        Logger rootLogger = (Logger) LoggerFactory.getLogger(Logger.ROOT_LOGGER_NAME);
+        rootLogger.setLevel(logLevel);
+    }
+}

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

@@ -0,0 +1,23 @@
+package org.embulk.executor.remoteserver;
+
+import com.github.kamatama41.nsocket.Command;
+import com.github.kamatama41.nsocket.Connection;
+
+class NotifyTaskStateCommand implements Command<UpdateTaskStateData> {
+    static final String ID = "notify_task_state";
+    private final SessionState sessionState;
+
+    NotifyTaskStateCommand(SessionState sessionState) {
+        this.sessionState = sessionState;
+    }
+
+    @Override
+    public void execute(UpdateTaskStateData data, Connection connection) throws Exception {
+        sessionState.update(data);
+    }
+
+    @Override
+    public String getId() {
+        return ID;
+    }
+}

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

@@ -0,0 +1,170 @@
+package org.embulk.executor.remoteserver;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.collect.ImmutableList;
+import com.google.common.io.ByteStreams;
+import org.jruby.embed.ScriptingContainer;
+
+import java.io.File;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.nio.file.DirectoryStream;
+import java.nio.file.Files;
+import java.nio.file.NoSuchFileException;
+import java.nio.file.NotDirectoryException;
+import java.nio.file.Path;
+import java.util.List;
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipInputStream;
+import java.util.zip.ZipOutputStream;
+
+// Copied from embulk-executor-mapreduce
+public class PluginArchive {
+    public static class GemSpec {
+        private final String name;
+        private final List<String> requirePaths;
+
+        @JsonCreator
+        public GemSpec(
+                @JsonProperty("name") String name,
+                @JsonProperty("requirePaths") List<String> requirePaths) {
+            this.name = name;
+            this.requirePaths = requirePaths;
+        }
+
+        @JsonProperty("name")
+        public String getName() {
+            return name;
+        }
+
+        @JsonProperty("requirePaths")
+        public List<String> getRequirePaths() {
+            return requirePaths;
+        }
+    }
+
+    private static class LocalGem
+            extends GemSpec {
+        private final File localPath;
+
+        public LocalGem(File localPath, String name, List<String> requirePaths) {
+            super(name, requirePaths);
+            this.localPath = localPath;
+        }
+
+        public File getLocalPath() {
+            return localPath;
+        }
+    }
+
+    public static class Builder {
+        private final ImmutableList.Builder<LocalGem> localGems = ImmutableList.builder();
+
+        @SuppressWarnings("unchecked")
+        public Builder addLoadedRubyGems(ScriptingContainer jruby) {
+            List<List<String>> tuples = (List<List<String>>) jruby.runScriptlet("Gem.loaded_specs.map {|k,v| [k, v.full_gem_path, v.require_paths].flatten }");
+            for (List<String> tuple : tuples) {
+                String name = tuple.remove(0);
+                String fullGemPath = tuple.remove(0);
+                List<String> requirePaths = ImmutableList.copyOf(tuple);
+                addSpec(new File(fullGemPath), name, requirePaths);
+            }
+            return this;
+        }
+
+        public Builder addSpec(File localPath, String name, List<String> requirePaths) {
+            localGems.add(new LocalGem(localPath, name, requirePaths));
+            return this;
+        }
+
+        public PluginArchive build() {
+            return new PluginArchive(localGems.build());
+        }
+    }
+
+    private final List<LocalGem> localGems;
+
+    private PluginArchive(List<LocalGem> localGems) {
+        this.localGems = localGems;
+    }
+
+    @SuppressWarnings("unchecked")
+    public void restoreLoadPathsTo(ScriptingContainer jruby) {
+        List<String> loadPaths = (List<String>) jruby.runScriptlet("$LOAD_PATH");
+        for (LocalGem localGem : localGems) {
+            Path localGemPath = localGem.getLocalPath().toPath();
+            for (String requirePath : localGem.getRequirePaths()) {
+                loadPaths.add(localGemPath.resolve(requirePath).toString());
+            }
+        }
+        jruby.setLoadPaths(loadPaths);
+    }
+
+    public List<GemSpec> dump(OutputStream out)
+            throws IOException {
+        ImmutableList.Builder<GemSpec> builder = ImmutableList.builder();
+        try (ZipOutputStream zip = new ZipOutputStream(out)) {
+            for (LocalGem localGem : localGems) {
+                zipDirectory(zip, localGem.getLocalPath().toPath(), localGem.getName() + "/");
+                builder.add(new GemSpec(localGem.getName(), localGem.getRequirePaths()));
+            }
+        }
+        return builder.build();
+    }
+
+    private static void zipDirectory(ZipOutputStream zip, Path directory, String name)
+            throws IOException {
+        try (DirectoryStream<Path> dirStream = Files.newDirectoryStream(directory)) {
+            for (Path path : dirStream) {
+                if (Files.isDirectory(path)) {
+                    zipDirectory(zip, path, name + path.getFileName() + "/");
+                } else {
+                    zip.putNextEntry(new ZipEntry(name + path.getFileName()));
+                    try (InputStream in = Files.newInputStream(path)) {
+                        ByteStreams.copy(in, zip);
+                    }
+                    zip.closeEntry();
+                }
+            }
+        } catch (NoSuchFileException | NotDirectoryException ex) {
+            // ignore
+        }
+    }
+
+    public static PluginArchive load(File localDirectory, List<GemSpec> gemSpecs,
+                                     InputStream in) throws IOException {
+        try (ZipInputStream zip = new ZipInputStream(in)) {
+            unzipDirectory(zip, localDirectory.toPath());
+        }
+
+        ImmutableList.Builder<LocalGem> builder = ImmutableList.builder();
+        for (GemSpec gemSpec : gemSpecs) {
+            builder.add(new LocalGem(
+                    new File(localDirectory, gemSpec.getName()),
+                    gemSpec.getName(),
+                    gemSpec.getRequirePaths()));
+        }
+        return new PluginArchive(builder.build());
+    }
+
+    private static void unzipDirectory(ZipInputStream zip, Path directory)
+            throws IOException {
+        while (true) {
+            ZipEntry entry = zip.getNextEntry();
+            if (entry == null) {
+                break;
+            }
+            Path path = directory.resolve(entry.getName());
+            if (entry.getName().endsWith("/")) {
+                Files.createDirectories(path);
+            } else {
+                Files.createDirectories(path.getParent());
+                try (OutputStream out = Files.newOutputStream(path)) {
+                    ByteStreams.copy(zip, out);
+                }
+            }
+        }
+    }
+}