embulk-filter-speedometer 0.2.1

data/lib/embulk/filter/speedometer.rb

@@ -0,0 +1,3 @@
+Embulk::JavaPlugin.register_filter(
+  "speedometer", "org.embulk.filter.SpeedometerFilterPlugin",
+  File.expand_path('../../../../classpath', __FILE__))

data/src/main/java/org/embulk/filter/SpeedometerFilterPlugin.java

@@ -0,0 +1,289 @@
+package org.embulk.filter;
+
+import javax.validation.constraints.Min;
+
+import org.embulk.config.Config;
+import org.embulk.config.ConfigDefault;
+import org.embulk.config.ConfigInject;
+import org.embulk.config.ConfigSource;
+import org.embulk.config.Task;
+import org.embulk.config.TaskSource;
+import org.embulk.spi.BufferAllocator;
+import org.embulk.spi.Column;
+import org.embulk.spi.ColumnVisitor;
+import org.embulk.spi.FilterPlugin;
+import org.embulk.spi.Page;
+import org.embulk.spi.PageBuilder;
+import org.embulk.spi.PageOutput;
+import org.embulk.spi.PageReader;
+import org.embulk.spi.Schema;
+import org.embulk.spi.time.Timestamp;
+import org.embulk.spi.time.TimestampFormatter;
+import org.embulk.spi.type.TimestampType;
+
+import com.google.common.collect.ImmutableMap;
+
+public class SpeedometerFilterPlugin
+        implements FilterPlugin
+{
+    private static final int TRUE_LENGTH = Boolean.toString(true).length();
+    private static final int FALSE_LENGTH = Boolean.toString(false).length();
+
+    public interface PluginTask
+            extends Task,  TimestampFormatter.FormatterTask
+    {
+        @Config("speed_limit")
+        @ConfigDefault("0")
+        @Min(0)
+        public long getSpeedLimit();
+
+        @Config("max_sleep_millisec")
+        @ConfigDefault("1000")
+        @Min(0)
+        public int getMaxSleepMillisec();
+
+        @Config("delimiter")
+        @ConfigDefault("\",\"")
+        public String getDelimiter();
+
+        @Config("record_padding_size")
+        @ConfigDefault("1")
+        public int getRecordPaddingSize();
+
+        @Config("log_interval_seconds")
+        @ConfigDefault("10")
+        @Min(0)
+        public int getLogIntervalSeconds();
+
+        @ConfigInject
+        public BufferAllocator getBufferAllocator();
+    }
+
+    @Override
+    public void transaction(ConfigSource config, Schema inputSchema,
+            FilterPlugin.Control control)
+    {
+        PluginTask task = config.loadConfig(PluginTask.class);
+        Schema outputSchema = inputSchema;
+        control.run(task.dump(), outputSchema);
+    }
+
+    @Override
+    public PageOutput open(TaskSource taskSource, Schema inputSchema,
+            Schema outputSchema, PageOutput output)
+    {
+        PluginTask task = taskSource.loadTask(PluginTask.class);
+
+        return new SpeedControlPageOutput(task, inputSchema, output);
+    }
+
+    static class SpeedControlPageOutput implements PageOutput {
+        private final SpeedometerSpeedController controller;
+        private final Schema schema;
+        private final ImmutableMap<Column, TimestampFormatter> timestampMap;
+        private final PageOutput pageOutput;
+        private final PageOutput addOnlyPageOutput;
+        private final PageReader pageReader;
+        private final BufferAllocator allocator;
+        private final int delimiterLength;
+        private final int recordPaddingSize;
+
+        SpeedControlPageOutput(PluginTask task, Schema schema, PageOutput pageOutput) {
+            this.controller = new SpeedometerSpeedController(task, SpeedometerSpeedAggregator.getInstance());
+            this.schema = schema;
+            this.pageOutput = pageOutput;
+            this.addOnlyPageOutput = new AddOnlyPageOutput(pageOutput);
+            this.allocator = task.getBufferAllocator();
+            this.delimiterLength = task.getDelimiter().length();
+            this.recordPaddingSize = task.getRecordPaddingSize();
+            pageReader = new PageReader(schema);
+            timestampMap = buildTimestampFormatterMap(task, schema);
+        }
+
+        @Override
+        public void add(Page page) {
+            try (final PageBuilder pageBuilder = new PageBuilder(allocator, schema, addOnlyPageOutput)) {
+                ColumnVisitorImpl visitor = new ColumnVisitorImpl(pageBuilder);
+                pageReader.setPage(page);
+                while (pageReader.nextRecord()) {
+                    visitor.speedMonitorStartRecord();
+                    schema.visitColumns(visitor);
+                    visitor.speedMonitorEndRecord();
+                    pageBuilder.addRecord();
+                }
+                pageBuilder.finish();
+            }
+        }
+
+        @Override
+        public void finish() {
+            pageOutput.finish();
+        }
+
+        @Override
+        public void close() {
+            controller.stop();
+            pageReader.close();
+            pageOutput.close();
+        }
+
+        private ImmutableMap<Column, TimestampFormatter> buildTimestampFormatterMap(final PluginTask task, Schema schema) {
+            final ImmutableMap.Builder<Column, TimestampFormatter> builder = new ImmutableMap.Builder<>();
+
+            schema.visitColumns(new ColumnVisitor() {
+                @Override
+                public void booleanColumn(Column column) { }
+
+                @Override
+                public void longColumn(Column column) { }
+
+                @Override
+                public void doubleColumn(Column column) { }
+
+                @Override
+                public void stringColumn(Column column) { }
+
+                @Override
+                public void timestampColumn(Column column) {
+                    if (column.getType() instanceof TimestampType) {
+                        TimestampType tt = (TimestampType) column.getType();
+                        builder.put(column, new TimestampFormatter(tt.getFormat(), task));
+                    } else {
+                        throw new RuntimeException("Timestamp should be TimestampType.");
+                    }
+                }
+            });
+
+            return builder.build();
+        }
+
+        // Ignore finish and close to avoid closing upper output stream.
+        static class AddOnlyPageOutput implements PageOutput {
+            protected final PageOutput output;
+
+            public AddOnlyPageOutput(PageOutput outptut) {
+                this.output = outptut;
+            }
+
+            @Override
+            public void add(Page page) {
+                output.add(page);
+            }
+
+            @Override
+            public void finish() { }
+
+            @Override
+            public void close() { }
+        }
+
+        class ColumnVisitorImpl implements ColumnVisitor {
+            private final PageBuilder pageBuilder;
+            private long startRecordTime;
+
+            ColumnVisitorImpl(PageBuilder pageBuilder) {
+                this.pageBuilder = pageBuilder;
+            }
+
+            @Override
+            public void booleanColumn(Column column) {
+                if (pageReader.isNull(column)) {
+                    speedMonitor(column);
+                    pageBuilder.setNull(column);
+                } else {
+                    pageBuilder.setBoolean(column, speedMonitor(column, pageReader.getBoolean(column)));
+                }
+            }
+
+            @Override
+            public void longColumn(Column column) {
+                if (pageReader.isNull(column)) {
+                    speedMonitor(column);
+                    pageBuilder.setNull(column);
+                } else {
+                    pageBuilder.setLong(column, speedMonitor(column, pageReader.getLong(column)));
+                }
+            }
+
+            @Override
+            public void doubleColumn(Column column) {
+                if (pageReader.isNull(column)) {
+                    speedMonitor(column);
+                    pageBuilder.setNull(column);
+                } else {
+                    pageBuilder.setDouble(column, speedMonitor(column, pageReader.getDouble(column)));
+                }
+            }
+
+            @Override
+            public void stringColumn(Column column) {
+                if (pageReader.isNull(column)) {
+                    speedMonitor(column);
+                    pageBuilder.setNull(column);
+                } else {
+                    pageBuilder.setString(column, speedMonitor(column, pageReader.getString(column)));
+                }
+            }
+
+            @Override
+            public void timestampColumn(Column column) {
+                if (pageReader.isNull(column)) {
+                    speedMonitor(column);
+                    pageBuilder.setNull(column);
+                } else {
+                    pageBuilder.setTimestamp(column, speedMonitor(column, pageReader.getTimestamp(column)));
+                }
+            }
+
+            private void speedMonitorStartRecord() {
+                startRecordTime = System.currentTimeMillis();
+            }
+
+            private void speedMonitorEndRecord() {
+                controller.checkSpeedLimit(startRecordTime, recordPaddingSize);
+            }
+
+            // For null column
+            private void speedMonitor(Column column) {
+                speedMonitorForDelimiter(column);
+            }
+
+            private boolean speedMonitor(Column column, boolean b) {
+                speedMonitorForDelimiter(column);
+                controller.checkSpeedLimit(startRecordTime, b ? TRUE_LENGTH : FALSE_LENGTH);
+                return b;
+            }
+
+            private long speedMonitor(Column column, long l) {
+                speedMonitorForDelimiter(column);
+                controller.checkSpeedLimit(startRecordTime, String.valueOf(l).length());
+                return l;
+            }
+
+            private double speedMonitor(Column column, double d) {
+                speedMonitorForDelimiter(column);
+                controller.checkSpeedLimit(startRecordTime, String.valueOf(d).length());
+                return d;
+            }
+
+            private String speedMonitor(Column column, String s) {
+                speedMonitorForDelimiter(column);
+                controller.checkSpeedLimit(startRecordTime, s.length());
+                return s;
+            }
+
+            private Timestamp speedMonitor(Column column, Timestamp t) {
+                speedMonitorForDelimiter(column);
+                TimestampFormatter formatter = timestampMap.get(column);
+                controller.checkSpeedLimit(startRecordTime, formatter.format(t).length());
+                return t;
+            }
+
+            private void speedMonitorForDelimiter(Column column) {
+                if (column.getIndex() > 0) {
+                    controller.checkSpeedLimit(startRecordTime, delimiterLength);
+                }
+            }
+        }
+    }
+}

