embulk-filter-copy 0.0.1

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

@@ -0,0 +1,17 @@
+package org.embulk.filter.copy.forward;
+
+import org.embulk.config.Config;
+import org.embulk.config.ConfigDefault;
+import org.embulk.config.Task;
+
+public interface ForwardBaseTask
+    extends Task
+{
+    @Config("shutdown_tag")
+    @ConfigDefault("\"shutdown\"")
+    String getShutdownTag();
+
+    @Config("message_tag")
+    @ConfigDefault("\"message\"")
+    String getMessageTag();
+}

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

@@ -0,0 +1,147 @@
+package org.embulk.filter.copy.forward;
+
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Maps;
+import influent.EventStream;
+import org.embulk.spi.Column;
+import org.embulk.spi.Schema;
+import org.embulk.spi.json.JsonParser;
+import org.embulk.spi.time.Timestamp;
+import org.embulk.spi.time.TimestampParser;
+import org.msgpack.value.Value;
+
+import java.util.Map;
+
+public class InForwardEventReader
+{
+    private final Schema schema;
+    private final Map<String, Column> columnMap;
+
+    private final JsonParser jsonParser = new JsonParser();
+    private final TimestampParser timestampParser;
+
+    private EventStream event = null;
+    private int eventMessageCount = 0;
+
+    private int readCount = 0;
+    private Map<String, Value> message = Maps.newHashMap();
+
+    public InForwardEventReader(Schema schema, TimestampParser timestampParser)
+    {
+        this.schema = schema;
+        ImmutableMap.Builder<String, Column> builder = ImmutableMap.builder();
+        schema.getColumns().forEach(column -> builder.put(column.getName(), column));
+        this.columnMap = builder.build();
+        this.timestampParser = timestampParser;
+    }
+
+    public Schema getSchema()
+    {
+        return schema;
+    }
+
+    private Column getColumn(String columnName)
+    {
+        return columnMap.get(columnName);
+    }
+
+    private Column getColumn(int columnIndex)
+    {
+        return getSchema().getColumn(columnIndex);
+    }
+
+    private Value getValue(Column column)
+    {
+        return message.get(column.getName());
+    }
+
+    public void setEvent(EventStream event)
+    {
+        this.event = event;
+        this.eventMessageCount = event.getEntries().size();
+    }
+
+    public boolean isNull(Column column)
+    {
+        return getValue(column).isNilValue();
+    }
+
+    public boolean isNull(int columnIndex)
+    {
+        return isNull(getColumn(columnIndex));
+    }
+
+    public boolean getBoolean(Column column)
+    {
+        return getValue(column).asBooleanValue().getBoolean();
+    }
+
+    public boolean getBoolean(int columnIndex)
+    {
+        return getBoolean(getColumn(columnIndex));
+    }
+
+    public long getLong(Column column)
+    {
+        return getValue(column).asIntegerValue().asLong();
+    }
+
+    public long getLong(int columnIndex)
+    {
+        return getLong(getColumn(columnIndex));
+    }
+
+    public double getDouble(Column column)
+    {
+        return getValue(column).asFloatValue().toDouble();
+    }
+
+    public double genDouble(int columnIndex)
+    {
+        return getDouble(getColumn(columnIndex));
+    }
+
+    public String getString(Column column)
+    {
+        return getValue(column).toString();
+    }
+
+    public String getString(int columnIndex)
+    {
+        return getString(getColumn(columnIndex));
+    }
+
+    public Timestamp getTimestamp(Column column)
+    {
+        return timestampParser.parse(getString(column));
+    }
+
+    public Timestamp getTimestamp(int columnIndex)
+    {
+        return getTimestamp(getColumn(columnIndex));
+    }
+
+    public Value getJson(Column column)
+    {
+        return jsonParser.parse(getString(column));
+    }
+
+    public Value getJson(int columnIndex)
+    {
+        return getJson(getColumn(columnIndex));
+    }
+
+    public boolean nextMessage()
+    {
+        if (eventMessageCount <= readCount) {
+            return false;
+        }
+
+        ImmutableMap.Builder<String, Value> builder = ImmutableMap.builder();
+        event.getEntries().get(readCount++).getRecord().entrySet().forEach(
+                entry -> builder.put(entry.getKey().toString(), entry.getValue())
+        );
+        message = builder.build();
+        return true;
+    }
+}

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

