embulk-filter-copy 0.0.1

data/src/main/java/org/embulk/filter/copy/forward/OutForwardService.java

@@ -0,0 +1,170 @@
+package org.embulk.filter.copy.forward;
+
+import com.google.common.base.Optional;
+import com.google.common.collect.Maps;
+import org.embulk.config.Config;
+import org.embulk.config.ConfigDefault;
+import org.embulk.spi.Exec;
+import org.komamitsu.fluency.Fluency;
+import org.slf4j.Logger;
+
+import java.io.IOException;
+import java.util.Map;
+
+public class OutForwardService
+{
+    private final static Logger logger = Exec.getLogger(OutForwardService.class);
+
+    public interface OutForwardTask
+            extends org.embulk.config.Task
+    {
+        @Config("host")
+        @ConfigDefault("\"localhost\"")
+        String getHost();
+
+        @Config("port")
+        @ConfigDefault("24224")
+        int getPort();
+
+        @Config("max_buffer_size")
+        @ConfigDefault("null")
+        Optional<Long> getMaxBufferSize();
+
+        @Config("buffer_chunk_initial_size")
+        @ConfigDefault("null")
+        Optional<Integer> getBufferChunkInitialSize();
+
+        @Config("buffer_chunk_retention_size")
+        @ConfigDefault("null")
+        Optional<Integer> getBufferChunkRetentionSize();
+
+        @Config("flush_interval_millis")
+        @ConfigDefault("null")
+        Optional<Integer> getFlushIntervalMillis();
+
+        @Config("sender_max_retry_count")
+        @ConfigDefault("null")
+        Optional<Integer> getSenderMaxRetryCount();
+
+        @Config("ack_response_mode")
+        @ConfigDefault("null")
+        Optional<Boolean> getAckResponseMode();
+
+        @Config("file_backup_dir")
+        @ConfigDefault("null")
+        Optional<String> getFileBackupDir();
+
+        @Config("wait_until_buffer_flushed")
+        @ConfigDefault("null")
+        Optional<Integer> getWaitUntilBufferFlushed();
+
+        @Config("wait_until_flusher_terminated")
+        @ConfigDefault("null")
+        Optional<Integer> getWaitUntilFlusherTerminated();
+    }
+
+    public interface Task
+            extends ForwardBaseTask
+    {
+        @Config("out_forward")
+        @ConfigDefault("{}")
+        OutForwardTask getOutForwardTask();
+    }
+
+    public static void sendShutdownMessage(Task task)
+    {
+        logger.info("out_forward: send shutdown message.");
+        OutForwardService outForward = new OutForwardService(task);
+        outForward.emit(task.getShutdownTag(), Maps.newHashMap());
+        outForward.finish();
+        outForward.close();
+    }
+
+    private final Task task;
+    private final Fluency client;
+
+    public OutForwardService(Task task)
+    {
+        this.task = task;
+        this.client = newFluency(task.getOutForwardTask());
+    }
+
+    private Fluency.Config configureFluencyConfig(OutForwardTask t)
+    {
+        Fluency.Config c = new Fluency.Config();
+        if (t.getMaxBufferSize().isPresent()) {
+            c.setMaxBufferSize(t.getMaxBufferSize().get());
+        }
+        if (t.getBufferChunkInitialSize().isPresent()) {
+            c.setBufferChunkInitialSize(t.getBufferChunkInitialSize().get());
+        }
+        if (t.getBufferChunkRetentionSize().isPresent()) {
+            c.setBufferChunkRetentionSize(t.getBufferChunkRetentionSize().get());
+        }
+        if (t.getFlushIntervalMillis().isPresent()) {
+            c.setFlushIntervalMillis(t.getFlushIntervalMillis().get());
+        }
+        if (t.getSenderMaxRetryCount().isPresent()) {
+            c.setSenderMaxRetryCount(t.getSenderMaxRetryCount().get());
+        }
+        if (t.getAckResponseMode().isPresent()) {
+            c.setAckResponseMode(t.getAckResponseMode().get());
+        }
+        if (t.getFileBackupDir().isPresent()) {
+            c.setFileBackupDir(t.getFileBackupDir().get());
+        }
+        if (t.getWaitUntilBufferFlushed().isPresent()) {
+            c.setWaitUntilBufferFlushed(t.getWaitUntilBufferFlushed().get());
+        }
+        if (t.getWaitUntilFlusherTerminated().isPresent()) {
+            c.setWaitUntilFlusherTerminated(t.getWaitUntilFlusherTerminated().get());
+        }
+        return c;
+    }
+
+    private Fluency newFluency(OutForwardTask t)
+    {
+        Fluency.Config c = configureFluencyConfig(t);
+        try {
+            return Fluency.defaultFluency(t.getHost(), t.getPort(), c);
+        }
+        catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    public void emit(String tag, Map<String, Object> message)
+    {
+        try {
+            client.emit(tag, message);
+        }
+        catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    public void emit(Map<String, Object> message)
+    {
+        emit(task.getMessageTag(), message);
+    }
+
+    public void finish()
+    {
+        try {
+            client.flush();
+        }
+        catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    public void close()
+    {
+        try {
+            client.close();
+        }
+        catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+}

data/src/main/java/org/embulk/filter/copy/forward/OutForwardVisitor.java

@@ -0,0 +1,63 @@
+package org.embulk.filter.copy.forward;
+
+import org.embulk.spi.Column;
+import org.embulk.spi.ColumnVisitor;
+import org.embulk.spi.PageReader;
+
+public class OutForwardVisitor
+    implements ColumnVisitor
+{
+    private final PageReader reader;
+    private final OutForwardEventBuilder builder;
+
+    public OutForwardVisitor(PageReader reader, OutForwardEventBuilder builder)
+    {
+        this.reader = reader;
+        this.builder = builder;
+    }
+
+    private void nullOr(Column column, Runnable r)
+    {
+        if (reader.isNull(column)) {
+            builder.setNull(column);
+            return;
+        }
+        r.run();
+    }
+
+    @Override
+    public void booleanColumn(Column column)
+    {
+        nullOr(column, () -> builder.setBoolean(column, reader.getBoolean(column)));
+    }
+
+    @Override
+    public void longColumn(Column column)
+    {
+        nullOr(column, () -> builder.setLong(column, reader.getLong(column)));
+    }
+
+    @Override
+    public void doubleColumn(Column column)
+    {
+        nullOr(column, () -> builder.setDouble(column, reader.getDouble(column)));
+    }
+
+    @Override
+    public void stringColumn(Column column)
+    {
+        nullOr(column, () -> builder.setString(column, reader.getString(column)));
+    }
+
+    @Override
+    public void timestampColumn(Column column)
+    {
+        nullOr(column, () -> builder.setTimestamp(column, reader.getTimestamp(column)));
+    }
+
+    @Override
+    public void jsonColumn(Column column)
+    {
+        nullOr(column, () -> builder.setJson(column, reader.getJson(column)));
+    }
+}

data/src/main/java/org/embulk/filter/copy/plugin/InternalForwardInputPlugin.java

@@ -0,0 +1,111 @@
+package org.embulk.filter.copy.plugin;
+
+import org.embulk.config.Config;
+import org.embulk.config.ConfigDiff;
+import org.embulk.config.ConfigInject;
+import org.embulk.config.ConfigSource;
+import org.embulk.config.Task;
+import org.embulk.config.TaskReport;
+import org.embulk.config.TaskSource;
+import org.embulk.filter.copy.forward.InForwardEventReader;
+import org.embulk.filter.copy.forward.InForwardService;
+import org.embulk.filter.copy.forward.InForwardVisitor;
+import org.embulk.spi.BufferAllocator;
+import org.embulk.spi.Exec;
+import org.embulk.spi.InputPlugin;
+import org.embulk.spi.PageBuilder;
+import org.embulk.spi.PageOutput;
+import org.embulk.spi.Schema;
+import org.embulk.spi.SchemaConfig;
+import org.embulk.spi.time.TimestampParser;
+import org.slf4j.Logger;
+
+import java.util.List;
+
+public class InternalForwardInputPlugin
+        implements InputPlugin
+{
+    public final static String PLUGIN_NAME = "internal_forward";
+    private final static Logger logger = Exec.getLogger(InternalForwardInputPlugin.class);
+
+    public interface PluginTask
+            extends Task, TimestampParser.Task, InForwardService.Task
+    {
+        @Config("columns")
+        SchemaConfig getColumns();
+
+        @ConfigInject
+        BufferAllocator getBufferAllocator();
+    }
+
+    @Override
+    public ConfigDiff transaction(ConfigSource config,
+            InputPlugin.Control control)
+    {
+        PluginTask task = config.loadConfig(PluginTask.class);
+
+        Schema schema = task.getColumns().toSchema();
+        int taskCount = 1;  // number of run() method calls
+
+        return resume(task.dump(), schema, taskCount, control);
+    }
+
+    @Override
+    public ConfigDiff resume(TaskSource taskSource,
+            Schema schema, int taskCount,
+            InputPlugin.Control control)
+    {
+        control.run(taskSource, schema, taskCount);
+        return Exec.newConfigDiff();
+    }
+
+    @Override
+    public void cleanup(TaskSource taskSource,
+            Schema schema, int taskCount,
+            List<TaskReport> successTaskReports)
+    {
+    }
+
+    @Override
+    public TaskReport run(TaskSource taskSource,
+            Schema schema, int taskIndex,
+            PageOutput output)
+    {
+        PluginTask task = taskSource.loadTask(PluginTask.class);
+
+        try (PageBuilder pageBuilder = new PageBuilder(task.getBufferAllocator(), schema, output)) {
+            TimestampParser timestampParser = new TimestampParser(
+                    task.getJRuby(),
+                    task.getDefaultTimestampFormat(),
+                    task.getDefaultTimeZone());
+            InForwardEventReader eventReader = new InForwardEventReader(schema, timestampParser);
+            InForwardVisitor inForwardVisitor = new InForwardVisitor(eventReader, pageBuilder);
+
+            InForwardService.builder()
+                    .task(task)
+                    .forEachEventCallback(
+                            event ->
+                            {
+                                // TODO: here is not thread-safe
+                                eventReader.setEvent(event);
+                                while (eventReader.nextMessage()) {
+                                    schema.visitColumns(inForwardVisitor);
+                                    pageBuilder.addRecord();
+                                }
+                            }
+                    )
+                    .build()
+                    .runUntilShouldShutdown();
+
+            pageBuilder.finish();
+        }
+
+        return Exec.newTaskReport(); // TODO
+    }
+
+    @Override
+    public ConfigDiff guess(ConfigSource config)
+    {
+        return Exec.newConfigDiff();
+    }
+}

data/src/main/java/org/embulk/filter/copy/service/EmbulkExecutorService.java

@@ -0,0 +1,111 @@
+package org.embulk.filter.copy.service;
+
+import com.google.common.util.concurrent.FutureCallback;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.ListeningExecutorService;
+import com.google.common.util.concurrent.MoreExecutors;
+import com.google.inject.Injector;
+import org.embulk.EmbulkEmbed;
+import org.embulk.config.ConfigException;
+import org.embulk.config.ConfigSource;
+import org.embulk.exec.ExecutionResult;
+import org.embulk.filter.copy.util.ElapsedTime;
+import org.embulk.guice.LifeCycleInjector;
+import org.embulk.spi.Exec;
+import org.slf4j.Logger;
+
+import javax.annotation.Nullable;
+
+import java.lang.reflect.Constructor;
+import java.lang.reflect.InvocationTargetException;
+import java.util.concurrent.Callable;
+import java.util.concurrent.Executors;
+
+public class EmbulkExecutorService
+{
+    private final static String THREAD_NAME = "embulk executor service";
+    private static final int NUM_THREADS = 1;
+    private final static Logger logger = Exec.getLogger(EmbulkExecutorService.class);
+    private final Injector injector;
+    private final ListeningExecutorService es;
+    private ListenableFuture<ExecutionResult> future;
+
+    public EmbulkExecutorService(Injector injector)
+    {
+        this.injector = injector;
+        this.es = MoreExecutors.listeningDecorator(
+                Executors.newFixedThreadPool(
+                        NUM_THREADS,
+                        r -> new Thread(r, THREAD_NAME)
+                ));
+    }
+
+    public void executeAsync(ConfigSource config)
+    {
+        logger.debug("execute with this config: {}", config);
+        if (future != null) {
+            throw new IllegalStateException("executeAsync is already called.");
+        }
+        future = es.submit(embulkRun(config));
+        Futures.addCallback(future, resultFutureCallback());
+    }
+
+    public void shutdown()
+    {
+        ElapsedTime.info(
+                logger,
+                "embulk executor service shutdown",
+                es::shutdown);
+    }
+
+    public void waitExecutionFinished()
+    {
+        if (future == null) {
+            throw new NullPointerException();
+        }
+
+        ElapsedTime.debugUntil(() -> future.isDone() || future.isCancelled(),
+                logger, "embulk executor", 3000L);
+    }
+
+    private Callable<ExecutionResult> embulkRun(ConfigSource config)
+    {
+        return () -> newEmbulkEmbed(injector).run(config);
+    }
+
+    private EmbulkEmbed newEmbulkEmbed(Injector injector)
+    {
+        try {
+            Constructor<EmbulkEmbed> constructor = EmbulkEmbed.class
+                    .getDeclaredConstructor(ConfigSource.class, LifeCycleInjector.class);
+            constructor.setAccessible(true);
+            return constructor.newInstance(Exec.newConfigSource(), injector);
+        }
+        catch (InstantiationException | IllegalAccessException | NoSuchMethodException | InvocationTargetException e) {
+            throw new ConfigException(e);
+        }
+    }
+
+    private FutureCallback<ExecutionResult> resultFutureCallback()
+    {
+        return new FutureCallback<ExecutionResult>()
+        {
+            @Override
+            public void onSuccess(@Nullable ExecutionResult result)
+            {
+                for (Throwable throwable : result.getIgnoredExceptions()) {
+                    logger.warn("Ignored error ", throwable);
+                }
+                logger.info("Config diff: {}", result.getConfigDiff());
+                logger.debug("ExecutionResult: {}", result);
+            }
+
+            @Override
+            public void onFailure(Throwable t)
+            {
+                throw new RuntimeException(t);
+            }
+        };
+    }
+}