MapComplete/scripts/generatePmTilesExtracts.ts

import Script from "./Script"
import { Tiles } from "../src/Models/TileRange"
import { OfflineBasemapManager } from "../src/service-worker/OfflineBasemapManager"

import { spawn } from "child_process"
import { existsSync } from "fs"
import { Utils } from "../src/Utils"

class GeneratePmTilesExtracts extends Script {
    private targetDir: string
    private skipped: number = 0
    constructor() {
        super(
            "Generates many pmtiles-archive from planet-latest.pmtiles. Must be started from the directory where planet-latest.pmtiles resides, archives will be created next to it"
        )
    }

    startProcess(script: string): Promise<void> {
        return new Promise((resolve, reject) => {
            const child = spawn("/data/pmtiles", script.split(" "), {
                stdio: "ignore",
                cwd: this.targetDir,
            })

            child.on("close", (code) => {
                if (code === 0) resolve()
                else reject(new Error(`Process exited with code ${code}`))
            })

            child.on("error", reject)
        })
    }

    private generateArchive(z: number, x: number, y: number, maxzoom?: number): Promise<void> {
        const [[max_lat, min_lon], [min_lat, max_lon]] = Tiles.tile_bounds(z, x, y)
        let maxzoomflag = ""
        if (maxzoom !== undefined) {
            maxzoomflag = " --maxzoom=" + maxzoom
        }
        return this.startProcess(
            `extract planet-latest.pmtiles --download-threads=1 --minzoom=${z}${maxzoomflag} --bbox=${[
                min_lon,
                min_lat + 0.0001,
                max_lon,
                max_lat,
            ].join(",")} ${this.getFilename(z, x, y)}`
        )
    }

    private *generateColumnIfNeeded(
        z: number,
        x: number,
        boundary: number,
        maxzoom?: number
    ): Generator<Promise<void>> {
        const lastFileForColumn = this.getFilename(z, x, boundary - 1)
        if (existsSync(this.targetDir + "/" + lastFileForColumn)) {
            // Skip this column, already exists
            console.log("Skipping column ", x, "at zoom", z)
            this.skipped += boundary
            return
        }
        console.log(
            "Starting column",
            x,
            "at zoom",
            z,
            "as",
            this.targetDir + "/" + lastFileForColumn,
            "does not exist"
        )

        for (let y = 0; y < boundary; y++) {
            yield this.generateArchive(z, x, y, maxzoom)
        }
    }

    private *generateField(z: number, maxzoom?: number): Generator<Promise<void>> {
        const boundary = 2 << (z - 1)

        for (let x = 0; x < boundary; x++) {
            for (const promise of this.generateColumnIfNeeded(z, x, boundary, maxzoom)) {
                yield promise
            }
        }
    }

    private getFilename(z: number, x: number, y: number) {
        return `${z}-${x}-${y}.pmtiles`
    }

    private *generateAll(): Generator<Promise<void>> {
        const zoomlevels: Record<number, number> = OfflineBasemapManager.zoomelevels
        for (const key in zoomlevels) {
            const minzoom: number = Number(key)
            const maxzoom: number | undefined = zoomlevels[key]
            for (const promise of this.generateField(minzoom, maxzoom)) {
                yield promise
            }
        }
    }

    createBatch<T>(generator: Generator<T>, length: number): T[] {
        const batch = []
        do {
            const next = generator.next()
            if (next.done) {
                return batch
            }
            batch.push(next.value)
        } while (batch.length < length)
        return batch
    }

    async main(args: string[]): Promise<void> {
        this.targetDir = args[0]
        if (!this.targetDir) {
            console.log("Please specify a target directory")
            return
        }
        let estimate = 0
        for (const key in OfflineBasemapManager.zoomelevels) {
            const z: number = Number(key)
            const boundary = 2 << z
            estimate += boundary * boundary
        }
        console.log("Target dir is:", this.targetDir)
        const numberOfThreads = 512
        const generator = this.generateAll()
        let batch: Promise<void>[] = []
        let done = 0
        const startDate = new Date()
        do {
            batch = this.createBatch(generator, numberOfThreads)
            await Promise.all(batch)
            done += batch.length
            const now = new Date()
            const timeElapsed = (now.getTime() - startDate.getTime()) / 1000
            const speed = ("" + done / timeElapsed).substring(0, 5)
            const perc = ("" + (100 * (done + this.skipped)) / estimate).substring(0, 5)
            const etaSecond = Math.floor(((estimate - done - this.skipped) * timeElapsed) / done)
            console.log(
                "Completed",
                numberOfThreads,
                `processes;  ${
                    done + this.skipped
                } / ${estimate}, ${perc}%, ${speed} tile/second, ETA: ${Utils.toHumanTime(
                    etaSecond
                )}`
            )
        } while (batch.length > 0)
    }
}

new GeneratePmTilesExtracts().run()