@@ -0,0 +1,187 @@
+package org.embulk.filter.copy.forward;
+
+import com.google.common.base.Optional;
+import influent.EventStream;
+import influent.Tag;
+import influent.forward.ForwardCallback;
+import influent.forward.ForwardServer;
+import org.embulk.config.Config;
+import org.embulk.config.ConfigDefault;
+import org.embulk.filter.copy.util.ElapsedTime;
+import org.embulk.spi.Exec;
+import org.slf4j.Logger;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Consumer;
+
+public class InForwardService
+{
+    private final static Logger logger = Exec.getLogger(InForwardService.class);
+
+    public interface InForwardTask
+            extends org.embulk.config.Task
+    {
+        @Config("port")
+        @ConfigDefault("24224")
+        int getPort();
+
+        @Config("chunk_size_limit")
+        @ConfigDefault("null")
+        Optional<Long> getChunkSizeLimit();
+
+        @Config("backlog")
+        @ConfigDefault("null")
+        Optional<Integer> getBacklog();
+
+        @Config("send_buffer_size")
+        @ConfigDefault("null")
+        Optional<Integer> getSendBufferSize();
+
+        @Config("receive_buffer_size")
+        @ConfigDefault("null")
+        Optional<Integer> getReceiveBufferSize();
+
+        @Config("keep_alive_enabled")
+        @ConfigDefault("null")
+        Optional<Boolean> getKeepAliveEnabled();
+
+        @Config("tcp_no_delay_enabled")
+        @ConfigDefault("null")
+        Optional<Boolean> getTcpNoDelayEnabled();
+    }
+
+    public interface Task
+            extends ForwardBaseTask
+    {
+        @Config("in_forward")
+        @ConfigDefault("{}")
+        InForwardTask getInForwardTask();
+    }
+
+    public static class Builder
+    {
+        private Task task;
+        private Consumer<EventStream> eventConsumer;
+
+        public Builder()
+        {
+        }
+
+        public Builder task(Task task)
+        {
+            this.task = task;
+            return this;
+        }
+
+        public Builder forEachEventCallback(Consumer<EventStream> eventConsumer)
+        {
+            this.eventConsumer = eventConsumer;
+            return this;
+        }
+
+        public InForwardService build()
+        {
+            return new InForwardService(task, eventConsumer);
+        }
+    }
+
+    public static Builder builder()
+    {
+        return new Builder();
+    }
+
+    // TODO: configure?
+    private static final String THREAD_NAME = "in_forward";
+    private static final int NUM_THREADS = 1;
+
+    private final Task task;
+    private final ForwardServer server;
+    private final ExecutorService callbackThread;
+    private final AtomicBoolean shouldShutdown = new AtomicBoolean(false);
+
+    private InForwardService(Task task, Consumer<EventStream> eventConsumer)
+    {
+        this.task = task;
+        this.callbackThread = Executors.newFixedThreadPool(
+                NUM_THREADS,
+                r -> new Thread(r, THREAD_NAME));
+        this.server = buildServer(task.getInForwardTask(), eventConsumer, callbackThread);
+    }
+
+    private ForwardServer buildServer(InForwardTask t, Consumer<EventStream> eventConsumer, Executor callbackThread)
+    {
+        ForwardServer.Builder builder = new ForwardServer.Builder(
+                ForwardCallback.ofSyncConsumer(
+                        wrapEventConsumer(eventConsumer),
+                        callbackThread));
+
+        builder.localAddress(t.getPort());
+
+        if (t.getChunkSizeLimit().isPresent()) {
+            builder.chunkSizeLimit(t.getChunkSizeLimit().get());
+        }
+        if (t.getBacklog().isPresent()) {
+            builder.backlog(t.getBacklog().get());
+        }
+        if (t.getSendBufferSize().isPresent()) {
+            builder.sendBufferSize(t.getSendBufferSize().get());
+        }
+        if (t.getReceiveBufferSize().isPresent()) {
+            builder.receiveBufferSize(t.getReceiveBufferSize().get());
+        }
+        if (t.getKeepAliveEnabled().isPresent()) {
+            builder.keepAliveEnabled(t.getKeepAliveEnabled().get());
+        }
+        if (t.getTcpNoDelayEnabled().isPresent()) {
+            builder.tcpNoDelayEnabled(t.getTcpNoDelayEnabled().get());
+        }
+
+        return builder.build();
+    }
+
+    private Consumer<EventStream> wrapEventConsumer(Consumer<EventStream> eventConsumer)
+    {
+        return eventStream ->
+        {
+            if (isShutdownTag(eventStream.getTag())) {
+                logger.info("Receive shutdown tag: {}", eventStream.getTag());
+                shouldShutdown.set(true);
+            }
+            else if (isMessageTag(eventStream.getTag())) {
+                eventConsumer.accept(eventStream);
+            }
+            else {
+                throw new RuntimeException(String.format("Unknown Tag received: %s", eventStream.getTag().getName()));
+            }
+        };
+    }
+
+    private boolean isShutdownTag(Tag tag)
+    {
+        return tag.getName().contentEquals(task.getShutdownTag());
+    }
+
+    private boolean isMessageTag(Tag tag)
+    {
+        return tag.getName().contentEquals(task.getMessageTag());
+    }
+
+    public void runUntilShouldShutdown()
+    {
+        server.start();
+        ElapsedTime.debugUntil(shouldShutdown::get, logger, "in_forward server", 1000L);
+        shutdown();
+    }
+
+    private void shutdown()
+    {
+        ElapsedTime.debug(logger, "shutdown in_forward callback", callbackThread::shutdown);
+        CompletableFuture<Void> shutdown = server.shutdown();
+        ElapsedTime.debugUntil(() -> shutdown.isCancelled() || shutdown.isCompletedExceptionally() || shutdown.isDone(),
+                logger, "shutdown in_forward server", 1000L);
+    }
+}

