org.apache.flink.runtime.state.gemini.engine.memstore.SegmentKListImpl Maven / Gradle / Ivy

/*
 * 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.gemini.engine.memstore;

import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.runtime.state.gemini.engine.GRegionContext;
import org.apache.flink.runtime.state.gemini.engine.exceptions.GeminiRuntimeException;
import org.apache.flink.runtime.state.gemini.engine.page.GValueType;
import org.apache.flink.runtime.state.gemini.engine.utils.SeqIDUtils;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * SegmentKListImpl.
 */
public class SegmentKListImpl implements SegmentKList {
	private final long segmentID;
	private final GRegionContext gRegionContext;
	private final long version;
	private final Map> dataMap;
	private int recordCount = 0;

	private boolean writeCopy;

	private TypeSerializer keySerializer;

	private TypeSerializer elementSerializer;

	@SuppressWarnings("unchecked")
	public SegmentKListImpl(long segmentID, GRegionContext gRegionContext) {
		this.segmentID = segmentID;
		this.gRegionContext = gRegionContext;
		this.dataMap = new HashMap<>();
		this.version = gRegionContext.getGContext().getCurVersion();
		this.writeCopy = gRegionContext.getGContext().getGConfiguration().isWriteCopy();
		this.keySerializer = gRegionContext.getPageSerdeFlink().getKeySerde();
		this.elementSerializer = gRegionContext.getPageSerdeFlink().getValueSerde();
	}

	@SuppressWarnings("unchecked")
	public SegmentKListImpl(long segmentID, GRegionContext gRegionContext, Map> dataMap) {
		this.segmentID = segmentID;
		this.gRegionContext = gRegionContext;
		this.dataMap = dataMap;
		this.version = gRegionContext.getGContext().getCurVersion();
		this.writeCopy = gRegionContext.getGContext().getGConfiguration().isWriteCopy();
		this.keySerializer = gRegionContext.getPageSerdeFlink().getKeySerde();
		this.elementSerializer = gRegionContext.getPageSerdeFlink().getValueSerde();
	}

	@Override
	public void add(K key, E element) {
		GSValueList gValueList = getOrCreateList(key);
		long seqID = gRegionContext.getNextSeqID();
		internalAdd(element, gValueList, seqID);
		updateListSeqID(gValueList, seqID);
	}

	private void internalAdd(E element, GSValueList gValueList, long seqID) {
		gValueList.value.add(GSValue.of(copyElementIfNeeded(element), GValueType.PutValue, seqID));
		recordCount++;
		gRegionContext.getWriteBufferStats().addTotalRecordCount(1);
	}

	@Override
	public void addAll(K key, Collection elements) {
		final GSValueList gValueList = getOrCreateList(key);
		long seqID = gRegionContext.getNextSeqID();
		Set duplicatedElementChecker = new HashSet<>();

		for (E element : elements) {
			//List will have same element, GeminiDB will gen different seqID for same Element.
			if (!duplicatedElementChecker.add(element)) {
				seqID = gRegionContext.getNextSeqID();
			}
			internalAdd(element, gValueList, seqID);
		}

		updateListSeqID(gValueList, seqID);
	}

	@Override
	public void remove(K key, E element) {
		throw new GeminiRuntimeException("not support remove element");
	}

	@Override
	public void removeAll(K key, Collection elements) {
		elements.forEach((v) -> remove(key, v));
	}

	@Override
	public void put(K key, List> value) {
		GSValueList glist = createPutList();
		long seqID = SeqIDUtils.INVALID_SEQID;
		for (GSValue gsValue : value) {
			seqID = Math.max(seqID, gsValue.getSeqID());
			// TODO better way is to replace put(K, List>) with put(K, List)
			gsValue.setValue(copyElementIfNeeded(gsValue.getValue()));
			glist.getValue().add(gsValue);
		}
		// TODO no need to copy key if the mapping has existed
		GSValueList old = dataMap.put(copyKeyIfNeeded(key), glist);
		updateListSeqID(glist, seqID);
		int oldSize = old == null ? 0 : old.getValue().size();
		int delta = value.size() - oldSize;
		recordCount = recordCount + delta;
		gRegionContext.getWriteBufferStats().addTotalRecordCount(delta);
	}

	@Override
	public GSValueList get(K key) {
		return dataMap.get(key);
	}

	@Override
	public void removeKey(K key) {
		GSValueList gsValueList = createDeleteList();
		long seqID = gRegionContext.getNextSeqID();
		// TODO no need to copy key if the mapping has existed
		dataMap.put(copyKeyIfNeeded(key), gsValueList);
		updateListSeqID(gsValueList, seqID);
	}

	@Override
	public long getSegmentID() {
		return this.segmentID;
	}

	@Override
	public int getRecordCount() {
		return this.recordCount;
	}

	@Override
	public long getVersion() {
		return this.version;
	}

	@Override
	public Segment>> copySegment() {
		Map> copyMap = new HashMap<>();
		for (Map.Entry> entry : dataMap.entrySet()) {
			copyMap.put(entry.getKey(), entry.getValue().copyGSValueList());
		}
		return new SegmentKListImpl<>(-1L, this.gRegionContext, copyMap);
	}

	@Override
	public Map> getData() {
		return dataMap;
	}

	private GSValueList getOrCreateList(K key) {
		GSValueList gValueList = dataMap.get(key);
		if (gValueList == null || gValueList.getSeqID() < gRegionContext.getRemoveAllSeqID()) {
			// if the list does not exist, or was created before removeAll, we create a new list.
			gValueList = createAddList();
			dataMap.put(copyKeyIfNeeded(key), gValueList);
		} else if (gValueList.getValueType() == GValueType.Delete) {
			gValueList = createPutList();
			dataMap.put(copyKeyIfNeeded(key), gValueList);
		}
		return gValueList;
	}

	private void updateListSeqID(GSValueList gsValueList, long seqID) {
		gsValueList.setSeqID(seqID);
	}

	private GSValueList createAddList() {
		return GSValueList.of(new ArrayList<>(), GValueType.AddList, SeqIDUtils.INVALID_SEQID);
	}

	private GSValueList createPutList() {
		return GSValueList.of(new ArrayList<>(), GValueType.PutList, SeqIDUtils.INVALID_SEQID);
	}

	private GSValueList createDeleteList() {
		return GSValueList.of(null, GValueType.Delete, SeqIDUtils.INVALID_SEQID);
	}

	private K copyKeyIfNeeded(K key) {
		return writeCopy ? keySerializer.copy(key) : key;
	}

	private E copyElementIfNeeded(E element) {
		return writeCopy ? elementSerializer.copy(element) : element;
	}
}