Make "to" range optional

Author: blootsvoets

src/main/java/org/radarbase/output/accounting/OffsetIntervals.kt

@@ -21,6 +21,7 @@ class OffsetIntervals {
     }
 
     fun contains(range: OffsetRangeSet.Range): Boolean {
+        checkNotNull(range.to)
         //  -index-1 if not found
         val searchIndex = offsetsFrom.binarySearch(range.from)
         val index = if (searchIndex >= 0) searchIndex else -searchIndex - 2
@@ -80,6 +81,7 @@ class OffsetIntervals {
 
     fun add(range: OffsetRangeSet.Range) {
         val (from, to, lastModified) = range
+        checkNotNull(to)
         var index = offsetsFrom.binarySearch(from)
         if (index < 0) {
             // index where this range would be entered

src/main/java/org/radarbase/output/accounting/OffsetRangeSet.kt

@@ -148,8 +148,8 @@ class OffsetRangeSet {
             .getOrDefault(this, EMPTY_VALUE)
             .read(function)
 
-    data class Range(val from: Long, val to: Long, val lastProcessed: Instant) {
-        val size: Long = to - from + 1
+    data class Range(val from: Long, val to: Long?, val lastProcessed: Instant) {
+        val size: Long? = to?.let { it - from + 1 }
         override fun toString() = "($from - $to, $lastProcessed)"
     }
 

src/main/java/org/radarbase/output/accounting/TopicPartitionOffsetRange.kt

@@ -30,22 +30,25 @@ data class TopicPartitionOffsetRange(
     val partition: Int = topicPartition.partition
 
     /** Full constructor.  */
-    constructor(topic: String, partition: Int, offsetFrom: Long, offsetTo: Long, lastModified: Instant = Instant.now()) : this(
+    constructor(topic: String, partition: Int, offsetFrom: Long, offsetTo: Long?, lastModified: Instant = Instant.now()) : this(
             TopicPartition(topic, partition),
             OffsetRangeSet.Range(offsetFrom, offsetTo, lastModified))
 
     override fun toString(): String = "$topic+$partition+${range.from}+${range.to} (${range.lastProcessed})"
 
     companion object {
+        private val filenameSplitRegex = "[+.]".toRegex()
+
         @Throws(NumberFormatException::class, IndexOutOfBoundsException::class)
         fun parseFilename(filename: String, lastModified: Instant): TopicPartitionOffsetRange {
-            val fileNameParts = filename.split("[+.]".toRegex()).dropLastWhile { it.isEmpty() }.toTypedArray()
+            val fileNameParts = filename.split(filenameSplitRegex)
+                    .dropLastWhile { it.isEmpty() || it == "avro" }
 
             return TopicPartitionOffsetRange(
                     fileNameParts[0],
-                    Integer.parseInt(fileNameParts[1]),
-                    java.lang.Long.parseLong(fileNameParts[2]),
-                    java.lang.Long.parseLong(fileNameParts[3]),
+                    fileNameParts[1].toInt(),
+                    fileNameParts[2].toLong(),
+                    fileNameParts.getOrNull(3)?.toLong(),
                     lastModified)
         }
     }

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

@@ -5,13 +5,15 @@ import java.nio.file.Path
 import java.time.Instant
 
 data class TopicFileList(val topic: String, val files: List<TopicFile>) {
-    val numberOfOffsets: Long = this.files.fold(0L) { sum, f -> sum + f.size }
+    val numberOfOffsets: Long? = this.files
+            .takeIf { fs -> fs.none { it.size == null } }
+            ?.fold(0L) { sum, f -> sum + f.size!! }
     val numberOfFiles: Int = this.files.size
 }
 
 data class TopicFile(val topic: String, val path: Path, val lastModified: Instant) {
     val range: TopicPartitionOffsetRange = TopicPartitionOffsetRange.parseFilename(path.fileName.toString(), lastModified)
-    val size: Long = range.range.size
+    val size: Long? = range.range.size
 }
 
 data class SimpleFileStatus(val path: Path, val isDirectory: Boolean, val lastModified: Instant)

src/main/java/org/radarbase/output/worker/RadarKafkaRestructure.kt

@@ -146,7 +146,12 @@ class RadarKafkaRestructure(
             .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.
-                    seenFiles.contains(f.range.copy(range = f.range.range.copy(to = f.range.range.to + 1))) }
+                    if (f.range.range.to != null) {
+                        seenFiles.contains(f.range.copy(range = f.range.range.copy(to = f.range.range.to + 1)))
+                    } else {
+                        seenFiles.contains(f.range.topicPartition, f.range.range.from, f.range.range.lastProcessed)
+                    }
+            }
             .take(maxFilesPerTopic)
             .map { kafkaStorage.delete(it.path) }
             .count()

src/main/java/org/radarbase/output/worker/RestructureWorker.kt

@@ -41,6 +41,7 @@ internal class RestructureWorker(
     fun processPaths(topicPaths: TopicFileList) {
         val numFiles = topicPaths.numberOfFiles
         val numOffsets = topicPaths.numberOfOffsets
+        val totalProgress = numOffsets ?: numFiles.toLong()
 
         val topic = topicPaths.topic
 
@@ -52,30 +53,35 @@ internal class RestructureWorker(
         val seenOffsets = accountant.offsets
                 .withFactory { ReadOnlyFunctionalValue(it) }
 
-        val progressBar = ProgressBar(topic, numOffsets, 50, 5, TimeUnit.SECONDS)
+        val progressBar = ProgressBar(topic, totalProgress, 50, 5, TimeUnit.SECONDS)
         progressBar.update(0)
 
         val batchSize = (BATCH_SIZE * ThreadLocalRandom.current().nextDouble(0.75, 1.25)).roundToLong()
-        var currentSize: Long = 0
+        var currentSize: Long = 0L
+        var currentFile: Long = 0L
         try {
             for (file in topicPaths.files) {
                 if (closed.get()) {
                     break
                 }
-                try {
+                val processedSize = try {
                     this.processFile(file, progressBar, seenOffsets)
                 } catch (exc: JsonMappingException) {
                     logger.error("Cannot map values", exc)
+                    0L
                 }
 
-                processedFileCount++
-                progressBar.update(currentSize)
-
-                currentSize += file.size
+                currentFile += 1
+                currentSize += processedSize
                 if (currentSize >= batchSize) {
                     currentSize = 0
                     cacheStore.flush()
                 }
+
+                processedFileCount++
+                if (numOffsets == null) {
+                    progressBar.update(currentFile)
+                }
             }
         } catch (ex: IOException) {
             logger.error("Failed to process file", ex)
@@ -85,26 +91,28 @@ internal class RestructureWorker(
             logger.warn("Shutting down")
         }
 
-        progressBar.update(numOffsets)
+        progressBar.update(totalProgress)
     }
 
     @Throws(IOException::class)
     private fun processFile(file: TopicFile,
-                            progressBar: ProgressBar, seenOffsets: OffsetRangeSet) {
+                            progressBar: ProgressBar, seenOffsets: OffsetRangeSet): Long {
         logger.debug("Reading {}", file.path)
 
         val offset = file.range.range.from
 
-        reader.newInput(file).use { input ->
+        return reader.newInput(file).use { input ->
             // processing zero-length files may trigger a stall. See:
             // https://github.com/RADAR-base/Restructure-HDFS-topic/issues/3
             if (input.length() == 0L) {
                 logger.warn("File {} has zero length, skipping.", file.path)
-                return
+                return 0L
             }
             val transaction = Accountant.Transaction(file.range.topicPartition, offset, file.lastModified)
+            var recordsSeen = 0L
             extractRecords(input) { relativeOffset, record ->
                 transaction.offset = offset + relativeOffset
+                recordsSeen += 1L
                 val alreadyContains = time("accounting.check") {
                     seenOffsets.contains(file.range.topicPartition, transaction.offset, transaction.lastModified)
                 }
@@ -113,8 +121,11 @@ internal class RestructureWorker(
                     this.writeRecord(transaction, record)
                 }
                 processedRecordsCount++
-                progressBar.update(processedRecordsCount)
+                if (file.size != null) {
+                    progressBar.update(processedRecordsCount)
+                }
             }
+            recordsSeen
         }
     }