org.apache.flink.api.common.state.OperatorStateStore Java Examples

private static <T> SourceOperator<T, MockSourceSplit> createTestOperator(
		SourceReader<T, MockSourceSplit> reader,
		WatermarkStrategy<T> watermarkStrategy,
		ProcessingTimeService timeService) throws Exception {

	final OperatorStateStore operatorStateStore =
			new MemoryStateBackend().createOperatorStateBackend(
					new MockEnvironmentBuilder().build(),
					"test-operator",
					Collections.emptyList(),
					new CloseableRegistry());

	final StateInitializationContext stateContext = new StateInitializationContextImpl(
		false, operatorStateStore, null, null, null);

	final SourceOperator<T, MockSourceSplit> sourceOperator =
			new TestingSourceOperator<>(reader, watermarkStrategy, timeService);
	sourceOperator.initializeState(stateContext);
	sourceOperator.open();

	return sourceOperator;
}
 

private TestPulsarConsumerSource createSource(Consumer<byte[]> testConsumer,
                                              long batchSize, boolean isCheckpointingEnabled) throws Exception {
    PulsarSourceBuilder<String> builder =
        PulsarSourceBuilder.builder(new SimpleStringSchema())
            .acknowledgementBatchSize(batchSize);
    TestPulsarConsumerSource source = new TestPulsarConsumerSource(builder, testConsumer, isCheckpointingEnabled);

    OperatorStateStore mockStore = mock(OperatorStateStore.class);
    FunctionInitializationContext mockContext = mock(FunctionInitializationContext.class);
    when(mockContext.getOperatorStateStore()).thenReturn(mockStore);
    when(mockStore.getSerializableListState(any(String.class))).thenReturn(null);

    source.initializeState(mockContext);

    return source;
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
  if (checkpointCoder == null) {
    // no checkpoint coder available in this source
    return;
  }

  OperatorStateStore stateStore = context.getOperatorStateStore();
  @SuppressWarnings("unchecked")
  CoderTypeInformation<KV<? extends UnboundedSource<OutputT, CheckpointMarkT>, CheckpointMarkT>>
      typeInformation = (CoderTypeInformation) new CoderTypeInformation<>(checkpointCoder);
  stateForCheckpoint =
      stateStore.getOperatorState(
          new ListStateDescriptor<>(
              DefaultOperatorStateBackend.DEFAULT_OPERATOR_STATE_NAME,
              typeInformation.createSerializer(new ExecutionConfig())));

  if (context.isRestored()) {
    isRestored = true;
    LOG.info("Restoring state in the UnboundedSourceWrapper.");
  } else {
    LOG.info("No restore state for UnboundedSourceWrapper.");
  }
}
 

