try (StateSnapshotContextSynchronousImpl snapshotContext = new StateSnapshotContextSynchronousImpl( checkpointId, timestamp, snapshotInProgress.setKeyedStateRawFuture(snapshotContext.getKeyedStateStreamFuture()); snapshotInProgress.setOperatorStateRawFuture(snapshotContext.getOperatorStateStreamFuture());
StateSnapshotContextSynchronousImpl context = new StateSnapshotContextSynchronousImpl( checkpointId, checkpointTimestamp, context.getRawKeyedOperatorStateOutput(); context.getRawOperatorStateOutput(); assertTrue(closableRegistry.contains(outputStream2)); context.close();
/** * Checks that the state snapshot context is closed after a successful snapshot operation. */ @Test public void testSnapshotMethod() throws Exception { final long checkpointId = 42L; final long timestamp = 1L; final CloseableRegistry closeableRegistry = new CloseableRegistry(); StateSnapshotContextSynchronousImpl context = spy(new StateSnapshotContextSynchronousImpl(0L, 0L)); whenNew(StateSnapshotContextSynchronousImpl.class).withAnyArguments().thenReturn(context); StreamTask<Void, AbstractStreamOperator<Void>> containingTask = mock(StreamTask.class); when(containingTask.getCancelables()).thenReturn(closeableRegistry); AbstractStreamOperator<Void> operator = mock(AbstractStreamOperator.class); when(operator.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class), any(CheckpointStreamFactory.class))).thenCallRealMethod(); doReturn(containingTask).when(operator).getContainingTask(); operator.snapshotState( checkpointId, timestamp, CheckpointOptions.forCheckpointWithDefaultLocation(), new MemCheckpointStreamFactory(Integer.MAX_VALUE)); verify(context).close(); }
RunnableFuture<SnapshotResult<OperatorStateHandle>> futureOperatorStateHandle = mock(RunnableFuture.class); StateSnapshotContextSynchronousImpl context = spy(new StateSnapshotContextSynchronousImpl(checkpointId, timestamp)); when(context.getKeyedStateStreamFuture()).thenReturn(futureKeyedStateHandle); when(context.getOperatorStateStreamFuture()).thenReturn(futureOperatorStateHandle); verify(context).close(); verify(operatorSnapshotResult).cancel();
@Nonnull public RunnableFuture<SnapshotResult<KeyedStateHandle>> getKeyedStateStreamFuture() throws IOException { KeyedStateHandle keyGroupsStateHandle = closeAndUnregisterStreamToObtainStateHandle(keyedStateCheckpointOutputStream); return toDoneFutureOfSnapshotResult(keyGroupsStateHandle); }
@Override public OperatorStateCheckpointOutputStream getRawOperatorStateOutput() throws Exception { if (null == operatorStateCheckpointOutputStream) { operatorStateCheckpointOutputStream = new OperatorStateCheckpointOutputStream(openAndRegisterNewStream()); } return operatorStateCheckpointOutputStream; }
@Before public void setUp() throws Exception { CloseableRegistry closableRegistry = new CloseableRegistry(); CheckpointStreamFactory streamFactory = new MemCheckpointStreamFactory(1024); KeyGroupRange keyGroupRange = new KeyGroupRange(0, 2); this.snapshotContext = new StateSnapshotContextSynchronousImpl(42, 4711, streamFactory, keyGroupRange, closableRegistry); }
@Override public void close() throws IOException { IOException exception = null; if (keyedStateCheckpointOutputStream != null) { try { closeAndUnregisterStream(keyedStateCheckpointOutputStream); } catch (IOException e) { exception = new IOException("Could not close the raw keyed state checkpoint output stream.", e); } } if (operatorStateCheckpointOutputStream != null) { try { closeAndUnregisterStream(operatorStateCheckpointOutputStream); } catch (IOException e) { exception = ExceptionUtils.firstOrSuppressed( new IOException("Could not close the raw operator state checkpoint output stream.", e), exception); } } if (exception != null) { throw exception; } } }
public RunnableFuture<KeyedStateHandle> getKeyedStateStreamFuture() throws IOException { KeyGroupsStateHandle keyGroupsStateHandle = closeAndUnregisterStreamToObtainStateHandle(keyedStateCheckpointOutputStream); return new DoneFuture<KeyedStateHandle>(keyGroupsStateHandle); }
@Nonnull public RunnableFuture<SnapshotResult<OperatorStateHandle>> getOperatorStateStreamFuture() throws IOException { OperatorStateHandle operatorStateHandle = closeAndUnregisterStreamToObtainStateHandle(operatorStateCheckpointOutputStream); return toDoneFutureOfSnapshotResult(operatorStateHandle); }
@Override public OperatorStateCheckpointOutputStream getRawOperatorStateOutput() throws Exception { if (null == operatorStateCheckpointOutputStream) { operatorStateCheckpointOutputStream = new OperatorStateCheckpointOutputStream(openAndRegisterNewStream()); } return operatorStateCheckpointOutputStream; }
/** * Tests that when taking a checkpoint when the fetcher is not running yet, * the checkpoint correctly contains the restored state instead. */ @Test public void checkRestoredCheckpointWhenFetcherNotReady() throws Exception { @SuppressWarnings("unchecked") final FlinkKafkaConsumerBase<String> consumer = new DummyFlinkKafkaConsumer<>(); final TestingListState<Tuple2<KafkaTopicPartition, Long>> restoredListState = new TestingListState<>(); setupConsumer(consumer, true, restoredListState, true, 0, 1); // snapshot before the fetcher starts running consumer.snapshotState(new StateSnapshotContextSynchronousImpl(17, 17)); // ensure that the list was cleared and refilled. while this is an implementation detail, we use it here // to figure out that snapshotState() actually did something. Assert.assertTrue(restoredListState.isClearCalled()); Set<Serializable> expected = new HashSet<>(); for (Serializable serializable : restoredListState.get()) { expected.add(serializable); } int counter = 0; for (Serializable serializable : restoredListState.get()) { assertTrue(expected.contains(serializable)); counter++; } assertEquals(expected.size(), counter); }
@Override public void close() throws IOException { IOException exception = null; if (keyedStateCheckpointOutputStream != null) { try { closeAndUnregisterStream(keyedStateCheckpointOutputStream); } catch (IOException e) { exception = new IOException("Could not close the raw keyed state checkpoint output stream.", e); } } if (operatorStateCheckpointOutputStream != null) { try { closeAndUnregisterStream(operatorStateCheckpointOutputStream); } catch (IOException e) { exception = ExceptionUtils.firstOrSuppressed( new IOException("Could not close the raw operator state checkpoint output stream.", e), exception); } } if (exception != null) { throw exception; } } }
public RunnableFuture<OperatorStateHandle> getOperatorStateStreamFuture() throws IOException { OperatorStateHandle operatorStateHandle = closeAndUnregisterStreamToObtainStateHandle(operatorStateCheckpointOutputStream); return new DoneFuture<>(operatorStateHandle); }
try (StateSnapshotContextSynchronousImpl snapshotContext = new StateSnapshotContextSynchronousImpl( checkpointId, timestamp, snapshotInProgress.setKeyedStateRawFuture(snapshotContext.getKeyedStateStreamFuture()); snapshotInProgress.setOperatorStateRawFuture(snapshotContext.getOperatorStateStreamFuture());
@Nonnull public RunnableFuture<SnapshotResult<KeyedStateHandle>> getKeyedStateStreamFuture() throws IOException { KeyedStateHandle keyGroupsStateHandle = closeAndUnregisterStreamToObtainStateHandle(keyedStateCheckpointOutputStream); return toDoneFutureOfSnapshotResult(keyGroupsStateHandle); }
@Override public OperatorStateCheckpointOutputStream getRawOperatorStateOutput() throws Exception { if (null == operatorStateCheckpointOutputStream) { operatorStateCheckpointOutputStream = new OperatorStateCheckpointOutputStream(openAndRegisterNewStream()); } return operatorStateCheckpointOutputStream; }
/** * Tests that no checkpoints happen when the fetcher is not running. */ @Test public void ignoreCheckpointWhenNotRunning() throws Exception { @SuppressWarnings("unchecked") final MockFetcher<String> fetcher = new MockFetcher<>(); final FlinkKafkaConsumerBase<String> consumer = new DummyFlinkKafkaConsumer<>( fetcher, mock(AbstractPartitionDiscoverer.class), false); final TestingListState<Tuple2<KafkaTopicPartition, Long>> listState = new TestingListState<>(); setupConsumer(consumer, false, listState, true, 0, 1); // snapshot before the fetcher starts running consumer.snapshotState(new StateSnapshotContextSynchronousImpl(1, 1)); // no state should have been checkpointed assertFalse(listState.get().iterator().hasNext()); // acknowledgement of the checkpoint should also not result in any offset commits consumer.notifyCheckpointComplete(1L); assertNull(fetcher.getAndClearLastCommittedOffsets()); assertEquals(0, fetcher.getCommitCount()); }
@Override public void close() throws IOException { IOException exception = null; if (keyedStateCheckpointOutputStream != null) { try { closeAndUnregisterStream(keyedStateCheckpointOutputStream); } catch (IOException e) { exception = new IOException("Could not close the raw keyed state checkpoint output stream.", e); } } if (operatorStateCheckpointOutputStream != null) { try { closeAndUnregisterStream(operatorStateCheckpointOutputStream); } catch (IOException e) { exception = ExceptionUtils.firstOrSuppressed( new IOException("Could not close the raw operator state checkpoint output stream.", e), exception); } } if (exception != null) { throw exception; } } }
try (StateSnapshotContextSynchronousImpl snapshotContext = new StateSnapshotContextSynchronousImpl( checkpointId, timestamp, snapshotInProgress.setKeyedStateRawFuture(snapshotContext.getKeyedStateStreamFuture()); snapshotInProgress.setOperatorStateRawFuture(snapshotContext.getOperatorStateStreamFuture());