Spark Structured Streaming中MemoryStream机制及测试问题排查

算力阿峰

2026-7-18

问题描述

我正在为一款基于Spark Structured Streaming将输入数据写入HBase表的应用编写测试。思路是Mock一个HBaseClient实例来创建BufferedMutator，其中TestingBufferedMutator不执行实际插入操作，仅将接收到的append操作保存到列表中，供测试后续验证。但目前代码执行从未进入forEachBatch{}逻辑块，推测是MemoryStream的输入未被解析，导致触发触发器时没有可处理的流数据。我已确认JSON字符串解析正常（println()语句能正确输出内容），也曾尝试重写空的afterSessionStart()方法，避免在流处理流程启动前添加数据。

我参考了相关指南实现了修改版代码，但问题仍未解决。想了解MemoryStream的具体工作机制，以及我在向流写入JSON数据时的错误操作，恳请提供帮助！

测试代码

class StreamingLoaderOutputSpec extends FunSpec with MockFactory with BeforeAndAfter {

  private val mockHBaseClient = mock[HBaseClient]
  private val testingBufferedMutator = new TestingBufferedMutator()

  private val streamingConfig = StreamingConfig(
    hbaseTableName = "sampleTable",
    sourceKafkaTopics = List(),
    checkpointDirectory="checkpointDir",
    batchInterval = 1,
    batchIntervalTimeUnit = TimeUnit.SECONDS,
    kerberosPrincipal = "kerberosPrincipal",
    kerberosKeyTabFile = "kerberosKeyTabFile",
    dryRun = false
  )
  private var jsonString: String = ""

  before {

    val fileSrc = scala.io.Source.fromFile("src/test/resources/sample_kafka_msg.json")
    jsonString = fileSrc.mkString
    fileSrc.close()

    (mockHBaseClient.getSerializableBufferedMutator _).expects(*).onCall((tableName: String) => {
      testingBufferedMutator
    }).once()
  }

  describe("StreamingLoader") {

    it("should create appends with correct content"){

      val spark: SparkSession = SparkSession.builder()
        .appName("StreamingLoader Test")
        .master("local[1]")
        .getOrCreate()

      implicit val sqlCtx: SQLContext = spark.sqlContext
      import spark.implicits._

      val events: MemoryStream[String] = MemoryStream[String]
      val sessions = events.toDS
      val sessionsDF = sessions.toDF

      assert(sessions.isStreaming, "sessions must be a streaming Dataset")

      println(s"Test JSON content: $jsonString")

      val cb = new CyclicBarrier(1, new Runnable {
        override def run(): Unit = {
          println("Adding dataSession has been started")
          val currentOffset = events.addData(jsonString)
          events.commit(currentOffset.asInstanceOf[LongOffset])
        }
      })

      val loader = new StreamingLoader(
        spark = spark, hbaseClient = mockHBaseClient, streamedDataFrame = sessionsDF, config = streamingConfig) {
        override def afterSessionStart(): Unit = {
          cb.await()
        }
      }

      assert(testingBufferedMutator.getTestList.nonEmpty)

      testingBufferedMutator.getTestList.foreach(println(_))

    }

  }

}

待测试应用代码