@Override
public final void initializeState(FunctionInitializationContext context) throws Exception {

	OperatorStateStore stateStore = context.getOperatorStateStore();

	this.unionOffsetStates = stateStore.getUnionListState(new ListStateDescriptor<>(OFFSETS_STATE_NAME,
		createStateSerializer(getRuntimeContext().getExecutionConfig())));

	if (context.isRestored()) {
		restoredState = new TreeMap<>(new KafkaTopicPartition.Comparator());

		// populate actual holder for restored state
		for (Tuple2<KafkaTopicPartition, Long> kafkaOffset : unionOffsetStates.get()) {
			restoredState.put(kafkaOffset.f0, kafkaOffset.f1);
		}

		LOG.info("Consumer subtask {} restored state: {}.", getRuntimeContext().getIndexOfThisSubtask(), restoredState);
	} else {
		LOG.info("Consumer subtask {} has no restore state.", getRuntimeContext().getIndexOfThisSubtask());
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {

	readyToFail = false;

	if (context.isRestored()) {
		isRunning = false;
	} else {
		isRunning = true;

		OperatorStateStore operatorStateStore = context.getOperatorStateStore();
		for (int i = 0; i < numListStates; ++i) {

			ListStateDescriptor<String> listStateDescriptor =
				new ListStateDescriptor<>("test-list-state-" + i, String.class);

			ListState<String> unionListState =
				operatorStateStore.getUnionListState(listStateDescriptor);

			for (int j = 0; j < numPartitionsPerListState; ++j) {
				unionListState.add(String.valueOf(j));
			}
		}
	}
}
 

@Before
public void beforeTest() throws Exception {

	OperatorStateStore mockStore = Mockito.mock(OperatorStateStore.class);
	FunctionInitializationContext mockContext = Mockito.mock(FunctionInitializationContext.class);
	Mockito.when(mockContext.getOperatorStateStore()).thenReturn(mockStore);
	Mockito.when(mockStore.getSerializableListState(any(String.class))).thenReturn(null);

	source = new RMQTestSource();
	source.initializeState(mockContext);
	source.open(config);

	messageId = 0;
	generateCorrelationIds = true;

	sourceThread = new Thread(new Runnable() {
		@Override
		public void run() {
			try {
				source.run(new DummySourceContext());
			} catch (Exception e) {
				exception = e;
			}
		}
	});
}
 

static PartitionCommitTrigger create(
		boolean isRestored,
		OperatorStateStore stateStore,
		Configuration conf,
		ClassLoader cl,
		List<String> partitionKeys,
		ProcessingTimeService procTimeService) throws Exception {
	String trigger = conf.get(SINK_PARTITION_COMMIT_TRIGGER);
	switch (trigger) {
		case PARTITION_TIME:
			return new PartitionTimeCommitTigger(
					isRestored, stateStore, conf, cl, partitionKeys);
		case PROCESS_TIME:
			return new ProcTimeCommitTigger(
					isRestored, stateStore, conf, procTimeService);
		default:
			throw new UnsupportedOperationException(
					"Unsupported partition commit trigger: " + trigger);
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
    OperatorStateStore stateStore = context.getOperatorStateStore();

    unionOffsetStates = stateStore.getUnionListState(
            new ListStateDescriptor<>(
                    OFFSETS_STATE_NAME,
                    TypeInformation.of(new TypeHint<Tuple2<String, MessageId>>() {
                    })));

    if (context.isRestored()) {
        restoredState = new TreeMap<>();
        unionOffsetStates.get().forEach(e -> restoredState.put(e.f0, e.f1));
        log.info("Source subtask {} restored state {}",
                taskIndex,
                StringUtils.join(restoredState.entrySet()));
    } else {
        log.info("Source subtask {} has no restore state", taskIndex);
    }
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {

	readyToFail = false;

	if (context.isRestored()) {
		isRunning = false;
	} else {
		isRunning = true;

		OperatorStateStore operatorStateStore = context.getOperatorStateStore();
		for (int i = 0; i < numListStates; ++i) {

			ListStateDescriptor<String> listStateDescriptor =
				new ListStateDescriptor<>("test-list-state-" + i, String.class);

			ListState<String> unionListState =
				operatorStateStore.getUnionListState(listStateDescriptor);

			for (int j = 0; j < numPartitionsPerListState; ++j) {
				unionListState.add(String.valueOf(j));
			}
		}
	}
}
 

public PartitionTimeCommitTigger(
		boolean isRestored,
		OperatorStateStore stateStore,
		Configuration conf,
		ClassLoader cl,
		List<String> partitionKeys) throws Exception {
	this.pendingPartitionsState = stateStore.getListState(PENDING_PARTITIONS_STATE_DESC);
	this.pendingPartitions = new HashSet<>();
	if (isRestored) {
		pendingPartitions.addAll(pendingPartitionsState.get().iterator().next());
	}

	this.partitionKeys = partitionKeys;
	this.commitDelay = conf.get(SINK_PARTITION_COMMIT_DELAY).toMillis();
	this.extractor = PartitionTimeExtractor.create(
			cl,
			conf.get(PARTITION_TIME_EXTRACTOR_KIND),
			conf.get(PARTITION_TIME_EXTRACTOR_CLASS),
			conf.get(PARTITION_TIME_EXTRACTOR_TIMESTAMP_PATTERN));

	this.watermarksState = stateStore.getListState(WATERMARKS_STATE_DESC);
	this.watermarks = new TreeMap<>();
	if (isRestored) {
		watermarks.putAll(watermarksState.get().iterator().next());
	}
}
 

public StreamingFileSinkHelper(
		Buckets<IN, ?> buckets,
		boolean isRestored,
		OperatorStateStore stateStore,
		ProcessingTimeService procTimeService,
		long bucketCheckInterval) throws Exception {
	this.bucketCheckInterval = bucketCheckInterval;
	this.buckets = buckets;
	this.bucketStates = stateStore.getListState(BUCKET_STATE_DESC);
	this.maxPartCountersState = stateStore.getUnionListState(MAX_PART_COUNTER_STATE_DESC);
	this.procTimeService = procTimeService;

	if (isRestored) {
		buckets.initializeState(bucketStates, maxPartCountersState);
	}

	long currentProcessingTime = procTimeService.getCurrentProcessingTime();
	procTimeService.registerTimer(currentProcessingTime + bucketCheckInterval, this);
}
 

@Before
public void beforeTest() throws Exception {

	OperatorStateStore mockStore = Mockito.mock(OperatorStateStore.class);
	FunctionInitializationContext mockContext = Mockito.mock(FunctionInitializationContext.class);
	Mockito.when(mockContext.getOperatorStateStore()).thenReturn(mockStore);
	Mockito.when(mockStore.getSerializableListState(any(String.class))).thenReturn(null);

	source = new RMQTestSource();
	source.initializeState(mockContext);
	source.open(config);

	messageId = 0;
	generateCorrelationIds = true;

	sourceThread = new Thread(new Runnable() {
		@Override
		public void run() {
			try {
				source.run(new DummySourceContext());
			} catch (Exception e) {
				exception = e;
			}
		}
	});
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {

	readyToFail = false;

	if (context.isRestored()) {
		isRunning = false;
	} else {
		isRunning = true;

		OperatorStateStore operatorStateStore = context.getOperatorStateStore();
		for (int i = 0; i < numListStates; ++i) {

			ListStateDescriptor<String> listStateDescriptor =
				new ListStateDescriptor<>("test-list-state-" + i, String.class);

			ListState<String> unionListState =
				operatorStateStore.getUnionListState(listStateDescriptor);

			for (int j = 0; j < numPartitionsPerListState; ++j) {
				unionListState.add(String.valueOf(j));
			}
		}
	}
}
 

private StateInitializationContext getStateContext() throws Exception {
	// Create a mock split.
	byte[] serializedSplitWithVersion = SimpleVersionedSerialization
			.writeVersionAndSerialize(new MockSourceSplitSerializer(), MOCK_SPLIT);

	// Crate the state context.
	OperatorStateStore operatorStateStore = createOperatorStateStore();
	StateInitializationContext stateContext = new StateInitializationContextImpl(
			false,
			operatorStateStore,
			null,
			null,
			null);

	// Update the context.
	stateContext.getOperatorStateStore()
				.getListState(SourceOperator.SPLITS_STATE_DESC)
				.update(Collections.singletonList(serializedSplitWithVersion));

	return stateContext;
}
 

private static <T> FunctionInitializationContext getInitializationContext(ListState<T> listState)
    throws Exception {
  FunctionInitializationContext mock = Mockito.mock(FunctionInitializationContext.class);
  OperatorStateStore mockOperatorState = getMockOperatorState(listState);
  when(mock.getOperatorStateStore()).thenReturn(mockOperatorState);
  return mock;
}
 

public StateInitializationContextImpl(
		boolean restored,
		OperatorStateStore operatorStateStore,
		KeyedStateStore keyedStateStore,
		Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs,
		Iterable<StatePartitionStreamProvider> rawOperatorStateInputs) {

	this.restored = restored;
	this.operatorStateStore = operatorStateStore;
	this.keyedStateStore = keyedStateStore;
	this.rawOperatorStateInputs = rawOperatorStateInputs;
	this.rawKeyedStateInputs = rawKeyedStateInputs;
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	final int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
	this.buckets = bucketsBuilder.createBuckets(subtaskIndex);

	final OperatorStateStore stateStore = context.getOperatorStateStore();
	bucketStates = stateStore.getListState(BUCKET_STATE_DESC);
	maxPartCountersState = stateStore.getUnionListState(MAX_PART_COUNTER_STATE_DESC);

	if (context.isRestored()) {
		buckets.initializeState(bucketStates, maxPartCountersState);
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	Preconditions.checkArgument(this.restoredBucketStates == null, "The operator has already been initialized.");

	try {
		initFileSystem();
	} catch (IOException e) {
		LOG.error("Error while creating FileSystem when initializing the state of the BucketingSink.", e);
		throw new RuntimeException("Error while creating FileSystem when initializing the state of the BucketingSink.", e);
	}

	if (this.refTruncate == null) {
		this.refTruncate = reflectTruncate(fs);
	}

	OperatorStateStore stateStore = context.getOperatorStateStore();
	restoredBucketStates = stateStore.getSerializableListState("bucket-states");

	int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
	if (context.isRestored()) {
		LOG.info("Restoring state for the {} (taskIdx={}).", getClass().getSimpleName(), subtaskIndex);

		for (State<T> recoveredState : restoredBucketStates.get()) {
			handleRestoredBucketState(recoveredState);
			if (LOG.isDebugEnabled()) {
				LOG.debug("{} idx {} restored {}", getClass().getSimpleName(), subtaskIndex, recoveredState);
			}
		}
	} else {
		LOG.info("No state to restore for the {} (taskIdx={}).", getClass().getSimpleName(), subtaskIndex);
	}
}
 

@Override
    public void initializeState(FunctionInitializationContext context)
            throws Exception
    {
        OperatorStateStore stateStore = context.getOperatorStateStore();
//        ListStateDescriptor<Tuple2<String, Long>> descriptor = new ListStateDescriptor<>(
//                "sink_partition_state",
//                TypeInformation.of(new TypeHint<Tuple2<String, Long>>() {}));
//        this.unionState = stateStore.getUnionListState(descriptor);
    }
 

private OperatorStateStore createOperatorStateStore() throws Exception {
	MockEnvironment env = new MockEnvironmentBuilder().build();
	final AbstractStateBackend abstractStateBackend = new MemoryStateBackend();
	CloseableRegistry cancelStreamRegistry = new CloseableRegistry();
	return abstractStateBackend.createOperatorStateBackend(
			env, "test-operator", Collections.emptyList(), cancelStreamRegistry);
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	Preconditions.checkArgument(this.restoredBucketStates == null, "The operator has already been initialized.");

	try {
		initFileSystem();
	} catch (IOException e) {
		LOG.error("Error while creating FileSystem when initializing the state of the BucketingSink.", e);
		throw new RuntimeException("Error while creating FileSystem when initializing the state of the BucketingSink.", e);
	}

	if (this.refTruncate == null) {
		this.refTruncate = reflectTruncate(fs);
	}

	// We are using JavaSerializer from the flink-runtime module here. This is very naughty and
	// we shouldn't be doing it because ideally nothing in the API modules/connector depends
	// directly on flink-runtime. We are doing it here because we need to maintain backwards
	// compatibility with old state and because we will have to rework/remove this code soon.
	OperatorStateStore stateStore = context.getOperatorStateStore();
	this.restoredBucketStates = stateStore.getListState(new ListStateDescriptor<>("bucket-states", new JavaSerializer<>()));

	int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
	if (context.isRestored()) {
		LOG.info("Restoring state for the {} (taskIdx={}).", getClass().getSimpleName(), subtaskIndex);

		for (State<T> recoveredState : restoredBucketStates.get()) {
			handleRestoredBucketState(recoveredState);
			if (LOG.isDebugEnabled()) {
				LOG.debug("{} idx {} restored {}", getClass().getSimpleName(), subtaskIndex, recoveredState);
			}
		}
	} else {
		LOG.info("No state to restore for the {} (taskIdx={}).", getClass().getSimpleName(), subtaskIndex);
	}
}
 

/**
 * Gets a mock context for initializing the source's state via {@link org.apache.flink.streaming.api.checkpoint.CheckpointedFunction#initializeState}.
 * @throws Exception
 */
FunctionInitializationContext getMockContext() throws Exception {
	OperatorStateStore mockStore = Mockito.mock(OperatorStateStore.class);
	FunctionInitializationContext mockContext = Mockito.mock(FunctionInitializationContext.class);
	Mockito.when(mockContext.getOperatorStateStore()).thenReturn(mockStore);
	Mockito.when(mockStore.getListState(any(ListStateDescriptor.class))).thenReturn(null);
	return mockContext;
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	Preconditions.checkArgument(this.restoredBucketStates == null,
		"The " + getClass().getSimpleName() + " has already been initialized.");

	try {
		initFileSystem();
	} catch (IOException e) {
		LOG.error("Error while creating FileSystem when initializing the state of the RollingSink.", e);
		throw new RuntimeException("Error while creating FileSystem when initializing the state of the RollingSink.", e);
	}

	if (this.refTruncate == null) {
		this.refTruncate = reflectTruncate(fs);
	}

	OperatorStateStore stateStore = context.getOperatorStateStore();
	restoredBucketStates = stateStore.getSerializableListState("rolling-states");

	int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
	if (context.isRestored()) {
		LOG.info("Restoring state for the {} (taskIdx={}).", getClass().getSimpleName(), subtaskIndex);

		for (BucketState bucketState : restoredBucketStates.get()) {
			handleRestoredBucketState(bucketState);
		}

		if (LOG.isDebugEnabled()) {
			LOG.debug("{} (taskIdx= {}) restored {}", getClass().getSimpleName(), subtaskIndex, bucketState);
		}
	} else {
		LOG.info("No state to restore for the {} (taskIdx= {}).", getClass().getSimpleName(), subtaskIndex);
	}
}
 

public ProcTimeCommitTigger(
		boolean isRestored,
		OperatorStateStore stateStore,
		Configuration conf,
		ProcessingTimeService procTimeService) throws Exception {
	this.pendingPartitionsState = stateStore.getListState(PENDING_PARTITIONS_STATE_DESC);
	this.pendingPartitions = new HashMap<>();
	if (isRestored) {
		pendingPartitions.putAll(pendingPartitionsState.get().iterator().next());
	}

	this.procTimeService = procTimeService;
	this.commitDelay = conf.get(SINK_PARTITION_COMMIT_DELAY).toMillis();
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	Preconditions.checkArgument(this.restoredBucketStates == null, "The operator has already been initialized.");

	try {
		initFileSystem();
	} catch (IOException e) {
		LOG.error("Error while creating FileSystem when initializing the state of the BucketingSink.", e);
		throw new RuntimeException("Error while creating FileSystem when initializing the state of the BucketingSink.", e);
	}

	if (this.refTruncate == null) {
		this.refTruncate = reflectTruncate(fs);
	}

	OperatorStateStore stateStore = context.getOperatorStateStore();
	restoredBucketStates = stateStore.getSerializableListState("bucket-states");

	int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
	if (context.isRestored()) {
		LOG.info("Restoring state for the {} (taskIdx={}).", getClass().getSimpleName(), subtaskIndex);

		for (State<T> recoveredState : restoredBucketStates.get()) {
			handleRestoredBucketState(recoveredState);
			if (LOG.isDebugEnabled()) {
				LOG.debug("{} idx {} restored {}", getClass().getSimpleName(), subtaskIndex, recoveredState);
			}
		}
	} else {
		LOG.info("No state to restore for the {} (taskIdx={}).", getClass().getSimpleName(), subtaskIndex);
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	final int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
	this.buckets = bucketsBuilder.createBuckets(subtaskIndex);

	final OperatorStateStore stateStore = context.getOperatorStateStore();
	bucketStates = stateStore.getListState(BUCKET_STATE_DESC);
	maxPartCountersState = stateStore.getUnionListState(MAX_PART_COUNTER_STATE_DESC);

	if (context.isRestored()) {
		buckets.initializeState(bucketStates, maxPartCountersState);
	}
}
 

public StateInitializationContextImpl(
		boolean restored,
		OperatorStateStore operatorStateStore,
		KeyedStateStore keyedStateStore,
		Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs,
		Iterable<StatePartitionStreamProvider> rawOperatorStateInputs) {

	this.restored = restored;
	this.operatorStateStore = operatorStateStore;
	this.keyedStateStore = keyedStateStore;
	this.rawOperatorStateInputs = rawOperatorStateInputs;
	this.rawKeyedStateInputs = rawKeyedStateInputs;
}
 

public StateInitializationContextImpl(
		boolean restored,
		OperatorStateStore operatorStateStore,
		KeyedStateStore keyedStateStore,
		Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs,
		Iterable<StatePartitionStreamProvider> rawOperatorStateInputs) {

	this.restored = restored;
	this.operatorStateStore = operatorStateStore;
	this.keyedStateStore = keyedStateStore;
	this.rawOperatorStateInputs = rawOperatorStateInputs;
	this.rawKeyedStateInputs = rawKeyedStateInputs;
}
 

@Override
public OperatorStateStore getOperatorStateStore() {
	return operatorStateStore;
}
 

private MockFunctionInitializationContext(boolean isRestored, OperatorStateStore operatorStateStore) {
	this.isRestored = isRestored;
	this.operatorStateStore = operatorStateStore;
}