data/src/main/java/org/embulk/filter/SpeedometerSpeedAggregator.java

@@ -0,0 +1,141 @@
+package org.embulk.filter;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+
+import org.embulk.spi.Exec;
+import org.slf4j.Logger;
+
+class SpeedometerSpeedAggregator {
+    private static class SpeedometerSpeedAggregatorHolder {
+        private static final SpeedometerSpeedAggregator INSTANCE = new SpeedometerSpeedAggregator();
+    }
+
+    private final long INITAL_START_TIME = 0;
+
+    private final AtomicInteger activeControllerCount = new AtomicInteger(0);
+    private final AtomicLong globalStartTime = new AtomicLong(INITAL_START_TIME);
+    private final AtomicLong globalTotalBytes = new AtomicLong(0);
+    private final AtomicLong previousLogReportTimeMillisec = new AtomicLong(INITAL_START_TIME);
+
+    // TODO: We can use google's library.
+    private final List<SpeedometerSpeedController> controllerList = new ArrayList<>();
+
+    public static SpeedometerSpeedAggregator getInstance() {
+        return SpeedometerSpeedAggregatorHolder.INSTANCE;
+    }
+
+    public SpeedometerSpeedAggregator() {
+        showLogMessage(activeControllerCount.get(), 0, 0, 0);
+    }
+
+    public void startController(SpeedometerSpeedController controller, long nowTime) {
+        globalStartTime.compareAndSet(INITAL_START_TIME, nowTime);
+        activeControllerCount.incrementAndGet();
+        synchronized (controllerList) {
+            controllerList.add(controller);
+        }
+    }
+
+    public void stopController(SpeedometerSpeedController controller) {
+        synchronized (controllerList) {
+            controllerList.remove(controller);
+        }
+        long runningCount = activeControllerCount.decrementAndGet();
+        globalTotalBytes.addAndGet(controller.getTotalBytes());
+
+        // NOTE: Sometimes, there is no running thread nevertheless there are remaining tasks.
+        // So, this message may be output while running tasks.
+        if (runningCount == 0) {
+            showOverallMessage();
+        }
+    }
+
+    public long getSpeedLimitForController(SpeedometerSpeedController controller) {
+        return controller.getSpeedLimit() / activeControllerCount.get();
+    }
+
+    public void checkProgress(long nowTime, int logIntervalMillisec) {
+        if (logIntervalMillisec <= 0) {
+            return;
+        }
+
+        long previousTime = previousLogReportTimeMillisec.get();
+        if (previousTime == INITAL_START_TIME) {
+            previousLogReportTimeMillisec.compareAndSet(INITAL_START_TIME, nowTime);
+        } else {
+            long nowInterval = previousTime + logIntervalMillisec;
+            if (nowInterval < nowTime) {
+                if (previousLogReportTimeMillisec.compareAndSet(previousTime, nowTime)) {
+                    showProgressMessage(nowTime);
+                    renewPeriods();
+                }
+            }
+        }
+    }
+
+    Logger getLogger() {
+        return Exec.getLogger(SpeedometerFilterPlugin.class);
+    }
+
+    long getGlobalStartTime() {
+        return globalStartTime.get();
+    }
+
+    int getActiveControllerCount() {
+        return activeControllerCount.get();
+    }
+
+    long getGlobalTotalBytes() {
+        return globalTotalBytes.get();
+    }
+
+    List<SpeedometerSpeedController> getControllerList() {
+        List<SpeedometerSpeedController> copyList;
+        synchronized (controllerList) {
+            copyList = new ArrayList<SpeedometerSpeedController>(controllerList);
+        }
+        return copyList;
+    }
+
+    private void renewPeriods() {
+        for (SpeedometerSpeedController controller : getControllerList()) {
+            controller.renewPeriod();
+        }
+    }
+
+    private void showProgressMessage(long nowTime) {
+        long currentTotalSize = globalTotalBytes.get();
+        long currentBytesPerSec = 0;
+        for (SpeedometerSpeedController controller : getControllerList()) {
+            currentTotalSize += controller.getTotalBytes();
+            currentBytesPerSec += controller.getPeriodBytesPerSec(nowTime);
+        }
+
+        long timeDelta = nowTime - globalStartTime.get();
+        timeDelta = timeDelta > 0 ? timeDelta : 1;
+
+        showLogMessage(activeControllerCount.get(), currentTotalSize, timeDelta, currentBytesPerSec);
+    }
+
+    private void showOverallMessage() {
+        long timeDelta = System.currentTimeMillis() - globalStartTime.get();
+        timeDelta = timeDelta > 0 ? timeDelta : 1;
+        long bytesPerSec = (globalTotalBytes.get() * 1000) / timeDelta;
+
+        showLogMessage(activeControllerCount.get(), globalTotalBytes.get(), timeDelta, bytesPerSec);
+    }
+
+    private void showLogMessage(int activeThreads, long totalBytes, long timeMilliSec, long bytesPerSec) {
+        Logger logger = getLogger();
+        if (logger != null) {
+            logger.info(String.format("{speedometer: {active: %d, total: %s, sec: %s, speed: %s/s}}",
+                    activeThreads,
+                    SpeedometerUtil.toNumberText(totalBytes),
+                    SpeedometerUtil.toTimeText(timeMilliSec),
+                    SpeedometerUtil.toNumberText(bytesPerSec)));
+        }
+    }
+}

