Flink - State 之读取 Kafka Offsets

  class MockKafkaSource extends RichParallelSourceFunction[String] with CheckpointedFunction {
    var unionOffsetStates: ListState[Tuple2[KafkaTopicPartition, java.lang.Long]] = _
    override def run(sourceContext: SourceFunction.SourceContext[String]): Unit = {
      Thread.sleep(300 * 1000)
    }

    override def cancel(): Unit = {}

    override def initializeState(functionInitializationContext: FunctionInitializationContext): Unit = {
      val stateStore = functionInitializationContext.getOperatorStateStore
      unionOffsetStates = stateStore.getUnionListState(new ListStateDescriptor(
        "topic-partition-offset-states",
        createTypeInformation[Tuple2[KafkaTopicPartition, java.lang.Long]]
      ))

      for (unionOffsetState <- unionOffsetStates.get()) {
        println(s"f0:${unionOffsetState.f0} f1:${unionOffsetState.f1}")
        println(s"stackTrace:\n${Thread.currentThread.getStackTrace.mkString("\n")}")
      }
    }

    override def snapshotState(functionSnapshotContext: FunctionSnapshotContext): Unit = {
    }
  }

object StateViewer extends App {
  val params = ParameterTool.fromArgs(args)
  val env = ExecutionEnvironment.getExecutionEnvironment
  val state = SavepointLoader.loadSavepoint(params.get("ckdir"))

  println(s"version:${state.getVersion} checkpointId:${state.getCheckpointId}")
  state.getOperatorStates.asScala.foreach{operatorState =>
    println(s"\toperatorState:${operatorState}")
    operatorState.getSubtaskStates.asScala.foreach{ case(subTaskIndex, subTaskState) =>
      println(s"\t\tsubTaskIndex:${subTaskIndex}")

      parseManagedOpState(subTaskState.getManagedOperatorState)
      parseManagedKeyedState(subTaskState.getManagedKeyedState)
    }
  }

  // unimplemented
  def parseManagedKeyedState(managedKeyedState: StateObjectCollection[KeyedStateHandle]) = {
    managedKeyedState.asScala.foreach{keyedStateHandle =>
      val keyGroupsStateHandle = keyedStateHandle.asInstanceOf[KeyGroupsStateHandle]
      val in = keyGroupsStateHandle.openInputStream()

      val inView = new DataInputViewStreamWrapper(in)
      val serializationProxy = new KeyedBackendSerializationProxy[Any](Thread.currentThread().getContextClassLoader)
//      val serializationProxy = new KeyedBackendSerializationProxy[_](Thread.currentThread().getContextClassLoader)
      serializationProxy.read(inView)

      serializationProxy.getStateMetaInfoSnapshots
    }
  }

  // copy from OperatorStateRestoreOperation.restore
  def parseManagedOpState(managedOperatorState: StateObjectCollection[OperatorStateHandle]) = {
    managedOperatorState.asScala.foreach{operatorStateHandle =>
      val serializationProxy = new OperatorBackendSerializationProxy(Thread.currentThread().getContextClassLoader)
      serializationProxy.read(
        new DataInputViewStreamWrapper(operatorStateHandle.openInputStream())
      )

      val nameToSerializer = serializationProxy.getOperatorStateMetaInfoSnapshots().asScala.map{stateMetaInfoSnapshot =>
        println(s"\t\t\t\tname:${stateMetaInfoSnapshot.getName} type:${stateMetaInfoSnapshot.getBackendStateType}")
        //          val restoredMetaInfo = new RegisteredOperatorStateBackendMetaInfo[_](stateMetaInfoSnapshot)
        val restoredMetaInfo = new RegisteredOperatorStateBackendMetaInfo[Any](stateMetaInfoSnapshot)
        restoredMetaInfo.getName -> restoredMetaInfo.getPartitionStateSerializer
      }.toMap

      val in = operatorStateHandle.openInputStream()
      operatorStateHandle.getStateNameToPartitionOffsets.asScala.foreach{case (name, metaInfo) =>
        val inView = new DataInputViewStreamWrapper(in)
        val serializer = nameToSerializer.get(name).get
        val offsets = metaInfo.getOffsets

        println(s"\t\t\t\tname:${name} serializer:${serializer} offsets:${offsets.mkString("[", ",", "]")}")
        offsets.foreach{offset =>
          in.seek(offset)
          println(s"\t\t\t\t\toffset:${offset} value:${serializer.deserialize(inView)}")
        }
      }
    }
  }
}