DAT-226: Use Netty's FastThreadLocal (#196)

This commit also gets rid of Stormpot BlazePool and replaces it
with the groups() operator, which showed noticeable thorughput
improvement.

Author: adutra

changelog/README.md

@@ -8,6 +8,7 @@
 - [improvement] DAT-108: Upgrade DSE driver to 1.6.3.
 - [improvement] DAT-227: Warn that continuous paging is not available when unloading with CL > ONE.
 - [new feature] DAT-224: Add support for numeric overflow and rounding.
+- [improvement] DAT-226: Use Netty's FastThreadLocal.
 
 
 ### 1.0.0-rc1

connectors/csv/pom.xml

@@ -44,8 +44,8 @@
     </dependency>
 
     <dependency>
-      <groupId>com.github.chrisvest</groupId>
-      <artifactId>stormpot</artifactId>
+      <groupId>io.netty</groupId>
+      <artifactId>netty-common</artifactId>
     </dependency>
 
     <dependency>

connectors/csv/src/main/java/com/datastax/dsbulk/connectors/csv/CSVConnector.java

@@ -29,6 +29,7 @@
 import com.univocity.parsers.csv.CsvParserSettings;
 import com.univocity.parsers.csv.CsvWriter;
 import com.univocity.parsers.csv.CsvWriterSettings;
+import io.netty.util.concurrent.DefaultThreadFactory;
 import java.io.IOException;
 import java.io.Reader;
 import java.io.UncheckedIOException;
@@ -44,7 +45,9 @@
 import java.nio.file.PathMatcher;
 import java.nio.file.Paths;
 import java.util.Arrays;
-import java.util.concurrent.TimeUnit;
+import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Function;
 import java.util.function.Supplier;
@@ -58,13 +61,6 @@
 import reactor.core.publisher.Signal;
 import reactor.core.scheduler.Scheduler;
 import reactor.core.scheduler.Schedulers;