class StreamingLoader(spark: SparkSession,
                      hbaseClient: HBaseClient,
                      streamedDataFrame: DataFrame,
                      config: StreamingConfig) extends Serializable {

  def execute(): Unit = {

    val labelSerializer = new LabelSerializer()

    val bufferedMutator = hbaseClient.getSerializableBufferedMutator(config.hbaseTableName)

    try {
      val query = streamedDataFrame
        .writeStream
        .outputMode(OutputMode.Append())
        .option("checkpointLocation", config.checkpointDirectory)
        .trigger(Trigger.ProcessingTime(config.batchInterval, config.batchIntervalTimeUnit))
        .foreachBatch { (receivedDataFrame: DataFrame, _: Long) =>

          val dataToWrite = Mapper(receivedDataFrame)

          dataToWrite.foreach(row => {

            val result = labelSerializer.process(row)

            if (result.isDefined) {
              val (rowKey, columnNameBytes, payload) = result.get

              // Append instead of Put HBaseClient Object
              val append = new Append(rowKey)
              append.addColumn(TargetColumnFamily, columnNameBytes, payload)

              try {
                bufferedMutator.mutate(append)
              } catch {
                case exception: Exception =>
                  // Swallow the exception and resume processing
                  logger.error("BufferedMutator.mutate raised Exception", exception)
                case throwable: Throwable =>
                  // Throwables are re-thrown, causing the application to fail
                  logger.error("BufferedMutator.mutate failed with Throwable", throwable)
                  throw throwable
              }
            }
          })

          try {
            bufferedMutator.flush()
          } catch {
            case exception: Exception =>
              // Swallow the exception and resume processing
              logger.error("BufferedMutator.flush raised Exception", exception)
            case throwable: Throwable =>
              // Throwables are re-thrown, causing the application to fail
              logger.error("BufferedMutator.flush failed with Throwable", throwable)
              throw throwable
          }
        }
        .start

      logger.info("Session has been started")
      afterSessionStart()

      query.awaitTermination()
      closeQuietly(bufferedMutator)

    } catch {
      case throwable: Throwable =>
        closeQuietly(bufferedMutator)

        throw throwable
    }
  }

  def afterSessionStart(): Unit = {

  }

  private def closeQuietly(resource: Closeable): Unit = {
    try {
      resource.close()
    } catch {
      case exception: Throwable => exception.printStackTrace()
    }
  }
}

编辑/更新

我发现代码中缺少processAllAvailable()调用，于是修改了afterSessionStart()方法，使其接收StreamingQuery参数，并将awaitTermination()语句移入其中：

def afterSessionStart(query: StreamingQuery): Unit = {
  query.awaitTermination()
}

同时我也修改了测试代码，补充了该方法的调用：

val loader = new StreamingLoader(
    spark = spark, hbaseClient = mockHBaseClient, streamedDataFrame = sessionsDF, config = streamingConfig) {
    override def afterSessionStart(query: StreamingQuery): Unit = {
        println("Adding dataSession has been started")
        val currentOffset = memoryStream.addData(jsonString)
        query.processAllAvailable()
        memoryStream.commit(currentOffset.asInstanceOf[LongOffset])
    }
}

遗憾的是问题仍未解决，但这种重构对于编写仅测试转换流类的测试文件至关重要。我还参考了相关线程来尝试修复代码。

解决方案

1. MemoryStream核心工作机制说明

MemoryStream是Spark Structured Streaming提供的测试专用内存数据源，核心逻辑：

addData()仅将数据写入内存缓冲区，不会自动触发流处理
必须调用commit()提交偏移量，数据才会被标记为待处理批次
流查询需通过processAllAvailable()显式触发处理，否则会一直等待触发器时间或外部信号

2. 代码核心问题分析

原代码问题点：

未启动流处理：仅创建StreamingLoader实例，但未调用loader.execute()，流处理流程完全没启动，自然不会进入forEachBatch
CyclicBarrier无意义：屏障数设为1无法实现“启动流后再添加数据”的同步逻辑
断言时机错误：流处理未执行就断言testingBufferedMutator.getTestList非空，此时数据还未被处理

更新后仍存在的问题：

processAllAvailable()顺序错误：先触发处理再添加数据，此时缓冲区为空，无数据可处理
awaitTermination()阻塞流程：该方法会让线程无限阻塞，无法执行后续断言
仍可能遗漏启动流处理：未确保loader.execute()被调用

3. 修复后的测试代码示例

class StreamingLoaderOutputSpec extends FunSpec with MockFactory with BeforeAndAfter {

  private val mockHBaseClient = mock[HBaseClient]
  private val testingBufferedMutator = new TestingBufferedMutator()

  private val streamingConfig = StreamingConfig(
    hbaseTableName = "sampleTable",
    sourceKafkaTopics = List(),
    checkpointDirectory="checkpointDir",
    batchInterval = 1,
    batchIntervalTimeUnit = TimeUnit.SECONDS,
    kerberosPrincipal = "kerberosPrincipal",
    kerberosKeyTabFile = "kerberosKeyTabFile",
    dryRun = false
  )
  private var jsonString: String = ""

