001// Licensed under the MIT license. See LICENSE file in the project root for full license information.
002
003package de.bytefish.pgbulkinsert.bulkprocessor;
004
005import de.bytefish.pgbulkinsert.bulkprocessor.handler.IBulkWriteHandler;
006
007import java.time.Duration;
008import java.util.ArrayList;
009import java.util.List;
010import java.util.Optional;
011import java.util.concurrent.Executors;
012import java.util.concurrent.ScheduledFuture;
013import java.util.concurrent.ScheduledThreadPoolExecutor;
014import java.util.concurrent.TimeUnit;
015
016public class BulkProcessor<TEntity> implements AutoCloseable {
017
018   private final ScheduledThreadPoolExecutor scheduler;
019
020    private final ScheduledFuture<?> scheduledFuture;
021
022    private volatile boolean closed = false;
023
024    private final IBulkWriteHandler<TEntity> handler;
025
026    private final int bulkSize;
027
028    private List<TEntity> batchedEntities;
029
030    public BulkProcessor(IBulkWriteHandler<TEntity> handler, int bulkSize) {
031        this(handler, bulkSize, null);
032    }
033
034    public BulkProcessor(IBulkWriteHandler<TEntity> handler, int bulkSize, Duration flushInterval) {
035
036        this.handler = handler;
037        this.bulkSize = bulkSize;
038
039        // Start with an empty List of batched entities:
040        this.batchedEntities = new ArrayList<>();
041
042        if(flushInterval != null) {
043            // Create a Scheduler for the time-based Flush Interval:
044            this.scheduler = (ScheduledThreadPoolExecutor) Executors.newScheduledThreadPool(1);
045            this.scheduler.setExecuteExistingDelayedTasksAfterShutdownPolicy(false);
046            this.scheduler.setContinueExistingPeriodicTasksAfterShutdownPolicy(false);
047            this.scheduledFuture = this.scheduler.scheduleWithFixedDelay(new Flush(), flushInterval.toMillis(), flushInterval.toMillis(), TimeUnit.MILLISECONDS);
048        } else {
049            this.scheduler = null;
050            this.scheduledFuture = null;
051        }
052    }
053
054    public synchronized BulkProcessor<TEntity> add(TEntity entity) {
055        batchedEntities.add(entity);
056        executeIfNeccessary();
057        return this;
058    }
059
060    @Override
061    public void close() throws Exception {
062        // If the Processor has already been closed, do not proceed:
063        if (closed) {
064            return;
065        }
066        closed = true;
067
068        // Quit the Scheduled FlushInterval Future:
069        Optional.ofNullable(this.scheduledFuture).ifPresent(future -> future.cancel(false));
070        Optional.ofNullable(this.scheduler).ifPresent(ScheduledThreadPoolExecutor::shutdown);
071
072        // Are there any entities left to write?
073        if (batchedEntities.size() > 0) {
074            execute();
075        }
076    }
077
078    private void executeIfNeccessary() {
079        if(batchedEntities.size() >= bulkSize) {
080            execute();
081        }
082    }
083
084    // (currently) needs to be executed under a lock
085    private void execute() {
086        // Assign to a new List:
087        final List<TEntity> entities = batchedEntities;
088        // We can restart batching entities:
089        batchedEntities = new ArrayList<>();
090        // Write the previously batched entities to PostgreSQL:
091        write(entities);
092    }
093
094    private void write(List<TEntity> entities) {
095        try {
096            handler.write(entities);
097        } catch(Exception e) {
098            throw new RuntimeException(e);
099        }
100    }
101
102    class Flush implements Runnable {
103
104        @Override
105        public void run() {
106            synchronized (BulkProcessor.this) {
107                if (closed) {
108                    return;
109                }
110                if (batchedEntities.size() == 0) {
111                    return;
112                }
113                execute();
114            }
115
116        }
117    }
118}