Spark Structured Streaming 中的容错

updateStatusMessage("Writing offsets to log")
reportTimeTaken("walCommit") {
  assert(offsetLog.add(
    currentBatchId,
    availableOffsets.toOffsetSeq(sources, offsetSeqMetadata)),
    s"Concurrent update to the log. Multiple streaming jobs detected for $currentBatchId")
    logInfo(s"Committed offsets for batch $currentBatchId. " +
    s"Metadata ${offsetSeqMetadata.toString}")
 
  // NOTE: The following code is correct because runStream() processes exactly one
  // batch at a time. If we add pipeline parallelism (multiple batches in flight at
  // the same time), this cleanup logic will need to change.
 
  // Now that we've updated the scheduler's persistent checkpoint, it is safe for the
  // sources to discard data from the previous batch.
  if (currentBatchId != 0) {
    val prevBatchOff = offsetLog.get(currentBatchId - 1)
    if (prevBatchOff.isDefined) {
      prevBatchOff.get.toStreamProgress(sources).foreach {
        case (src: Source, off) => src.commit(off)
        case (reader: MicroBatchReader, off) =>
          reader.commit(reader.deserializeOffset(off.json))
      }
    } else {
      throw new IllegalStateException(s"batch $currentBatchId doesn't exist")
    }
  }
   
  // It is now safe to discard the metadata beyond the minimum number to retain.
  // Note that purge is exclusive, i.e. it purges everything before the target ID.
  if (minLogEntriesToMaintain < currentBatchId) {
    offsetLog.purge(currentBatchId - minLogEntriesToMaintain)
    commitLog.purge(currentBatchId - minLogEntriesToMaintain)
  }
}

override def beforeAll(): Unit = {
  Path(TestConfiguration.TestDirInput).createDirectory()
  Path(TestConfiguration.TestDirOutput).createDirectory()
  for (i <- 1 to 10) {
    val file = s"file${i}"
    val content =
      s"""
        |{"id": 1, "name": "content1=${i}"}
        |{"id": 2, "name": "content2=${i}"}
      """.stripMargin
    File(s"${TestConfiguration.TestDirInput}/${file}").writeAll(content)
  }
}
 
override def afterAll(): Unit = {
  Path(TestConfiguration.TestDirInput).deleteRecursively()
  Path(TestConfiguration.TestDirOutput).deleteRecursively()
}
 
 
"after one failure" should "all rows should be processed and output in idempotent manner" in {
  for (i <- 0 until 2) {
    val sparkSession: SparkSession = SparkSession.builder()
      .appName("Spark Structured Streaming fault tolerance example")
      .master("local[2]").getOrCreate()
    import sparkSession.implicits._
    try {
      val schema = StructType(StructField("id", LongType, nullable = false) ::
        StructField("name", StringType, nullable = false) :: Nil)
 
      val idNameInputStream = sparkSession.readStream.format("json").schema(schema).option("maxFilesPerTrigger", 1)
        .load(TestConfiguration.TestDirInput)
        .toDF("id", "name")
        .map(row => {
          val fieldName = row.getAs[String]("name")
          if (fieldName == "content1=7" && !GlobalFailureFlag.alreadyFailed.get() && GlobalFailureFlag.mustFail.get()) {
            GlobalFailureFlag.alreadyFailed.set(true)
            GlobalFailureFlag.mustFail.set(false)
            throw new RuntimeException("Something went wrong")
          }
          fieldName
        })
 
      val query = idNameInputStream.writeStream.outputMode("update").foreach(new ForeachWriter[String] {
        override def process(value: String): Unit = {
          val fileName = value.replace("=", "_")
          File(s"${TestConfiguration.TestDirOutput}/${fileName}").writeAll(value)
        }
 
        override def close(errorOrNull: Throwable): Unit = {}
 
        override def open(partitionId: Long, version: Long): Boolean = true
      }).start()
      query.awaitTermination(15000)
    } catch {
      case re: StreamingQueryException => {
        println("As expected, RuntimeException was thrown")
      }
    }
  }
 
  val outputFiles = Path(TestConfiguration.TestDirOutput).toDirectory.files.toSeq
  outputFiles should have size 20
  val filesContent = outputFiles.map(file => Source.fromFile(file.path).getLines.mkString(""))
  val expected = (1 to 2).flatMap(id => {
    (1 to 10).map(nr => s"content${id}=${nr}")
  })
  filesContent should contain allElementsOf expected
}
 
object TestConfiguration {
  val TestDirInput = "/tmp/spark-fault-tolerance"
  val TestDirOutput = s"${TestDirInput}-output"
}
 
object GlobalFailureFlag {
  var mustFail = new AtomicBoolean(true)
  var alreadyFailed = new AtomicBoolean(false)
}