data/src/main/java/org/embulk/filter/copy/forward/InForwardVisitor.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.PageBuilder;
+
+public class InForwardVisitor
+    implements ColumnVisitor
+{
+    private final InForwardEventReader reader;
+    private final PageBuilder builder;
+
+    public InForwardVisitor(InForwardEventReader reader, PageBuilder 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/forward/OutForwardEventBuilder.java

@@ -0,0 +1,135 @@
+package org.embulk.filter.copy.forward;
+
+import com.google.common.collect.Maps;
+import org.embulk.spi.Column;
+import org.embulk.spi.Schema;
+import org.embulk.spi.time.Timestamp;
+import org.embulk.spi.time.TimestampFormatter;
+import org.msgpack.value.Value;
+
+import java.util.Map;
+
+public class OutForwardEventBuilder
+{
+    private final Schema schema;
+    private final TimestampFormatter timestampFormatter;
+
+    private Map<String, Object> message;
+
+    public OutForwardEventBuilder(
+            Schema schema,
+            TimestampFormatter timestampFormatter)
+    {
+
+        this.schema = schema;
+        this.timestampFormatter = timestampFormatter;
+
+        setNewMessage();
+    }
+
+    private void setNewMessage()
+    {
+        this.message = Maps.newHashMap();
+    }
+
+    public void emitMessage(OutForwardService outForward)
+    {
+        outForward.emit(message);
+        setNewMessage();
+    }
+
+    public Schema getSchema()
+    {
+        return schema;
+    }
+
+    private Column getColumn(int columnIndex)
+    {
+        return getSchema().getColumn(columnIndex);
+    }
+
+    private void setValue(String columnName, Object value)
+    {
+        message.put(columnName, value);
+    }
+
+    private void setValue(Column column, Object value)
+    {
+        setValue(column.getName(), value);
+    }
+
+    public void setValue(int columnIndex, Object value)
+    {
+        setValue(getColumn(columnIndex), value);
+    }
+
+    public void setNull(Column column)
+    {
+        setValue(column, null);
+    }
+
+    private void setNull(int columnIndex)
+    {
+        setNull(getColumn(columnIndex));
+    }
+
+    public void setBoolean(Column column, boolean value)
+    {
+        setValue(column, value);
+    }
+
+    public void setBoolean(int columnIndex, boolean value)
+    {
+        setBoolean(getColumn(columnIndex), value);
+    }
+
+    public void setLong(Column column, long value)
+    {
+        setValue(column, value);
+    }
+
+    public void setLong(int columnIndex, long value)
+    {
+        setLong(getColumn(columnIndex), value);
+    }
+
+    public void setDouble(Column column, double value)
+    {
+        setValue(column, value);
+    }
+
+    public void setDouble(int columnIndex, double value)
+    {
+        setDouble(getColumn(columnIndex), value);
+    }
+
+    public void setString(Column column, String value)
+    {
+        setValue(column, value);
+    }
+
+    public void setString(int columnIndex, String value)
+    {
+        setString(getColumn(columnIndex), value);
+    }
+
+    public void setTimestamp(Column column, Timestamp value)
+    {
+        setValue(column, timestampFormatter.format(value));
+    }
+
+    public void setTimestamp(int columnIndex, Timestamp value)
+    {
+        setTimestamp(getColumn(columnIndex), value);
+    }
+
+    public void setJson(Column column, Value value)
+    {
+        setValue(column, value.toJson());
+    }
+
+    public void setJson(int columnIndex, Value value)
+    {
+        setJson(getColumn(columnIndex), value);
+    }
+}