Skip to content

B-Prod-P8: async state API (Flink 2.x) — spec §6e #8

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jackylee-ch wants to merge 1 commit into
base: b-prod-p7-tuning
Choose a base branch
from
Open

B-Prod-P8: async state API (Flink 2.x) — spec §6e #8

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
119 changes: 119 additions & 0 deletions ...t-rs/src/main/java/org/apache/flink/state/forstrs/async/ForStRsAsyncAggregatingState.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,119 @@
/*
* 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.state.forstrs.async;

import org.apache.flink.annotation.Internal;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.state.forstrs.keyed.ForStRsKeyedStateBackend;
import org.apache.flink.state.forstrs.state.ForStRsAggregatingState;

import java.util.concurrent.CompletableFuture;
import java.util.function.Function;

/**
* Async wrapper around {@link ForStRsAggregatingState} for Flink 2.x async stateful operators (spec
* §6e).
*
* @param <K> backend key type
* @param <IN> input value type
* @param <ACC> accumulator type (persisted)
* @param <OUT> result type
*/
@Internal
public final class ForStRsAsyncAggregatingState<K, IN, ACC, OUT> {

private final ForStRsKeyedStateBackend<K> backend;
private final String stateName;
private final TypeSerializer<ACC> accSerializer;
private final AggregateFunction<IN, ACC, OUT> aggregateFunction;
private final PerKeyFuturesChain<K> chain;

public ForStRsAsyncAggregatingState(
ForStRsKeyedStateBackend<K> backend,
String stateName,
TypeSerializer<ACC> accSerializer,
AggregateFunction<IN, ACC, OUT> aggregateFunction,
PerKeyFuturesChain<K> chain) {
this.backend = backend;
this.stateName = stateName;
this.accSerializer = accSerializer;
this.aggregateFunction = aggregateFunction;
this.chain = chain;
}

/** Async {@link ForStRsAggregatingState#get()}. */
public CompletableFuture<OUT> get() {
return runOnKey(
state -> {
try {
return state.get();
} catch (Exception e) {
throw new RuntimeException(e);
}
});
}

/** Async {@link ForStRsAggregatingState#add(Object)}. */
public CompletableFuture<Void> add(IN value) {
return runOnKey(
state -> {
try {
state.add(value);
return null;
} catch (Exception e) {
throw new RuntimeException(e);
}
});
}

/** Async {@link ForStRsAggregatingState#clear()}. */
public CompletableFuture<Void> clear() {
return runOnKey(
state -> {
state.clear();
return null;
});
}

private <R> CompletableFuture<R> runOnKey(
Function<ForStRsAggregatingState<IN, ACC, OUT>, R> op) {
final K capturedKey = backend.getCurrentKey();
if (capturedKey == null) {
CompletableFuture<R> failed = new CompletableFuture<>();
failed.completeExceptionally(
new IllegalStateException(
"ForStRsAsyncAggregatingState: getCurrentKey() returned null"));
return failed;
}
return chain.enqueue(
capturedKey,
() -> {
// See ForStRsAsyncValueState.runOnKey for the rationale on locking the
// delegate during the setCurrentKey + buildPrefix + state-op window.
synchronized (backend) {
backend.setCurrentKey(capturedKey);
ForStRsAggregatingState<IN, ACC, OUT> state =
backend.getAggregatingState(
stateName, accSerializer, aggregateFunction);
return op.apply(state);
}
});
}
}
146 changes: 146 additions & 0 deletions ...nd-forst-rs/src/main/java/org/apache/flink/state/forstrs/async/ForStRsAsyncListState.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,146 @@
/*
* 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.state.forstrs.async;

import org.apache.flink.annotation.Internal;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.state.forstrs.keyed.ForStRsKeyedStateBackend;
import org.apache.flink.state.forstrs.state.ForStRsListState;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.function.Function;

/**
* Async wrapper around {@link ForStRsListState} for Flink 2.x async stateful operators (spec §6e).
*
* <p>Note on {@code get()}: the underlying sync state returns an {@link Iterable}, but to make the
* result safe to consume on any thread (and immune to subsequent backend mutations from other keys)
* we materialize it into a {@link List} before completing the future.
*
* @param <K> backend key type
* @param <T> element type
*/
@Internal
public final class ForStRsAsyncListState<K, T> {

private final ForStRsKeyedStateBackend<K> backend;
private final String stateName;
private final TypeSerializer<T> elementSerializer;
private final PerKeyFuturesChain<K> chain;

public ForStRsAsyncListState(
ForStRsKeyedStateBackend<K> backend,
String stateName,
TypeSerializer<T> elementSerializer,
PerKeyFuturesChain<K> chain) {
this.backend = backend;
this.stateName = stateName;
this.elementSerializer = elementSerializer;
this.chain = chain;
}

/** Async {@link ForStRsListState#get()} — materialized into a snapshot {@link List}. */
public CompletableFuture<List<T>> get() {
return runOnKey(
state -> {
try {
List<T> out = new ArrayList<>();
for (T t : state.get()) {
out.add(t);
}
return out;
} catch (Exception e) {
throw new RuntimeException(e);
}
});
}

/** Async {@link ForStRsListState#add(Object)}. */
public CompletableFuture<Void> add(T value) {
return runOnKey(
state -> {
try {
state.add(value);
return null;
} catch (Exception e) {
throw new RuntimeException(e);
}
});
}

/** Async {@link ForStRsListState#update(List)}. */
public CompletableFuture<Void> update(List<T> values) {
return runOnKey(
state -> {
try {
state.update(values);
return null;
} catch (Exception e) {
throw new RuntimeException(e);
}
});
}

/** Async {@link ForStRsListState#addAll(List)}. */
public CompletableFuture<Void> addAll(List<T> values) {
return runOnKey(
state -> {
try {
state.addAll(values);
return null;
} catch (Exception e) {
throw new RuntimeException(e);
}
});
}

/** Async {@link ForStRsListState#clear()}. */
public CompletableFuture<Void> clear() {
return runOnKey(
state -> {
state.clear();
return null;
});
}

private <R> CompletableFuture<R> runOnKey(Function<ForStRsListState<T>, R> op) {
final K capturedKey = backend.getCurrentKey();
if (capturedKey == null) {
CompletableFuture<R> failed = new CompletableFuture<>();
failed.completeExceptionally(
new IllegalStateException(
"ForStRsAsyncListState: getCurrentKey() returned null"));
return failed;
}
return chain.enqueue(
capturedKey,
() -> {
// See ForStRsAsyncValueState.runOnKey for the rationale on locking the
// delegate during the setCurrentKey + buildPrefix + state-op window.
synchronized (backend) {
backend.setCurrentKey(capturedKey);
ForStRsListState<T> state =
backend.getListState(stateName, elementSerializer);
return op.apply(state);
}
});
}
}
Loading
Loading