Added a StorageIndex for the source storage to reduce LIST calls

Before every restructuring or cleaning operation, the index is updated.
The StorageIndex can be updated with a configurable sync time to make a
full sync, otherwise it just updates directories that have files in them.
A separate sync time can be set to also scan empty directories.
The first implementation is only a InMemoryStorageIndex. For
very large datasets, a file-based index might be needed.

Author: blootsvoets

README.md

@@ -118,8 +118,16 @@ source:
   # only actually needed if source type is hdfs
   azure:
     # azure options
+  index:
+    # Interval to fully synchronize the index with the source storage
+    fullSyncInterval: 3600
+    # Interval to sync empty directories with.
+    # They are also synced during a full sync.
+    emptyDirectorySyncInterval: 900
 ```
 
+The index makes a scan of the source before any operations. Further list operations are done on the index only. This is especially relevant for S3 storage where list operations are priced.
+
 The target is similar, and in addition supports the local file system (`local`).
 
 ```yaml

restructure.yml

@@ -34,6 +34,12 @@ source:
   # only actually needed if source type is hdfs
   hdfs:
     nameNodes: [hdfs-namenode-1, hdfs-namenode-2]
+  index:
+    # Interval to fully synchronize the index with the storage
+    fullSyncInterval: 3600
+    # Interval to sync empty directories with.
+    # They are also synced during a full sync.
+    emptyDirectorySyncInterval: 900
 
 # Target data resource
 # @since: 0.7.0

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

@@ -28,8 +28,7 @@ import org.radarbase.output.config.CommandLineArgs
 import org.radarbase.output.config.RestructureConfig
 import org.radarbase.output.format.RecordConverterFactory
 import org.radarbase.output.path.RecordPathFactory
-import org.radarbase.output.source.SourceStorage
-import org.radarbase.output.source.SourceStorageFactory
+import org.radarbase.output.source.*
 import org.radarbase.output.target.TargetStorage
 import org.radarbase.output.target.TargetStorageFactory
 import org.radarbase.output.util.Timer
@@ -39,7 +38,9 @@ import org.radarbase.output.worker.RadarKafkaRestructure
 import org.slf4j.LoggerFactory
 import redis.clients.jedis.JedisPool
 import java.io.IOException
+import java.nio.file.Path
 import java.text.NumberFormat
+import java.time.Duration
 import java.time.LocalDateTime
 import java.time.format.DateTimeFormatter
 import java.util.concurrent.atomic.LongAdder