  before {
    val fileSrc = scala.io.Source.fromFile("src/test/resources/sample_kafka_msg.json")
    jsonString = fileSrc.mkString
    fileSrc.close()

    (mockHBaseClient.getSerializableBufferedMutator _).expects(*).returns(testingBufferedMutator).once()
  }

  describe("StreamingLoader") {
    it("should create appends with correct content") {
      val spark: SparkSession = SparkSession.builder()
        .appName("StreamingLoader Test")
        .master("local[1]")
        .getOrCreate()
      import spark.implicits._

      val events: MemoryStream[String] = MemoryStream[String]
      val sessionsDF = events.toDS.toDF
      assert(sessionsDF.isStreaming, "sessionsDF must be a streaming DataFrame")

      // 启动流处理并同步数据添加逻辑
      val loader = new StreamingLoader(
        spark = spark, hbaseClient = mockHBaseClient, streamedDataFrame = sessionsDF, config = streamingConfig) {
        override def afterSessionStart(query: StreamingQuery): Unit = {
          // 先添加数据并提交偏移量
          val offset = events.addData(jsonString)
          events.commit(offset.asInstanceOf[LongOffset])
          // 触发所有可用数据的处理
          query.processAllAvailable()
          // 停止查询避免阻塞
          query.stop()
        }
      }

      // 在新线程中启动流处理，避免主线程被awaitTermination阻塞
      new Thread(() => loader.execute()).start()
      // 等待流处理完成（可根据环境调整时间）
      Thread.sleep(2000)

      // 执行断言
      assert(testingBufferedMutator.getTestList.nonEmpty, "No append operations were captured")
      testingBufferedMutator.getTestList.foreach(println(_))

      // 关闭Spark会话
      spark.stop()
    }
  }
}

4. 待测试代码的配套调整

修改StreamingLoader的afterSessionStart()方法，传入StreamingQuery实例：

class StreamingLoader(spark: SparkSession,
                      hbaseClient: HBaseClient,
                      streamedDataFrame: DataFrame,
                      config: StreamingConfig) extends Serializable {

  def execute(): Unit = {
    val labelSerializer = new LabelSerializer()
    val bufferedMutator = hbaseClient.getSerializableBufferedMutator(config.hbaseTableName)

    try {
      val query = streamedDataFrame
        .writeStream
        .outputMode(OutputMode.Append())
        .option("checkpointLocation", config.checkpointDirectory)
        .trigger(Trigger.ProcessingTime(config.batchInterval, config.batchIntervalTimeUnit))
        .foreachBatch { (receivedDataFrame: DataFrame, _: Long) =>
          // 原有逻辑保持不变
        }
        .start

      logger.info("Session has been started")
      afterSessionStart(query) // 传入query实例
      query.awaitTermination()
      closeQuietly(bufferedMutator)
    } catch {
      case throwable: Throwable =>
        closeQuietly(bufferedMutator)
        throw throwable
    }
  }

  // 更新方法签名，接收StreamingQuery参数
  def afterSessionStart(query: StreamingQuery): Unit = {}

  private def closeQuietly(resource: Closeable): Unit = {
    try {
      resource.close()
    } catch {
      case exception: Throwable => exception.printStackTrace()
    }
  }
}

5. 关键注意事项

必须启动流处理：测试中务必调用loader.execute()，且因awaitTermination()会阻塞主线程，需在新线程中执行
数据处理顺序：先调用addData()和commit()，再触发processAllAvailable()，确保数据已进入待处理队列
终止流查询：处理完成后调用query.stop()，避免测试无限阻塞
调整等待时间：根据测试环境性能，适当延长Thread.sleep()的时间，确保流处理完全执行

内容的提问来源于stack exchange，提问作者Peter Ivony

Spark Structured Streaming中MemoryStream机制及测试问题排查

测试代码 #

待测试应用代码 #

编辑/更新 #

1. MemoryStream核心工作机制说明 #

2. 代码核心问题分析 #