Structured Streaming 源码剖析（一）- Source

// 返回此 Source 的数据的 schema
def schema: StructType

// 返回此 Source 的最大可用 offset
// 如果此 Source 从未接收过任何数据，则返回 None
def getOffset: Option[Offset]


// 返回 （start，end] 偏移量之间的数据。当 start 为 None 时，批处理应以第一个记录开头。此方法必须始终为特定的 start 和 end 对返回相同的数据; 即使在另一个节点上重新启动 Source 之后也是如此。
// 更上层总是调用此方法，其值 start 大于或等于传递给 commit 的最后一个值，而 end 值小于或等于 getOffset 返回的最后一个值
// 当从日志中获取数据时，offset 的类型可能是 SerializedOffset。此外，StreamExecution 仅比较 Offset JSON 以确定两个对象是否相等。修改 Offset JSON 格式时可能会产生冲突，在这种情况下，Source应该返回一个空的DataFrame
def getBatch(start: Option[Offset], end: Offset): DataFrame

// 通知 Source 已完成处理到 end 偏移量的所有数据，并且将来只会请求大于 end 的偏移量
def commit(end: Offset) : Unit

// 停止此 Source 并释放它已分配的所有资源
def stop(): Unit

// Offset 的 JSON 序列化表示，用于将偏移量保存到 offsetLog
// 注意：我们假设 等效/相等 offset 序列化为相同的 JSON 字符串
public abstract String json();

@Override
public boolean equals(Object obj) {
    if (obj instanceof Offset) {
        return this.json().equals(((Offset) obj).json());
    } else {
        return false;
    }
}

case class SerializedOffset(override val json: String) extends Offset

def kafkaSchema: StructType = StructType(Seq(
    StructField("key", BinaryType),
    StructField("value", BinaryType),
    StructField("topic", StringType),
    StructField("partition", IntegerType),
    StructField("offset", LongType),
    StructField("timestamp", TimestampType),
    StructField("timestampType", IntegerType)
  ))

// 包含各个 topic partition 对应的 offset
case class KafkaSourceOffset(partitionToOffsets: Map[TopicPartition, Long]) extends OffsetV2 {
  override val json = JsonUtils.partitionOffsets(partitionToOffsets)
}

override def getOffset: Option[Offset] = {
    // 获取 init offsets
    initialPartitionOffsets
    
    // 获取 lastest offsets
    val latest = kafkaReader.fetchLatestOffsets()
    // maxOffsetsPerTrigger 为每次 trigger 拉取的 offset 数
    val offsets = maxOffsetsPerTrigger match {
      case None =>
        // 若为指定，则拉取到 lastest
        latest
      // currentPartitionOffsets 上一次消费到的 offsets
      case Some(limit) if currentPartitionOffsets.isEmpty =>
        rateLimit(limit, initialPartitionOffsets, latest)
      case Some(limit) =>
        rateLimit(limit, currentPartitionOffsets.get, latest)
    }

    currentPartitionOffsets = Some(offsets)
    logDebug(s"GetOffset: ${offsets.toSeq.map(_.toString).sorted}")
    Some(KafkaSourceOffset(offsets))
  }

private lazy val initialPartitionOffsets = {
  val metadataLog =
    // metadataPath 为该 KafkaSource 对应的 meta 持久化到 hdfs 的地址
    new HDFSMetadataLog[KafkaSourceOffset](sqlContext.sparkSession, metadataPath) {
      // 将 KafkaSourceOffset 写到 hdfs 
      override def serialize(metadata: KafkaSourceOffset, out: OutputStream): Unit = {
        out.write(0) // A zero byte is written to support Spark 2.1.0 (SPARK-19517)
        val writer = new BufferedWriter(new OutputStreamWriter(out, StandardCharsets.UTF_8))
        writer.write("v" + VERSION + "\n")
        writer.write(metadata.json)
        writer.flush
      }

      // 从 hdfs 中读出 KafkaSourceOffset
      override def deserialize(in: InputStream): KafkaSourceOffset = {
        in.read() // A zero byte is read to support Spark 2.1.0 (SPARK-19517)
        val content = IOUtils.toString(new InputStreamReader(in, StandardCharsets.UTF_8))
        // HDFSMetadataLog guarantees that it never creates a partial file.
        assert(content.length != 0)
        if (content(0) == 'v') {
          val indexOfNewLine = content.indexOf("\n")
          if (indexOfNewLine > 0) {
            val version = parseVersion(content.substring(0, indexOfNewLine), VERSION)
            KafkaSourceOffset(SerializedOffset(content.substring(indexOfNewLine + 1)))
          } else {
            throw new IllegalStateException(
              s"Log file was malformed: failed to detect the log file version line.")
          }
        } else {
          // The log was generated by Spark 2.1.0
          KafkaSourceOffset(SerializedOffset(content))
        }
      }
    }

  // EarliestOffsetRangeLimit: 希望返回最小的 offset
  // LatestOffsetRangeLimit: 希望返回最大的 offset
  // SpecificOffsetRangeLimit: 希望返回指定的 offset
  metadataLog.get(0).getOrElse {
    val offsets = startingOffsets match {
      case EarliestOffsetRangeLimit => KafkaSourceOffset(kafkaReader.fetchEarliestOffsets())
      case LatestOffsetRangeLimit => KafkaSourceOffset(kafkaReader.fetchLatestOffsets())
      case SpecificOffsetRangeLimit(p) => kafkaReader.fetchSpecificOffsets(p, reportDataLoss)
    }
    metadataLog.add(0, offsets)
    logInfo(s"Initial offsets: $offsets")
    offsets
  }.partitionToOffsets
}

private def rateLimit(
      limit: Long,
      from: Map[TopicPartition, Long],
      until: Map[TopicPartition, Long]): Map[TopicPartition, Long]

def getBatch(start: Option[Offset], end: Offset): DataFrame

override def stop(): Unit = synchronized {
  kafkaReader.close()
}