main.app.cash.backfila.client.stat.internal.StaticDatasourceBackfillOperator.kt Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of client-static Show documentation
Backfila is a service that manages backfill state, calling into other services to do batched work.
There is a newer version: 2024.10.28.205607-fab304f
package app.cash.backfila.client.stat.internal

import app.cash.backfila.client.spi.BackfilaParametersOperator
import app.cash.backfila.client.spi.BackfillOperator
import app.cash.backfila.client.stat.StaticDatasourceBackfillBase
import app.cash.backfila.protos.clientservice.GetNextBatchRangeRequest
import app.cash.backfila.protos.clientservice.GetNextBatchRangeResponse
import app.cash.backfila.protos.clientservice.KeyRange
import app.cash.backfila.protos.clientservice.PrepareBackfillRequest
import app.cash.backfila.protos.clientservice.PrepareBackfillResponse
import app.cash.backfila.protos.clientservice.RunBatchRequest
import app.cash.backfila.protos.clientservice.RunBatchResponse
import com.google.common.base.Preconditions.checkArgument
import okio.ByteString.Companion.encodeUtf8

class StaticDatasourceBackfillOperator(
  override val backfill: StaticDatasourceBackfillBase,
  val parametersOperator: BackfilaParametersOperator,
) : BackfillOperator {

  override fun name(): String = backfill.javaClass.toString()

  override fun prepareBackfill(request: PrepareBackfillRequest): PrepareBackfillResponse {
    val config = parametersOperator.constructBackfillConfig(request)
    backfill.validate(config)
    val datasource = backfill.getStaticDatasource(config)

    val start = request.range?.start?.utf8()?.let {
      it.toIntOrNull() ?: error("Start of range must be a number")
    } ?: 0
    val end = request.range?.end?.utf8()?.let {
      it.toIntOrNull() ?: error("End of range must be a number")
    } ?: datasource.size

    // Sanity check that this backfill will actually process something
    require(start >= 0 && end >= 0) {
      "Start and end must be positive integers, start: $start end: $end"
    }
    require(start <= end) {
      "Start must be less than or equal to end, start: $start end: $end"
    }
    require(start <= datasource.size) {
      "Start is greater than the static datasource size, start: $start size: ${datasource.size}"
    }

    val onlyPartition = PrepareBackfillResponse.Partition.Builder()
      .partition_name(PARTITION)
      .backfill_range(
        KeyRange.Builder()
          .start(start.toString().encodeUtf8())
          .end(end.toString().encodeUtf8())
          .build(),
      )
      .build()

    return PrepareBackfillResponse.Builder()
      .partitions(listOf(onlyPartition))
      .build()
  }

  override fun getNextBatchRange(request: GetNextBatchRangeRequest): GetNextBatchRangeResponse {
    checkArgument(request.partition_name == PARTITION, "Attempting to get batch for unknown partition ${request.partition_name}")
    val batchSize = request.batch_size.toInt()
    val scanSize = request.scan_size.toInt()

    val backfillRange = request.backfill_range.decode()

    val batchingStart = request.previous_end_key?.utf8()?.toInt() ?: backfillRange.start

    val batches = mutableListOf()
    for (batchStart in batchingStart until minOf(batchingStart + scanSize, backfillRange.end) step batchSize) {
      val batchEnd = minOf(batchStart + batchSize, backfillRange.end)
      batches += GetNextBatchRangeResponse.Batch.Builder()
        .batch_range(
          KeyRange.Builder()
            .start(batchStart.toString().encodeUtf8())
            .end(batchEnd.toString().encodeUtf8())
            .build(),
        )
        .matching_record_count((batchEnd - batchStart).toLong())
        .scanned_record_count((batchEnd - batchStart).toLong())
        .build()
    }

    return GetNextBatchRangeResponse.Builder()
      .batches(batches)
      .build()
  }

  override fun runBatch(request: RunBatchRequest): RunBatchResponse {
    checkArgument(request.partition_name == PARTITION, "Attempting to get batch for unknown partition ${request.partition_name}")
    val batchRange = request.batch_range.decode()
    val config = parametersOperator.constructBackfillConfig(request)

    val batch = backfill.getStaticDatasource(config.prepareConfig()).subList(batchRange.start, batchRange.end)

    backfill.runBatch(batch, config)

    return RunBatchResponse.Builder()
      .build()
  }

  data class DecodedRange(
    val start: Int,
    val end: Int,
  )
  private fun KeyRange.decode(): DecodedRange {
    val start = this.start.utf8().toInt()
    val end = this.end.utf8().toInt()
    return DecodedRange(start, end)
  }

  companion object {
    private const val PARTITION = "only"
  }
}