org.apache.flink.runtime.state.ttl.TtlMapState Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of flink-runtime_2.11 Show documentation
There is a newer version: 1.13.6
/*
 * 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.ttl;

import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.common.typeutils.base.MapSerializer;
import org.apache.flink.runtime.state.internal.InternalMapState;
import org.apache.flink.util.FlinkRuntimeException;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;

import java.util.AbstractMap;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.function.Function;

/**
 * This class wraps map state with TTL logic.
 *
 * @param  The type of key the state is associated to
 * @param  The type of the namespace
 * @param  Type of the user entry key of state with TTL
 * @param  Type of the user entry value of state with TTL
 */
class TtlMapState
	extends AbstractTtlState, Map>, InternalMapState>>
	implements InternalMapState {
	TtlMapState(TtlStateContext>, Map> ttlStateContext) {
		super(ttlStateContext);
	}

	@Override
	public UV get(UK key) throws Exception {
		TtlValue ttlValue = getWrapped(key);
		return ttlValue == null ? null : ttlValue.getUserValue();
	}

	private TtlValue getWrapped(UK key) throws Exception {
		accessCallback.run();
		return getWrappedWithTtlCheckAndUpdate(
			() -> original.get(key), v -> original.put(key, v), () -> original.remove(key));
	}

	@Override
	public void put(UK key, UV value) throws Exception {
		accessCallback.run();
		original.put(key, wrapWithTs(value));
	}

	@Override
	public void putAll(Map map) throws Exception {
		accessCallback.run();
		if (map == null) {
			return;
		}
		Map> ttlMap = new HashMap<>(map.size());
		long currentTimestamp = timeProvider.currentTimestamp();
		for (Map.Entry entry : map.entrySet()) {
			UK key = entry.getKey();
			ttlMap.put(key, TtlUtils.wrapWithTs(entry.getValue(), currentTimestamp));
		}
		original.putAll(ttlMap);
	}

	@Override
	public void remove(UK key) throws Exception {
		accessCallback.run();
		original.remove(key);
	}

	@Override
	public boolean contains(UK key) throws Exception {
		TtlValue ttlValue = getWrapped(key);
		return ttlValue != null;
	}

	@Override
	public Iterable> entries() throws Exception {
		return entries(e -> e);
	}

	private  Iterable entries(
		Function, R> resultMapper) throws Exception {
		accessCallback.run();
		Iterable>> withTs = original.entries();
		return () -> new EntriesIterator<>(withTs == null ? Collections.emptyList() : withTs, resultMapper);
	}

	@Override
	public Iterable keys() throws Exception {
		return entries(Map.Entry::getKey);
	}

	@Override
	public Iterable values() throws Exception {
		return entries(Map.Entry::getValue);
	}

	@Override
	public Iterator> iterator() throws Exception {
		return entries().iterator();
	}

	@Nullable
	@Override
	public Map> getUnexpiredOrNull(@Nonnull Map> ttlValue) {
		Map> unexpired = new HashMap<>();
		TypeSerializer> valueSerializer =
			((MapSerializer>) original.getValueSerializer()).getValueSerializer();
		for (Map.Entry> e : ttlValue.entrySet()) {
			if (!expired(e.getValue())) {
				// we have to do the defensive copy to update the value
				unexpired.put(e.getKey(), valueSerializer.copy(e.getValue()));
			}
		}
		return ttlValue.size() == unexpired.size() ? ttlValue : unexpired;
	}

	@Override
	public void clear() {
		original.clear();
	}

	private class EntriesIterator implements Iterator {
		private final Iterator>> originalIterator;
		private final Function, R> resultMapper;
		private Map.Entry nextUnexpired = null;
		private boolean rightAfterNextIsCalled = false;

		private EntriesIterator(
			@Nonnull Iterable>> withTs,
			@Nonnull Function, R> resultMapper) {
			this.originalIterator = withTs.iterator();
			this.resultMapper = resultMapper;
		}

		@Override
		public boolean hasNext() {
			rightAfterNextIsCalled = false;
			while (nextUnexpired == null && originalIterator.hasNext()) {
				nextUnexpired = getUnexpiredAndUpdateOrCleanup(originalIterator.next());
			}
			return nextUnexpired != null;
		}

		@Override
		public R next() {
			if (hasNext()) {
				rightAfterNextIsCalled = true;
				R result = resultMapper.apply(nextUnexpired);
				nextUnexpired = null;
				return result;
			}
			throw new NoSuchElementException();
		}

		@Override
		public void remove() {
			if (rightAfterNextIsCalled) {
				originalIterator.remove();
			} else {
				throw new IllegalStateException("next() has not been called or hasNext() has been called afterwards," +
					" remove() is supported only right after calling next()");
			}
		}

		private Map.Entry getUnexpiredAndUpdateOrCleanup(Map.Entry> e) {
			TtlValue unexpiredValue;
			try {
				unexpiredValue = getWrappedWithTtlCheckAndUpdate(
					e::getValue,
					v -> original.put(e.getKey(), v),
					originalIterator::remove);
			} catch (Exception ex) {
				throw new FlinkRuntimeException(ex);
			}
			return unexpiredValue == null ? null : new AbstractMap.SimpleEntry<>(e.getKey(), unexpiredValue.getUserValue());
		}
	}
}