diff --git a/planetiler-core/src/main/java/com/onthegomap/planetiler/Planetiler.java b/planetiler-core/src/main/java/com/onthegomap/planetiler/Planetiler.java index 56ee3dcb97..2426b1fedd 100644 --- a/planetiler-core/src/main/java/com/onthegomap/planetiler/Planetiler.java +++ b/planetiler-core/src/main/java/com/onthegomap/planetiler/Planetiler.java @@ -1,5 +1,7 @@ package com.onthegomap.planetiler; +import static com.onthegomap.planetiler.worker.Worker.joinFutures; + import com.onthegomap.planetiler.archive.TileArchiveConfig; import com.onthegomap.planetiler.archive.TileArchiveMetadata; import com.onthegomap.planetiler.archive.TileArchiveWriter; @@ -39,6 +41,7 @@ import com.onthegomap.planetiler.util.Wikidata; import com.onthegomap.planetiler.validator.JavaProfileValidator; import com.onthegomap.planetiler.worker.RunnableThatThrows; +import com.onthegomap.planetiler.worker.Worker.NamedThreadFactory; import java.io.IOException; import java.nio.file.FileSystem; import java.nio.file.Files; @@ -49,6 +52,7 @@ import java.util.Map; import java.util.Optional; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.function.Function; @@ -954,11 +958,16 @@ public void run() { featureGroup.loadStringEncoders(stringEncoderPath); featureGroup.initFromManifest(chunkManifestPath); } else { - for (Stage stage : stages) { - try { - stage.task.run(); - } catch (Exception e) { - throw new PlanetilerException("Error occurred during stage " + stage.id, e); + int requestedParallelism = config.sourceParallelism(); + if (requestedParallelism > 1 && stages.size() > 1) { + runStagesInParallel(stages, Math.min(requestedParallelism, stages.size())); + } else { + for (Stage stage : stages) { + try { + stage.task.run(); + } catch (Exception e) { + throw new PlanetilerException("Error occurred during stage " + stage.id, e); + } } } @@ -1135,6 +1144,27 @@ private void ensureInputFilesExist() { } } + private void runStagesInParallel(List stages, int parallelism) { + LOGGER.info("Running {} source stages with parallelism={}", stages.size(), parallelism); + try (var es = Executors.newFixedThreadPool(parallelism, new NamedThreadFactory("stage-runner"))) { + joinFutures(stages.stream() + .>map(stage -> CompletableFuture.runAsync(() -> { + try { + stage.task.run(); + } catch (Exception e) { + throw new PlanetilerException("Error occurred during stage " + stage.id, e); + } + }, es)) + .toList() + ).join(); + } catch (CompletionException ce) { + if (ce.getCause() instanceof PlanetilerException pe) { + throw pe; + } + throw ce; + } + } + private record Stage(String id, List details, RunnableThatThrows task) { Stage(String id, String description, RunnableThatThrows task) { diff --git a/planetiler-core/src/main/java/com/onthegomap/planetiler/config/PlanetilerConfig.java b/planetiler-core/src/main/java/com/onthegomap/planetiler/config/PlanetilerConfig.java index 3a826153c9..0df5bd0a8d 100644 --- a/planetiler-core/src/main/java/com/onthegomap/planetiler/config/PlanetilerConfig.java +++ b/planetiler-core/src/main/java/com/onthegomap/planetiler/config/PlanetilerConfig.java @@ -73,7 +73,8 @@ public record PlanetilerConfig( int featureSourceIdMultiplier, List extraNameTags, boolean reuseFeatureDb, - boolean parallelTempIO + boolean parallelTempIO, + int sourceParallelism ) { public static final int MIN_MINZOOM = 0; @@ -256,7 +257,11 @@ public static PlanetilerConfig from(Arguments arguments) { arguments.getBoolean("reuse_featuredb", "Reuse existing feature DB on disk, skipping source reading stages (for iterating on post-processing logic)", false), - parallelTempIO + parallelTempIO, + Math.max(1, arguments.getInteger("source_parallelism", + "number of input source stages to run in parallel (1 = sequential, current behavior). " + + "Useful for schemas with many small sources (e.g. shapefiles)", + 1)) ); } diff --git a/planetiler-core/src/main/java/com/onthegomap/planetiler/worker/Worker.java b/planetiler-core/src/main/java/com/onthegomap/planetiler/worker/Worker.java index aec95fe535..b5ca7bb5de 100644 --- a/planetiler-core/src/main/java/com/onthegomap/planetiler/worker/Worker.java +++ b/planetiler-core/src/main/java/com/onthegomap/planetiler/worker/Worker.java @@ -144,13 +144,13 @@ public void await() { } /** A thread factory that prepends {@code name-} to all thread names. */ - private static class NamedThreadFactory implements ThreadFactory { + public static class NamedThreadFactory implements ThreadFactory { private final ThreadGroup group; private final AtomicInteger threadNumber = new AtomicInteger(1); private final String namePrefix; - private NamedThreadFactory(String name) { + public NamedThreadFactory(String name) { group = Thread.currentThread().getThreadGroup(); namePrefix = name + "-"; } diff --git a/planetiler-core/src/test/java/com/onthegomap/planetiler/PlanetilerTests.java b/planetiler-core/src/test/java/com/onthegomap/planetiler/PlanetilerTests.java index 1f402eb4e3..fa321d077e 100644 --- a/planetiler-core/src/test/java/com/onthegomap/planetiler/PlanetilerTests.java +++ b/planetiler-core/src/test/java/com/onthegomap/planetiler/PlanetilerTests.java @@ -2880,12 +2880,17 @@ public void processFeature(SourceFeature source, FeatureCollector features) { } } - @Test - void testPlanetilerRunnerShapefile() throws Exception { + @ParameterizedTest + @ValueSource(strings = { + "", + "--source_parallelism=2", + "--source_parallelism=4" + }) + void testPlanetilerRunnerShapefile(String args) throws Exception { Path mbtiles = tempDir.resolve("output.mbtiles"); Path resourceDir = TestUtils.pathToResource(""); - Planetiler.create(Arguments.fromArgs("--tmpdir=" + tempDir.resolve("data"))) + Planetiler.create(Arguments.fromArgs((args + " --tmpdir=" + tempDir.resolve("data")).split("\\s+"))) .setProfile(new Profile.NullProfile() { @Override public void processFeature(SourceFeature source, FeatureCollector features) { diff --git a/planetiler-core/src/test/java/com/onthegomap/planetiler/config/PlanetilerConfigTest.java b/planetiler-core/src/test/java/com/onthegomap/planetiler/config/PlanetilerConfigTest.java new file mode 100644 index 0000000000..218ada628a --- /dev/null +++ b/planetiler-core/src/test/java/com/onthegomap/planetiler/config/PlanetilerConfigTest.java @@ -0,0 +1,24 @@ +package com.onthegomap.planetiler.config; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import org.junit.jupiter.api.Test; + +class PlanetilerConfigTest { + + @Test + void testSourceParallelismDefaultsToOne() { + assertEquals(1, PlanetilerConfig.defaults().sourceParallelism()); + } + + @Test + void testSourceParallelismParsesArgument() { + assertEquals(4, PlanetilerConfig.from(Arguments.fromArgs("--source_parallelism=4")).sourceParallelism()); + } + + @Test + void testSourceParallelismClampedToOne() { + assertEquals(1, PlanetilerConfig.from(Arguments.fromArgs("--source_parallelism=0")).sourceParallelism()); + assertEquals(1, PlanetilerConfig.from(Arguments.fromArgs("--source_parallelism=-5")).sourceParallelism()); + } +}