@@ -78,9 +79,27 @@ class Application(
 
     override val workerSemaphore = Semaphore(config.worker.numThreads * 2)
 
+    override val storageIndexManagers: Map<Path, StorageIndexManager>
+
     private val jobs: List<Job>
 
     init {
+        val indexConfig = config.source.index
+        val (fullScan, emptyScan) = if (indexConfig == null) {
+            listOf(3600L, 900L)
+        } else {
+            listOf(indexConfig.fullSyncInterval, indexConfig.emptyDirectorySyncInterval)
+        }.map { Duration.ofSeconds(it) }
+
+        storageIndexManagers = config.paths.inputs.associateWith { input ->
+            MutableStorageIndexManager(
+                InMemoryStorageIndex(),
+                sourceStorage,
+                fullScan,
+                emptyScan,
+                input,
+            )
+        }
         val serviceMutex = Mutex()
         jobs = listOfNotNull(
             RadarKafkaRestructure.job(config, serviceMutex),

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

@@ -26,9 +26,11 @@ import org.radarbase.output.config.RestructureConfig
 import org.radarbase.output.format.RecordConverterFactory
 import org.radarbase.output.path.RecordPathFactory
 import org.radarbase.output.source.SourceStorage
+import org.radarbase.output.source.StorageIndexManager
 import org.radarbase.output.target.TargetStorage
 import org.radarbase.output.worker.FileCacheStore
 import java.io.IOException
+import java.nio.file.Path
 
 /** Factory for all factory classes and settings.  */
 interface FileStoreFactory {
@@ -42,6 +44,7 @@ interface FileStoreFactory {
     val redisHolder: RedisHolder
     val offsetPersistenceFactory: OffsetPersistenceFactory
     val workerSemaphore: Semaphore
+    val storageIndexManagers: Map<Path, StorageIndexManager>
 
     @Throws(IOException::class)
     fun newFileCacheStore(accountant: Accountant): FileCacheStore

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

@@ -11,6 +11,8 @@ 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.source.StorageIndex
+import org.radarbase.output.source.StorageNode
 import org.radarbase.output.util.ResourceContext.Companion.resourceContext
 import org.radarbase.output.util.SuspendedCloseable.Companion.useSuspended
 import org.radarbase.output.util.Timer
@@ -43,11 +45,11 @@ class SourceDataCleaner(
     private val supervisor = SupervisorJob()
 
     @Throws(IOException::class, InterruptedException::class)
-    suspend fun process(directoryName: String) {
+    suspend fun process(storageIndex: StorageIndex, directoryName: String) {
         // Get files and directories
         val absolutePath = Paths.get(directoryName)
 
-        val paths = topicPaths(absolutePath)
+        val paths = topicPaths(storageIndex, absolutePath)
 
         logger.info("{} topics found", paths.size)
 
@@ -56,7 +58,7 @@ class SourceDataCleaner(
                 launch {
                     try {
                         val deleteCount = fileStoreFactory.workerSemaphore.withPermit {
-                            mapTopic(p)
+                            mapTopic(storageIndex, p)
                         }
                         if (deleteCount > 0) {
                             logger.info("Removed {} files in topic {}", deleteCount, p.fileName)
@@ -70,7 +72,7 @@ class SourceDataCleaner(
         }
     }
 
-    private suspend fun mapTopic(topicPath: Path): Long {
+    private suspend fun mapTopic(storageIndex: StorageIndex, topicPath: Path): Long {
         val topic = topicPath.fileName.toString()
         return try {
             lockManager.tryWithLock(topic) {
@@ -84,7 +86,7 @@ class SourceDataCleaner(
                                 fileStoreFactory,
                             )
                         }
-                        deleteOldFiles(accountant, extractionCheck, topic, topicPath).toLong()
+                        deleteOldFiles(storageIndex, accountant, extractionCheck, topic, topicPath).toLong()
                     }
                 }
             }
@@ -95,14 +97,15 @@ class SourceDataCleaner(
     }
 
     private suspend fun deleteOldFiles(
+        storageIndex: StorageIndex,
         accountant: Accountant,
         extractionCheck: ExtractionCheck,
         topic: String,
         topicPath: Path,
     ): Int {
         val offsets = accountant.offsets.copyForTopic(topic)
 
-        val paths = sourceStorage.listTopicFiles(topic, topicPath, maxFilesPerTopic) { f ->
+        val paths = sourceStorage.listTopicFiles(storageIndex, 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.
@@ -115,6 +118,7 @@ class SourceDataCleaner(
                     logger.info("Removing {}", file.path)
                     Timer.time("cleaner.delete") {
                         sourceStorage.delete(file.path)
+                        storageIndex.remove(StorageNode.StorageFile(file.path, Instant.MIN))
                     }
                     true
                 } else {
@@ -127,8 +131,8 @@ class SourceDataCleaner(
             }
     }
 
-    private suspend fun topicPaths(path: Path): List<Path> =
-        sourceStorage.listTopics(path, excludeTopics)
+    private suspend fun topicPaths(storageIndex: StorageIndex, path: Path): List<Path> =
+        sourceStorage.listTopics(storageIndex, path, excludeTopics)
             // different services start on different topics to decrease lock contention
             .shuffled()
 
@@ -147,9 +151,10 @@ class SourceDataCleaner(
 
         private suspend fun runCleaner(factory: FileStoreFactory) {
             SourceDataCleaner(factory).useSuspended { cleaner ->
-                for (input in factory.config.paths.inputs) {
+                for ((input, indexManager) in factory.storageIndexManagers) {
+                    indexManager.update()
                     logger.info("Cleaning {}", input)
-                    cleaner.process(input.toString())
+                    cleaner.process(indexManager.storageIndex, input.toString())
                 }
                 logger.info("Cleaned up {} files", cleaner.deletedFileCount.format())
             }

src/main/java/org/radarbase/output/config/ResourceConfig.kt

@@ -11,6 +11,7 @@ data class ResourceConfig(
     val hdfs: HdfsConfig? = null,
     val local: LocalConfig? = null,
     val azure: AzureConfig? = null,
+    val index: StorageIndexConfig? = null,
 ) {
     @get:JsonIgnore
     val sourceType: ResourceType by lazy {
@@ -33,3 +34,4 @@ data class ResourceConfig(
         ResourceType.AZURE -> copyOnChange(azure, { it?.withEnv(prefix) }) { copy(azure = it) }
     }
 }
+

src/main/java/org/radarbase/output/config/StorageIndexConfig.kt

@@ -0,0 +1,12 @@
+package org.radarbase.output.config
+
+data class StorageIndexConfig(
+    /** How often to fully sync the storage index, in seconds. */
+    val fullSyncInterval: Long = 3600L,
+    /**
+     * How often to sync empty directories with the storage index, in seconds.
+     * If this is very large, empty directories will only be scanned during
+     * full sync.
+     */
+    val emptyDirectorySyncInterval: Long = 900L,
+)

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

@@ -11,6 +11,7 @@ import org.radarbase.output.util.TemporaryDirectory
 import org.radarbase.output.util.withoutFirstSegment
 import java.nio.file.Path
 import java.nio.file.Paths
+import java.time.Instant
 import kotlin.io.path.createTempFile
 import kotlin.io.path.deleteIfExists
 
@@ -24,22 +25,28 @@ class AzureSourceStorage(
 
     private fun blobClient(path: Path) = blobContainerClient.getBlobClient(path.withoutFirstSegment())
 
-    override suspend fun list(path: Path, maxKeys: Int?): List<SimpleFileStatus> =
+    override suspend fun list(path: Path, startAfter: Path?, maxKeys: Int?): List<StorageNode> =
         withContext(Dispatchers.IO) {
             var iterable: Iterable<BlobItem> = blobContainerClient.listBlobsByHierarchy("$path/")
+            if (startAfter != null) {
+                iterable = iterable.filter { Paths.get(it.name) > startAfter }
+            }
             if (maxKeys != null) {
                 iterable = iterable.take(maxKeys)
             }
             iterable.map {
-                SimpleFileStatus(
-                    Paths.get(it.name),
-                    it.isPrefix ?: false,
-                    it.properties?.lastModified?.toInstant(),
-                )
+                if (it.isPrefix == true) {
+                    StorageNode.StorageFile(
+                        Paths.get(it.name),
+                        it.properties?.lastModified?.toInstant() ?: Instant.now(),
+                    )
+                } else {
+                    StorageNode.StorageDirectory(Paths.get(it.name))
+                }
             }
         }
 
-    override suspend fun createTopicFile(topic: String, status: SimpleFileStatus): TopicFile {
+    override suspend fun createTopicFile(topic: String, status: StorageNode): TopicFile {
         var topicFile = super.createTopicFile(topic, status)
 
         if (readOffsetFromMetadata && topicFile.range.range.to == null) {

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

@@ -0,0 +1,10 @@
+package org.radarbase.output.source
+
+class DelegatingStorageIndex(
+    private val sourceStorage: SourceStorage,
+) : StorageIndex {
+    override suspend fun list(dir: StorageNode.StorageDirectory, maxKeys: Int?): List<StorageNode> =
+        sourceStorage.list(dir.path, maxKeys = maxKeys)
+
+    override suspend fun remove(file: StorageNode.StorageFile) = Unit
+}

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

@@ -0,0 +1,121 @@
+package org.radarbase.output.source
+
+import org.radarbase.output.source.StorageIndex.Companion.ROOT
+import java.nio.file.Path
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.ConcurrentMap
+
+class InMemoryStorageIndex : MutableStorageIndex {
+    private val fileIndex: ConcurrentMap<StorageNode.StorageDirectory, Map<Path, StorageNode>> = ConcurrentHashMap()
+    private val rootSet = ConcurrentHashMap<Path, StorageNode>()
+
+    init {
+        fileIndex[ROOT] = rootSet
+    }
+
+    override suspend fun list(dir: StorageNode.StorageDirectory, maxKeys: Int?): List<StorageNode> {
+        val listing = if (dir === ROOT) {
+            rootSet
+        } else {
+            fileIndex[dir] ?: return listOf()
+        }
+
+        return if (maxKeys != null) {
+            listing.values.take(maxKeys)
+        } else {
+            listing.values.toList()
+        }
+    }
+
+    private fun add(node: StorageNode) {
+        var currentNode = node
+        var parent = currentNode.parent()
+        if (currentNode is StorageNode.StorageDirectory) {
+            fileIndex.computeIfAbsent(currentNode) {
+                mapOf()
+            }
+        }
+        while (parent != null) {
+            fileIndex.compute(parent) { _, map ->
+                if (map == null) {
+                    mapOf(currentNode.path to currentNode)
+                } else {
+                    val newMap = map.toMutableMap()
+                    newMap[currentNode.path] = currentNode
+                    newMap
+                }
+            }
+            currentNode = parent
+            parent = currentNode.parent()
+        }
+        rootSet[currentNode.path] = currentNode
+    }
+
+    override suspend fun addAll(parent: StorageNode.StorageDirectory, nodes: List<StorageNode>): Collection<StorageNode> {
+        add(parent)
+        nodes.asSequence()
+            .filterIsInstance<StorageNode.StorageDirectory>()
+            .forEach { node ->
+                fileIndex.computeIfAbsent(node) {
+                    mapOf()
+                }
+            }
+        val newMap = fileIndex.compute(parent) { _, map ->
+            if (map == null) {
+                buildMap(nodes.size) {
+                    nodes.forEach { put(it.path, it) }
+                }
+            } else {
+                buildMap(nodes.size + map.size) {
+                    putAll(map)
+                    nodes.forEach { put(it.path, it) }
+                }
+            }
+        } ?: mapOf()
+
+        return newMap.values
+    }
+
+    override suspend fun sync(parent: StorageNode.StorageDirectory, nodes: List<StorageNode>): Collection<StorageNode> {
+        add(parent)
+        val newMap = fileIndex.compute(parent) { _, map ->
+            if (map == null) {
+                buildMap(nodes.size) {
+                    nodes.forEach { put(it.path, it) }
+                }
+            } else {
+                buildMap(nodes.size) {
+                    nodes.forEach { put(it.path, it) }
+                }
+            }
+        } ?: mapOf()
+
+        nodes.asSequence()
+            .filterIsInstance<StorageNode.StorageDirectory>()
+            .filter { it.path !in newMap }
+            .forEach { removeRecursive(it) }
+
+        return newMap.values
+    }
+
+    override suspend fun remove(file: StorageNode.StorageFile) {
+        val parent = file.parent()
+
+        if (parent != null) {
+            fileIndex.computeIfPresent(parent) { _, map ->
+                (map - file.path).takeIf { it.isNotEmpty() }
+            }
+        } else {
+            rootSet.remove(file.path)
+        }
+    }
+
+    private fun removeRecursive(node: StorageNode.StorageDirectory) {
+        val directoriesToRemove = ArrayDeque<StorageNode.StorageDirectory>()
+        fileIndex.remove(node)?.values?.filterIsInstanceTo(directoriesToRemove)
+        while (directoriesToRemove.isNotEmpty()) {
+            val first = directoriesToRemove.removeFirst()
+            fileIndex.remove(first)?.values?.filterIsInstanceTo(directoriesToRemove)
+        }
+    }
+}