<commit_before>package io.korobi.mongotoelastic.mongo;

import com.mongodb.client.FindIterable;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoCursor;
import com.mongodb.client.MongoDatabase;
import io.korobi.mongotoelastic.exception.UserIsAtFaultException;
import io.korobi.mongotoelastic.logging.InjectLogger;
import io.korobi.mongotoelastic.opt.IOptions;
import io.korobi.mongotoelastic.processor.IDocumentProcessor;
import io.korobi.mongotoelastic.processor.RunnableProcessor;
import io.korobi.mongotoelastic.util.NumberUtil;
import org.apache.logging.log4j.Logger;
import org.bson.Document;

import javax.inject.Inject;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CountDownLatch;

public class MongoRetriever {

    @InjectLogger
    private Logger logger;
    private MongoDatabase database;
    private IOptions opts;
    private IDocumentProcessor processor;
    private int itemsPerThread;

    @Inject
    public MongoRetriever(MongoDatabase database, IOptions opts, IDocumentProcessor processor) {
        this.database = database;
        this.opts = opts;
        this.processor = processor;

        this.checkUserThreadParameters();
    }

    private void checkUserThreadParameters() {
        double itemsPerThread = this.opts.getBatchSize() / (double) this.opts.getThreadCap();
        if (NumberUtil.isValueFractional(itemsPerThread)) {
            throw new UserIsAtFaultException("Disallowing poor thread-related option choice.");
        }
        this.itemsPerThread = (int) itemsPerThread;
    }

    public void processData() {
        MongoCollection<Document> collection = this.database.getCollection("chats");
        logger.info("There are " + String.valueOf(collection.count()) + " chats!");

        List<Document> currentBatch;
        FindIterable<Document> iterable = collection.find();
        iterable.batchSize(this.opts.getBatchSize());
        MongoCursor<Document> cursor = iterable.iterator();
        while (!(currentBatch = this.buildBatch(cursor)).isEmpty()) {
            // great! We have a bunch of documents in RAM now :D
            CountDownLatch latch = new CountDownLatch(this.opts.getThreadCap());
            logger.info(String.format("Got a new batch of %d documents", currentBatch.size()));
            if (currentBatch.size() != this.opts.getBatchSize()) {
                logger.info("Last batch!");
            }

            for (int threadNumber = 1; threadNumber <= this.opts.getThreadCap(); threadNumber++) {
                int fromIndex = (threadNumber - 1) * this.itemsPerThread;
                int endIndex = fromIndex + this.itemsPerThread;
                int lastIndexInBatch = currentBatch.size() - 1;
                if (endIndex > lastIndexInBatch) {
                    endIndex = lastIndexInBatch; // this doesn't work :D
                }
                List<Document> forThread = currentBatch.subList(fromIndex, endIndex);
                Thread thread = new Thread(new RunnableProcessor(forThread, this.processor, latch));
                thread.start();
            }


            try {
                latch.await();
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
        cursor.close();
    }

    private List<Document> buildBatch(MongoCursor<Document> cursor) {
        List<Document> currentBatch = new ArrayList<>(this.opts.getBatchSize());

        while (cursor.hasNext() && currentBatch.size() < this.opts.getBatchSize()) {
            currentBatch.add(cursor.next()); // yield return cursor.next() :(
        }
        return currentBatch;
    }
}<commit_msg>Use a thread pool instead of spawning new threads<commit_after>package io.korobi.mongotoelastic.mongo;

import com.mongodb.client.FindIterable;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoCursor;
import com.mongodb.client.MongoDatabase;
import io.korobi.mongotoelastic.logging.InjectLogger;
import io.korobi.mongotoelastic.opt.IOptions;
import io.korobi.mongotoelastic.processor.IDocumentProcessor;
import org.apache.logging.log4j.Logger;
import org.bson.Document;

import javax.inject.Inject;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

public class MongoRetriever {

    @InjectLogger
    private Logger logger;
    private final MongoDatabase database;
    private final IOptions opts;
    private final IDocumentProcessor processor;
    private final int itemsPerThread;
    private final ExecutorService threadPool;

    @Inject
    public MongoRetriever(MongoDatabase database, IOptions opts, IDocumentProcessor processor) {
        this.database = database;
        this.opts = opts;
        this.processor = processor;
        this.itemsPerThread = opts.getBatchSize();
        this.threadPool = Executors.newFixedThreadPool(opts.getThreadCap());
    }

    public void processData() {
        MongoCollection<Document> collection = this.database.getCollection("chats");
        logger.info("There are " + String.valueOf(collection.count()) + " chats!");

        FindIterable<Document> iterable = collection.find();
        iterable.batchSize(this.opts.getBatchSize());
        MongoCursor<Document> cursor = iterable.iterator();

        List<Document> documents = new ArrayList<>(this.itemsPerThread);
        int current = 0;
        while(cursor.hasNext()) {
            if(current < this.itemsPerThread) {
                ++current;
                documents.add(cursor.next());
            } else {
                logger.info(String.format("Got a new batch of %d documents", documents.size()));
                processBatch(documents);
                documents = new ArrayList<>(this.itemsPerThread);
                current = 0;
            }
        }
        cursor.close();

        if(!documents.isEmpty()) {
            logger.info("Last batch!");
            processBatch(documents);
        }

        this.threadPool.shutdown();
        try {
            this.threadPool.awaitTermination(0, TimeUnit.MILLISECONDS);
        } catch(InterruptedException e) {
            e.printStackTrace();
        }
    }

    private void processBatch(List<Document> documents) {
        this.threadPool.submit(() -> this.processor.run(documents));
    }

}
{'fraction_non_alphanumeric': 0.10512339893783193, 'fraction_numerical': 0.0012496094970321774, 'mean_word_length': 2.854906682721252, 'pattern_counts': {'":': 0, '<': 22, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1115249', 'n_tokens_mistral': 1760, 'n_tokens_neox': 1705, 'n_words': 351}