-import reactor.util.concurrent.Queues;
-import stormpot.Allocator;
-import stormpot.BlazePool;
-import stormpot.Config;
-import stormpot.Poolable;
-import stormpot.Slot;
-import stormpot.Timeout;
 
 /**
  * A connector for CSV files.
@@ -115,7 +111,7 @@ public class CSVConnector implements Connector {
   private CsvWriterSettings writerSettings;
   private AtomicInteger counter;
   private Scheduler scheduler;
-  private BlazePool<PoolableCSVWriter> pool;
+  private List<CSVWriter> writers;
 
   @Override
   public void configure(LoaderConfig settings, boolean read) {
@@ -176,12 +172,12 @@ public void init() throws URISyntaxException, IOException {
   }
 
   @Override
-  public void close() throws InterruptedException {
+  public void close() {
     if (scheduler != null) {
       scheduler.dispose();
     }
-    if (pool != null) {
-      pool.shutdown().await(new Timeout(1, TimeUnit.MINUTES));
+    if (writers != null) {
+      writers.forEach(CSVWriter::close);
     }
   }
 
@@ -219,38 +215,21 @@ public Supplier<? extends Publisher<Publisher<Record>>> readByResource() {
     assert !read;
     if (root != null && maxConcurrentFiles > 1) {
       return upstream -> {
-        scheduler = Schedulers.newParallel("csv-connector", maxConcurrentFiles);
-        Config<PoolableCSVWriter> config =
-            new Config<PoolableCSVWriter>()
-                .setSize(maxConcurrentFiles)
-                .setAllocator(new CSVWriterAllocator());
-        pool = new BlazePool<>(config);
-        Timeout timeout = new Timeout(Long.MAX_VALUE, TimeUnit.SECONDS);
+        ThreadFactory threadFactory = new DefaultThreadFactory("csv-connector");
+        scheduler = Schedulers.newParallel(maxConcurrentFiles, threadFactory);
+        writers = new CopyOnWriteArrayList<>();
+        for (int i = 0; i < maxConcurrentFiles; i++) {
+          writers.add(new CSVWriter());
+        }
         return Flux.from(upstream)
-            .window(Queues.SMALL_BUFFER_SIZE)
             .parallel(maxConcurrentFiles)
             .runOn(scheduler)
-            .flatMap(
-                records -> {
-                  try {
-                    PoolableCSVWriter writer = pool.claim(timeout);
-                    if (writer != null) {
-                      try {
-                        return records.transform(writeRecords(writer));
-                      } finally {
-                        writer.release();
-                      }
-                    }
-                  } catch (InterruptedException e) {
-                    Thread.currentThread().interrupt();
-                  }
-                  return Flux.empty();
-                })
-            .sequential();
+            .groups()
+            .flatMap(records -> records.transform(writeRecords(writers.get(records.key()))));
       };
     } else {
       return upstream -> {
-        PoolableCSVWriter writer = new PoolableCSVWriter(null);
+        CSVWriter writer = new CSVWriter();
         return Flux.from(upstream).transform(writeRecords(writer)).doOnTerminate(writer::close);
       };
     }
@@ -396,7 +375,7 @@ private Flux<URL> scanRootDirectory() {
             });
   }
 
-  private Function<Flux<Record>, Flux<Record>> writeRecords(PoolableCSVWriter writer) {
+  private Function<Flux<Record>, Flux<Record>> writeRecords(CSVWriter writer) {
     return upstream ->
         upstream
             .materialize()
@@ -414,34 +393,11 @@ private Function<Flux<Record>, Flux<Record>> writeRecords(PoolableCSVWriter writ
             .dematerialize();
   }
 
-  private class CSVWriterAllocator implements Allocator<PoolableCSVWriter> {
-    @Override
-    public PoolableCSVWriter allocate(Slot slot) {
-      return new PoolableCSVWriter(slot);
-    }
-
-    @Override
-    public void deallocate(PoolableCSVWriter writer) {
-      writer.close();
-    }
-  }
-
-  private class PoolableCSVWriter implements Poolable {
-
-    private final Slot slot;
+  private class CSVWriter {
 
     private URL url;
     private CsvWriter writer;
 
-    private PoolableCSVWriter(Slot slot) {
-      this.slot = slot;
-    }
-
-    @Override
-    public void release() {
-      slot.release(this);
-    }
-
     private void write(Record record) {
       try {
         if (writer == null) {

connectors/json/pom.xml

@@ -49,8 +49,8 @@
     </dependency>
 
     <dependency>
-      <groupId>com.github.chrisvest</groupId>
-      <artifactId>stormpot</artifactId>
+      <groupId>io.netty</groupId>
+      <artifactId>netty-common</artifactId>
     </dependency>
 
     <dependency>

connectors/json/src/main/java/com/datastax/dsbulk/connectors/json/JsonConnector.java

@@ -26,6 +26,7 @@
 import com.fasterxml.jackson.core.JsonParseException;
 import com.fasterxml.jackson.core.JsonParser;
 import com.fasterxml.jackson.core.JsonToken;
+import com.fasterxml.jackson.core.TreeNode;
 import com.fasterxml.jackson.core.io.SerializedString;
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.core.util.DefaultPrettyPrinter;
@@ -38,6 +39,7 @@
 import com.fasterxml.jackson.databind.SerializationFeature;
 import com.google.common.base.Suppliers;
 import com.typesafe.config.ConfigException;
+import io.netty.util.concurrent.DefaultThreadFactory;
 import java.io.BufferedReader;
 import java.io.IOException;
 import java.io.UncheckedIOException;
@@ -53,8 +55,10 @@
 import java.nio.file.PathMatcher;
 import java.nio.file.Paths;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
-import java.util.concurrent.TimeUnit;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Function;
 import java.util.function.Supplier;
@@ -67,13 +71,6 @@
 import reactor.core.publisher.Signal;
 import reactor.core.scheduler.Scheduler;
 import reactor.core.scheduler.Schedulers;
-import reactor.util.concurrent.Queues;
-import stormpot.Allocator;
-import stormpot.BlazePool;
-import stormpot.Config;
-import stormpot.Poolable;
-import stormpot.Slot;
-import stormpot.Timeout;
 
 /**
  * A connector for Json files.
@@ -133,10 +130,9 @@ enum DocumentMode {
   private Map<MapperFeature, Boolean> mapperFeatures;
   private Map<SerializationFeature, Boolean> serializationFeatures;
   private Map<DeserializationFeature, Boolean> deserializationFeatures;
-
   private boolean prettyPrint;
   private Scheduler scheduler;
-  private BlazePool<PoolableJsonWriter> pool;
+  private List<JsonWriter> writers;
 
   @Override
   public RecordMetadata getRecordMetadata() {
@@ -211,12 +207,12 @@ public void init() throws URISyntaxException, IOException {
   }
 
   @Override
-  public void close() throws InterruptedException {
+  public void close() {
     if (scheduler != null) {
       scheduler.dispose();
     }
-    if (pool != null) {
-      pool.shutdown().await(new Timeout(1, TimeUnit.MINUTES));
+    if (writers != null) {
+      writers.forEach(JsonWriter::close);
     }
   }
 
@@ -254,38 +250,21 @@ public Supplier<? extends Publisher<Publisher<Record>>> readByResource() {
     assert !read;
     if (root != null && maxConcurrentFiles > 1) {
       return upstream -> {
-        scheduler = Schedulers.newParallel("json-connector", maxConcurrentFiles);
-        Config<PoolableJsonWriter> config =
-            new Config<PoolableJsonWriter>()
-                .setSize(maxConcurrentFiles)
-                .setAllocator(new JsonWriterAllocator());
-        pool = new BlazePool<>(config);
-        Timeout timeout = new Timeout(Long.MAX_VALUE, TimeUnit.SECONDS);
+        ThreadFactory threadFactory = new DefaultThreadFactory("json-connector");
+        scheduler = Schedulers.newParallel(maxConcurrentFiles, threadFactory);
+        writers = new CopyOnWriteArrayList<>();
+        for (int i = 0; i < maxConcurrentFiles; i++) {
+          writers.add(new JsonWriter());
+        }
         return Flux.from(upstream)
-            .window(Queues.SMALL_BUFFER_SIZE)
             .parallel(maxConcurrentFiles)
             .runOn(scheduler)
-            .flatMap(
-                records -> {
-                  try {
-                    PoolableJsonWriter writer = pool.claim(timeout);
-                    if (writer != null) {
-                      try {
-                        return records.transform(writeRecords(writer));
-                      } finally {
-                        writer.release();
-                      }
-                    }
-                  } catch (InterruptedException e) {
-                    Thread.currentThread().interrupt();
-                  }
-                  return Flux.empty();
-                })
-            .sequential();
+            .groups()
+            .flatMap(records -> records.transform(writeRecords(writers.get(records.key()))));
       };
     } else {
       return upstream -> {
-        PoolableJsonWriter writer = new PoolableJsonWriter(null);
+        JsonWriter writer = new JsonWriter();
         return Flux.from(upstream).transform(writeRecords(writer)).doOnTerminate(writer::close);
       };
     }
@@ -415,7 +394,7 @@ private Flux<URL> scanRootDirectory() {
             });
   }
 
-  private Function<Flux<Record>, Flux<Record>> writeRecords(PoolableJsonWriter writer) {
+  private Function<Flux<Record>, Flux<Record>> writeRecords(JsonWriter writer) {
     return upstream ->
         upstream
             .materialize()
@@ -433,35 +412,12 @@ private Function<Flux<Record>, Flux<Record>> writeRecords(PoolableJsonWriter wri
             .dematerialize();
   }
 
-  private class JsonWriterAllocator implements Allocator<PoolableJsonWriter> {
-    @Override
-    public PoolableJsonWriter allocate(Slot slot) {
-      return new PoolableJsonWriter(slot);
-    }
-
-    @Override
-    public void deallocate(PoolableJsonWriter writer) {
-      writer.close();
-    }
-  }
-
-  private class PoolableJsonWriter implements Poolable {
-
-    private final Slot slot;
+  private class JsonWriter {
 
     private URL url;
     private JsonGenerator writer;
     private long currentLine;
 
-    private PoolableJsonWriter(Slot slot) {
-      this.slot = slot;
-    }
-
-    @Override
-    public void release() {
-      slot.release(this);
-    }
-
     private void write(Record record) {
       try {
         if (writer == null) {
@@ -477,7 +433,7 @@ private void write(Record record) {
         writer.writeStartObject();
         for (String field : record.fields()) {
           writer.writeFieldName(field);
-          writer.writeObject(record.getFieldValue(field));
+          writer.writeTree((TreeNode) record.getFieldValue(field));
         }
         writer.writeEndObject();
         currentLine++;
@@ -495,7 +451,8 @@ private void open() throws IOException {
       url = getOrCreateDestinationURL();
       writer = createJsonWriter(url);
       if (mode == DocumentMode.SINGLE_DOCUMENT) {
-        // do not use writer.writeStartArray(): we need to fool the parser into thinking it's on multi doc mode,
+        // do not use writer.writeStartArray(): we need to fool the parser into thinking it's on
+        // multi doc mode,
         // to get a better-looking result
         writer.writeRaw('[');
         writer.writeRaw(System.lineSeparator());

connectors/json/src/test/java/com/datastax/dsbulk/connectors/json/JsonConnectorTest.java

@@ -48,7 +48,6 @@
 import org.slf4j.LoggerFactory;
 import reactor.core.publisher.Flux;
 
-/** */
 @SuppressWarnings("Duplicates")
 class JsonConnectorTest {
 
@@ -198,7 +197,15 @@ void should_write_to_stdout_with_special_encoding() throws Exception {
       connector.configure(settings, false);
       connector.init();
       assertThat(connector.isWriteToStandardOutput()).isTrue();
-      Flux.<Record>just(new DefaultRecord(null, null, -1, null, "fóô", "bàr", "qïx"))
+      Flux.<Record>just(
+          new DefaultRecord(
+              null,
+              null,
+              -1,
+              null,
+              factory.textNode("fóô"),
+              factory.textNode("bàr"),
+              factory.textNode("qïx")))
           .transform(connector.write())
           .blockLast();
       assertThat(new String(baos.toByteArray(), "ISO-8859-1"))