flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateCheckpointOutputStream.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.flink.runtime.state;

import org.apache.flink.annotation.PublicEvolving;
import org.apache.flink.util.Preconditions;

import java.io.IOException;
import java.util.Arrays;

/**
 * Checkpoint output stream that allows to write raw keyed state in a partitioned way, split into
 * key-groups.
 */
@PublicEvolving
public final class KeyedStateCheckpointOutputStream
        extends NonClosingCheckpointOutputStream<KeyGroupsStateHandle> {

    public static final long NO_OFFSET_SET = -1L;
    public static final int NO_CURRENT_KEY_GROUP = -1;

    private int currentKeyGroup;
    private final KeyGroupRangeOffsets keyGroupRangeOffsets;

    public KeyedStateCheckpointOutputStream(
            CheckpointStreamFactory.CheckpointStateOutputStream delegate,
            KeyGroupRange keyGroupRange) {

        super(delegate);
        Preconditions.checkNotNull(keyGroupRange);
        Preconditions.checkArgument(keyGroupRange != KeyGroupRange.EMPTY_KEY_GROUP_RANGE);

        this.currentKeyGroup = NO_CURRENT_KEY_GROUP;
        long[] emptyOffsets = new long[keyGroupRange.getNumberOfKeyGroups()];
        // mark offsets as currently not set
        Arrays.fill(emptyOffsets, NO_OFFSET_SET);
        this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(keyGroupRange, emptyOffsets);
    }

    @Override
    public void close() throws IOException {
        // users should not be able to actually close the stream, it is closed by the system.
        // TODO if we want to support async writes, this call could trigger a callback to the
        // snapshot context that a handle is available.
    }

    /** Returns a list of all key-groups which can be written to this stream. */
    public KeyGroupsList getKeyGroupList() {
        return keyGroupRangeOffsets.getKeyGroupRange();
    }

    /**
     * User code can call this method to signal that it begins to write a new key group with the
     * given key group id. This id must be within the {@link KeyGroupsList} provided by the stream.
     * Each key-group can only be started once and is considered final/immutable as soon as this
     * method is called again.
     */
    public void startNewKeyGroup(int keyGroupId) throws IOException {
        if (isKeyGroupAlreadyStarted(keyGroupId)) {
            throw new IOException("Key group " + keyGroupId + " already registered!");
        }
        keyGroupRangeOffsets.setKeyGroupOffset(keyGroupId, delegate.getPos());
        currentKeyGroup = keyGroupId;
    }

    /**
     * Returns true, if the key group with the given id was already started. The key group might not
     * yet be finished, if it's id is equal to the return value of {@link #getCurrentKeyGroup()}.
     */
    public boolean isKeyGroupAlreadyStarted(int keyGroupId) {
        return NO_OFFSET_SET != keyGroupRangeOffsets.getKeyGroupOffset(keyGroupId);
    }

    /**
     * Returns true if the key group is already completely written and immutable. It was started and
     * since then another key group has been started.
     */
    public boolean isKeyGroupAlreadyFinished(int keyGroupId) {
        return isKeyGroupAlreadyStarted(keyGroupId) && keyGroupId != getCurrentKeyGroup();
    }

    /**
     * Returns the key group that is currently being written. The key group was started but not yet
     * finished, i.e. data can still be added. If no key group was started, this returns {@link
     * #NO_CURRENT_KEY_GROUP}.
     */
    public int getCurrentKeyGroup() {
        return currentKeyGroup;
    }

    @Override
    KeyGroupsStateHandle closeAndGetHandle() throws IOException {
        StreamStateHandle streamStateHandle = super.closeAndGetHandleAfterLeasesReleased();
        return streamStateHandle != null
                ? new KeyGroupsStateHandle(keyGroupRangeOffsets, streamStateHandle)
                : null;
    }
}