@Override
    public void replace(String key, StringResourceVersion resourceVersion, T state)
            throws Exception {
        checkNotNull(key, "Key in ConfigMap.");
        checkNotNull(state, "State.");

        final RetrievableStateHandle<T> newStateHandle = storage.store(state);

        final byte[] serializedStateHandle =
                serializeOrDiscard(new StateHandleWithDeleteMarker<>(newStateHandle));

        // initialize flags to serve the failure case
        boolean discardOldState = false;
        boolean discardNewState = true;
        // We don't want to greedily pull the old state handle as we have to do that anyway in
        // replaceEntry method for check of delete markers.
        final AtomicReference<RetrievableStateHandle<T>> oldStateHandleRef =
                new AtomicReference<>();
        try {
            final boolean success =
                    updateConfigMap(
                                    cm -> {
                                        try {
                                            return replaceEntry(
                                                    cm,
                                                    key,
                                                    serializedStateHandle,
                                                    oldStateHandleRef);
                                        } catch (NotExistException e) {
                                            throw new CompletionException(e);
                                        }
                                    })
                            .get();

            // swap subject for deletion in case of success
            discardOldState = success;
            discardNewState = !success;
        } catch (Exception ex) {
            final Optional<PossibleInconsistentStateException> possibleInconsistentStateException =
                    ExceptionUtils.findThrowable(ex, PossibleInconsistentStateException.class);
            if (possibleInconsistentStateException.isPresent()) {
                // it's unclear whether the state handle metadata was written to the ConfigMap -
                // hence, we don't discard any data
                discardNewState = false;
                throw possibleInconsistentStateException.get();
            }

            throw ExceptionUtils.findThrowable(ex, NotExistException.class).orElseThrow(() -> ex);
        } finally {
            if (discardNewState) {
                newStateHandle.discardState();
            }

            if (discardOldState) {
                Objects.requireNonNull(
                                oldStateHandleRef.get(),
                                "state handle should have been set on success")
                        .discardState();
            }
        }
    }
@Test
    void testReplace() throws Exception {
        new Context() {
            {
                runTest(
                        () -> {
                            leaderCallbackGrantLeadership();

                            final KubernetesStateHandleStore<
                                            TestingLongStateHandleHelper.LongStateHandle>
                                    store =
                                            new KubernetesStateHandleStore<>(
                                                    flinkKubeClient,
                                                    LEADER_CONFIGMAP_NAME,
                                                    longStateStorage,
                                                    filter,
                                                    LOCK_IDENTITY);

                            store.addAndLock(key, state);

                            final TestingLongStateHandleHelper.LongStateHandle newState =
                                    new TestingLongStateHandleHelper.LongStateHandle(23456L);
                            final StringResourceVersion resourceVersion = store.exists(key);
                            store.replace(key, resourceVersion, newState);

                            assertThat(store.getAllAndLock()).hasSize(1);
                            assertThat(store.getAndLock(key).retrieveState()).isEqualTo(newState);
                        });
            }
        };
    }