Added dedicated coroutine tree lister

Author: blootsvoets

src/main/java/org/radarbase/output/Application.kt

@@ -75,10 +75,15 @@ class Application(
 
     override val workerSemaphore = Semaphore(config.worker.numThreads * 2)
 
-    private val jobs = listOfNotNull(
-        Job("restructure", config.service.interval, ::runRestructure).takeIf { config.worker.enable },
-        Job("clean", config.cleaner.interval, ::runCleaner).takeIf { config.cleaner.enable },
-    )
+    private val jobs: List<Job>
+
+    init {
+        val serviceMutex = Mutex()
+        jobs = listOfNotNull(
+            RadarKafkaRestructure.job(config, serviceMutex),
+            SourceDataCleaner.job(config, serviceMutex),
+        )
+    }
 
     @Throws(IOException::class)
     override fun newFileCacheStore(accountant: Accountant) = FileCacheStore(this, accountant)
@@ -105,15 +110,10 @@ class Application(
         if (config.service.enable) {
             runService()
         } else {
-            val serviceMutex = Mutex()
             runBlocking {
-                jobs.map {
-                    async {
-                        serviceMutex.withLock {
-                            it.run()
-                        }
-                    }
-                }.awaitAll()
+                jobs.forEach { job ->
+                    launch { job.run(this@Application) }
+                }
             }
         }
     }
@@ -122,47 +122,18 @@ class Application(
         logger.info("Running as a Service with poll interval of {} seconds", config.service.interval)
         logger.info("Press Ctrl+C to exit...")
 
-        val serviceMutex = Mutex()
-
         runBlocking {
-            jobs.forEach {
-                launch { it.schedule(serviceMutex) }
-            }
-        }
-    }
-
-    private suspend fun runCleaner() {
-        SourceDataCleaner(this).useSuspended { cleaner ->
-            for (input in config.paths.inputs) {
-                logger.info("Cleaning {}", input)
-                cleaner.process(input.toString())
-            }
-            logger.info("Cleaned up {} files",
-                    cleaner.deletedFileCount.format())
-        }
-    }
-
-    private suspend fun runRestructure() {
-        RadarKafkaRestructure(this).useSuspended { restructure ->
-            for (input in config.paths.inputs) {
-                logger.info("In:  {}", input)
-                logger.info("Out: {}", pathFactory.root)
-                restructure.process(input.toString())
+            jobs.forEach { job ->
+                launch { job.schedule(this@Application) }
             }
-
-            logger.info(
-                "Processed {} files and {} records",
-                restructure.processedFileCount.format(),
-                restructure.processedRecordsCount.format(),
-            )
         }
     }
 
     companion object {
         private val logger = LoggerFactory.getLogger(Application::class.java)
         const val CACHE_SIZE_DEFAULT = 100
 
-        private fun LongAdder.format(): String =
+        internal fun LongAdder.format(): String =
                 NumberFormat.getNumberInstance().format(sum())
 
         private fun parseArgs(args: Array<String>): CommandLineArgs {

src/main/java/org/radarbase/output/cleaner/SourceDataCleaner.kt

@@ -1,13 +1,18 @@
 package org.radarbase.output.cleaner
 
 import kotlinx.coroutines.*
-import kotlinx.coroutines.flow.*
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
 import kotlinx.coroutines.sync.withPermit
+import org.radarbase.output.Application.Companion.format
 import org.radarbase.output.FileStoreFactory
 import org.radarbase.output.accounting.Accountant
 import org.radarbase.output.accounting.AccountantImpl
+import org.radarbase.output.config.RestructureConfig
 import org.radarbase.output.util.ResourceContext.Companion.resourceContext
+import org.radarbase.output.util.SuspendedCloseable.Companion.useSuspended
 import org.radarbase.output.util.Timer
+import org.radarbase.output.worker.Job
 import org.slf4j.LoggerFactory
 import java.io.Closeable
 import java.io.IOException
@@ -33,6 +38,7 @@ class SourceDataCleaner(
             .minus(fileStoreFactory.config.cleaner.age.toLong(), ChronoUnit.DAYS)
 
     val deletedFileCount = LongAdder()
+    private val scope = CoroutineScope(Dispatchers.Default)
 
     @Throws(IOException::class, InterruptedException::class)
     suspend fun process(directoryName: String) {
@@ -43,23 +49,21 @@ class SourceDataCleaner(
 
         logger.info("{} topics found", paths.size)
 
-        coroutineScope {
-            paths.forEach { p ->
-                launch {
-                    try {
-                        val deleteCount = fileStoreFactory.workerSemaphore.withPermit {
-                            mapTopic(p)
-                        }
-                        if (deleteCount > 0) {
-                            logger.info("Removed {} files in topic {}", deleteCount, p.fileName)
-                            deletedFileCount.add(deleteCount)
-                        }
-                    } catch (ex: Exception) {
-                        logger.warn("Failed to map topic", ex)
+        paths.map { p ->
+            scope.launch {
+                try {
+                    val deleteCount = fileStoreFactory.workerSemaphore.withPermit {
+                        mapTopic(p)
                     }
+                    if (deleteCount > 0) {
+                        logger.info("Removed {} files in topic {}", deleteCount, p.fileName)
+                        deletedFileCount.add(deleteCount)
+                    }
+                } catch (ex: Exception) {
+                    logger.warn("Failed to map topic", ex)
                 }
             }
-        }
+        }.joinAll()
     }
 
     private suspend fun mapTopic(topicPath: Path): Long {
@@ -93,41 +97,65 @@ class SourceDataCleaner(
     ): Int {
         val offsets = accountant.offsets.copyForTopic(topic)
 
-        return sourceStorage.walker.walkRecords(topic, topicPath).consumeAsFlow()
-            .filter { f ->
-                f.lastModified.isBefore(deleteThreshold) &&
-                    // ensure that there is a file with a larger offset also
-                    // processed, so the largest offset is never removed.
-                    offsets.contains(f.range.mapRange { r -> r.incrementTo() })
-            }
-            .take(maxFilesPerTopic)
-            .takeWhile { !isClosed.get() }
-            .count { file ->
-                if (extractionCheck.isExtracted(file)) {
-                    logger.info("Removing {}", file.path)
-                    Timer.time("cleaner.delete") {
-                        sourceStorage.delete(file.path)
+        val paths = sourceStorage.listTopicFiles(topic, topicPath, maxFilesPerTopic) { f ->
+            f.lastModified.isBefore(deleteThreshold) &&
+                // ensure that there is a file with a larger offset also
+                // processed, so the largest offset is never removed.
+                offsets.contains(f.range.mapRange { r -> r.incrementTo() })
+        }
+
+        val accountantMutex = Mutex()
+
+        return coroutineScope {
+            paths
+                .map { file ->
+                    async {
+                        if (extractionCheck.isExtracted(file)) {
+                            logger.info("Removing {}", file.path)
+                            Timer.time("cleaner.delete") {
+                                sourceStorage.delete(file.path)
+                            }
+                            true
+                        } else {
+                            // extract the file again at a later time
+                            logger.warn("Source file was not completely extracted: {}", file.path)
+                            val fullRange = file.range.mapRange { it.ensureToOffset() }
+                            accountantMutex.withLock {
+                                accountant.remove(fullRange)
+                            }
+                            false
+                        }
                     }
-                    true
-                } else {
-                    logger.warn("Source file was not completely extracted: {}", file.path)
-                    // extract the file again at a later time
-                    accountant.remove(file.range.mapRange { it.ensureToOffset() })
-                    false
                 }
-            }
+                .awaitAll()
+                .count { it }
+        }
     }
 
-    private suspend fun topicPaths(path: Path): List<Path> = sourceStorage.walker.walkTopics(path, excludeTopics)
+    private suspend fun topicPaths(path: Path): List<Path> = sourceStorage.listTopics(path, excludeTopics)
             .toMutableList()
             // different services start on different topics to decrease lock contention
             .also { it.shuffle() }
 
     override fun close() {
-        isClosed.set(true)
+        scope.cancel()
     }
 
     companion object {
         private val logger = LoggerFactory.getLogger(SourceDataCleaner::class.java)
+
+        fun job(config: RestructureConfig, serviceMutex: Mutex): Job? = if (config.cleaner.enable) {
+            Job("cleaner", config.cleaner.interval, ::runCleaner, serviceMutex)
+        } else null
+
+        private suspend fun runCleaner(factory: FileStoreFactory) {
+            SourceDataCleaner(factory).useSuspended { cleaner ->
+                for (input in factory.config.paths.inputs) {
+                    logger.info("Cleaning {}", input)
+                    cleaner.process(input.toString())
+                }
+                logger.info("Cleaned up {} files", cleaner.deletedFileCount.format())
+            }
+        }
     }
 }

src/main/java/org/radarbase/output/source/AzureSourceStorage.kt

@@ -54,8 +54,6 @@ class AzureSourceStorage(
         blobClient(path).delete()
     }
 
-    override val walker: SourceStorageWalker = GeneralSourceStorageWalker(this)
-
     override fun createReader(): SourceStorage.SourceStorageReader = AzureSourceStorageReader()
 
     private inner class AzureSourceStorageReader: SourceStorage.SourceStorageReader {

src/main/java/org/radarbase/output/source/GeneralSourceStorageWalker.kt

@@ -28,7 +28,6 @@ class S3SourceStorage(
         config: S3Config,
         private val tempPath: Path
 ): SourceStorage {
-    override val walker: SourceStorageWalker = GeneralSourceStorageWalker(this)
     private val bucket = config.bucket
     private val readEndOffset = config.endOffsetFromTags
 

src/main/java/org/radarbase/output/source/S3SourceStorage.kt

@@ -1,11 +1,10 @@
 package org.radarbase.output.source
 
-import kotlinx.coroutines.FlowPreview
-import kotlinx.coroutines.flow.Flow
 import org.apache.avro.file.SeekableInput
 import org.radarbase.output.accounting.TopicPartitionOffsetRange
-import org.radarbase.output.util.SuspendedCloseable
-import java.io.Closeable
+import org.radarbase.output.util.*
+import org.radarbase.output.util.AvroFileLister.Companion.avroFileTreeLister
+import org.radarbase.output.util.AvroTopicLister.Companion.avroTopicTreeLister
 import java.nio.file.Path
 import java.time.Instant
 
@@ -24,8 +23,28 @@ interface SourceStorage {
         return TopicFile(topic, status.path, lastModified, range)
     }
 
-    /** Find records and topics. */
-    val walker: SourceStorageWalker
+    /**
+     * Recursively returns all record files in a sequence of a given topic with path.
+     * The path must only contain records of a single topic, this is not verified.
+     */
+    suspend fun listTopicFiles(
+        topic: String,
+        topicPath: Path,
+        limit: Int,
+        predicate: (TopicFile) -> Boolean,
+    ): List<TopicFile> = avroFileTreeLister()
+        .list(TopicPath(topic, topicPath), limit, predicate)
+
+    /**
+     * Recursively find all topic root paths of records in the given path.
+     * Exclude paths belonging to the set of given excluded topics.
+     */
+    suspend fun listTopics(
+        root: Path,
+        exclude: Set<String>,
+    ): List<Path> = avroTopicTreeLister()
+        .listTo(LinkedHashSet(), root)
+        .filter { it.fileName.toString() !in exclude }
 
     /**
      * File reader for the storage medium.