data/src/main/java/org/embulk/filter/SpeedometerSpeedController.java

@@ -0,0 +1,107 @@
+package org.embulk.filter;
+
+import org.embulk.filter.SpeedometerFilterPlugin.PluginTask;
+
+
+class SpeedometerSpeedController {
+    private final SpeedometerSpeedAggregator aggregator;
+
+    private final long limitBytesPerSec;
+    private final int maxSleepMillisec;
+    private final int logIntervalMillisec;
+
+    private long startTime;
+    private volatile long periodStartTime;
+    private volatile long periodTotalBytes;
+    private volatile long threadTotalBytes;
+    private volatile boolean renewFlag = true;
+
+    SpeedometerSpeedController(PluginTask task, SpeedometerSpeedAggregator aggregator) {
+        this.limitBytesPerSec = task.getSpeedLimit();
+        this.maxSleepMillisec = task.getMaxSleepMillisec();
+        this.logIntervalMillisec = task.getLogIntervalSeconds() * 1000;
+        this.aggregator = aggregator;
+    }
+
+    public void stop() {
+        startNewPeriod(0);
+        aggregator.stopController(this);
+    }
+
+    public long getSpeedLimit() {
+        return limitBytesPerSec;
+    }
+
+    public int getMaxSleepMillisec() {
+        return maxSleepMillisec;
+    }
+
+    public int getLogIntervalMillisec() {
+        return logIntervalMillisec;
+    }
+
+    public long getTotalBytes() {
+        return threadTotalBytes + periodTotalBytes;
+    }
+
+    public long getPeriodBytesPerSec(long nowTime) {
+        long timeDeltaMillisec = nowTime - periodStartTime;
+        if (timeDeltaMillisec <= 0) {
+            timeDeltaMillisec = 1;
+        }
+        return (periodTotalBytes * 1000) / timeDeltaMillisec;
+    }
+
+    public void checkSpeedLimit(long nowTime, long newDataSize) {
+        if (startTime == 0) {
+            startTime = nowTime;
+            aggregator.startController(this, startTime);
+        }
+
+        if (renewFlag) {
+            renewFlag = false;
+            startNewPeriod(nowTime);
+        }
+
+        periodTotalBytes += newDataSize;
+        aggregator.checkProgress(nowTime, logIntervalMillisec);
+
+        if (limitBytesPerSec <= 0) {
+            return;
+        }
+
+        long speedLimitForThread = aggregator.getSpeedLimitForController(this);
+        long timeDeltaMillisec = nowTime > periodStartTime ? nowTime - periodStartTime : 1;
+        long bytesPerSec = (periodTotalBytes * 1000) / timeDeltaMillisec;
+        long overBytes = bytesPerSec - speedLimitForThread;
+
+        if (overBytes > 0) {
+            try {
+                long sleepTime = (periodTotalBytes * 1000) / speedLimitForThread - timeDeltaMillisec;
+                sleepTime = sleepTime > maxSleepMillisec ? maxSleepMillisec : sleepTime > 0 ? sleepTime : 0;
+                Thread.sleep(sleepTime);
+            } catch (InterruptedException e) {
+                // TODO: Do I need to throw an exception ?
+            }
+        }
+    }
+
+    void renewPeriod() {
+        renewFlag = true;
+    }
+
+    boolean isRenewPeriodSet() {
+        return renewFlag;
+    }
+
+    SpeedometerSpeedAggregator getAggregator() {
+        return aggregator;
+    }
+
+    private void startNewPeriod(long newPeriodTime) {
+        threadTotalBytes += periodTotalBytes;
+        periodTotalBytes = 0;
+        periodStartTime = newPeriodTime;
+    }
+}
+