diff --git a/flink-state-backends/flink-statebackend-forst-rs/src/main/java/org/apache/flink/state/forstrs/async/ForStRsAsyncAggregatingState.java b/flink-state-backends/flink-statebackend-forst-rs/src/main/java/org/apache/flink/state/forstrs/async/ForStRsAsyncAggregatingState.java new file mode 100644 index 00000000000000..2df625ad2bcfbc --- /dev/null +++ b/flink-state-backends/flink-statebackend-forst-rs/src/main/java/org/apache/flink/state/forstrs/async/ForStRsAsyncAggregatingState.java @@ -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 backend key type + * @param input value type + * @param accumulator type (persisted) + * @param result type + */ +@Internal +public final class ForStRsAsyncAggregatingState { + + private final ForStRsKeyedStateBackend backend; + private final String stateName; + private final TypeSerializer accSerializer; + private final AggregateFunction aggregateFunction; + private final PerKeyFuturesChain chain; + + public ForStRsAsyncAggregatingState( + ForStRsKeyedStateBackend backend, + String stateName, + TypeSerializer accSerializer, + AggregateFunction aggregateFunction, + PerKeyFuturesChain chain) { + this.backend = backend; + this.stateName = stateName; + this.accSerializer = accSerializer; + this.aggregateFunction = aggregateFunction; + this.chain = chain; + } + + /** Async {@link ForStRsAggregatingState#get()}. */ + public CompletableFuture get() { + return runOnKey( + state -> { + try { + return state.get(); + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + } + + /** Async {@link ForStRsAggregatingState#add(Object)}. */ + public CompletableFuture 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 clear() { + return runOnKey( + state -> { + state.clear(); + return null; + }); + } + + private CompletableFuture runOnKey( + Function, R> op) { + final K capturedKey = backend.getCurrentKey(); + if (capturedKey == null) { + CompletableFuture 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 state = + backend.getAggregatingState( + stateName, accSerializer, aggregateFunction); + return op.apply(state); + } + }); + } +} diff --git a/flink-state-backends/flink-statebackend-forst-rs/src/main/java/org/apache/flink/state/forstrs/async/ForStRsAsyncListState.java b/flink-state-backends/flink-statebackend-forst-rs/src/main/java/org/apache/flink/state/forstrs/async/ForStRsAsyncListState.java new file mode 100644 index 00000000000000..4e418bed43f748 --- /dev/null +++ b/flink-state-backends/flink-statebackend-forst-rs/src/main/java/org/apache/flink/state/forstrs/async/ForStRsAsyncListState.java @@ -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). + * + *

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 backend key type + * @param element type + */ +@Internal +public final class ForStRsAsyncListState { + + private final ForStRsKeyedStateBackend backend; + private final String stateName; + private final TypeSerializer elementSerializer; + private final PerKeyFuturesChain chain; + + public ForStRsAsyncListState( + ForStRsKeyedStateBackend backend, + String stateName, + TypeSerializer elementSerializer, + PerKeyFuturesChain 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> get() { + return runOnKey( + state -> { + try { + List 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 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 update(List values) { + return runOnKey( + state -> { + try { + state.update(values); + return null; + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + } + + /** Async {@link ForStRsListState#addAll(List)}. */ + public CompletableFuture addAll(List values) { + return runOnKey( + state -> { + try { + state.addAll(values); + return null; + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + } + + /** Async {@link ForStRsListState#clear()}. */ + public CompletableFuture clear() { + return runOnKey( + state -> { + state.clear(); + return null; + }); + } + + private CompletableFuture runOnKey(Function, R> op) { + final K capturedKey = backend.getCurrentKey(); + if (capturedKey == null) { + CompletableFuture 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 state = + backend.getListState(stateName, elementSerializer); + return op.apply(state); + } + }); + } +} diff --git a/flink-state-backends/flink-statebackend-forst-rs/src/main/java/org/apache/flink/state/forstrs/async/ForStRsAsyncMapState.java b/flink-state-backends/flink-statebackend-forst-rs/src/main/java/org/apache/flink/state/forstrs/async/ForStRsAsyncMapState.java new file mode 100644 index 00000000000000..8d0f1f6a4c7b46 --- /dev/null +++ b/flink-state-backends/flink-statebackend-forst-rs/src/main/java/org/apache/flink/state/forstrs/async/ForStRsAsyncMapState.java @@ -0,0 +1,222 @@ +/* + * 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.ForStRsMapState; + +import java.util.AbstractMap; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.function.Function; + +/** + * Async wrapper around {@link ForStRsMapState} for Flink 2.x async stateful operators (spec §6e). + * + *

Iteration semantics. {@link ForStRsMapState#iterator()} returns a stateful, FFM-backed + * cursor; calling that off-thread is unsafe because the underlying ForSt-RS iterator + Arena + * lifetime is owned by the worker thread. To keep the async surface easy to reason about we + * materialize iterations into snapshot {@link List}s on the worker thread before completing the + * future. This costs O(N) memory per call but is correct under concurrent {@code setCurrentKey} + * calls and outlives the worker arena. Callers who need streaming-iteration semantics can fall back + * to the sync state via the keyed-backend. + * + * @param backend key type + * @param user-key type + * @param user-value type + */ +@Internal +public final class ForStRsAsyncMapState { + + private final ForStRsKeyedStateBackend backend; + private final String stateName; + private final TypeSerializer userKeySerializer; + private final TypeSerializer userValueSerializer; + private final PerKeyFuturesChain chain; + + public ForStRsAsyncMapState( + ForStRsKeyedStateBackend backend, + String stateName, + TypeSerializer userKeySerializer, + TypeSerializer userValueSerializer, + PerKeyFuturesChain chain) { + this.backend = backend; + this.stateName = stateName; + this.userKeySerializer = userKeySerializer; + this.userValueSerializer = userValueSerializer; + this.chain = chain; + } + + /** Async {@link ForStRsMapState#get(Object)}. */ + public CompletableFuture get(UK uk) { + return runOnKey( + state -> { + try { + return state.get(uk); + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + } + + /** Async {@link ForStRsMapState#put(Object, Object)}. */ + public CompletableFuture put(UK uk, UV uv) { + return runOnKey( + state -> { + try { + state.put(uk, uv); + return null; + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + } + + /** Async {@link ForStRsMapState#remove(Object)}. */ + public CompletableFuture remove(UK uk) { + return runOnKey( + state -> { + state.remove(uk); + return null; + }); + } + + /** Async {@link ForStRsMapState#contains(Object)}. */ + public CompletableFuture contains(UK uk) { + return runOnKey( + state -> { + try { + return state.contains(uk); + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + } + + /** Async {@link ForStRsMapState#isEmpty()}. */ + public CompletableFuture isEmpty() { + return runOnKey( + state -> { + try { + return state.isEmpty(); + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + } + + /** + * Async snapshot of {@link ForStRsMapState#keys()}: returns a fully-materialized list of user + * keys captured on the worker thread. Iteration is safe to consume on any thread. + */ + public CompletableFuture> keys() { + return runOnKey( + state -> { + try { + List out = new ArrayList<>(); + for (UK k : state.keys()) { + out.add(k); + } + return out; + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + } + + /** + * Async snapshot of {@link ForStRsMapState#values()}: returns a fully-materialized list of user + * values captured on the worker thread. + */ + public CompletableFuture> values() { + return runOnKey( + state -> { + try { + List out = new ArrayList<>(); + for (UV v : state.values()) { + out.add(v); + } + return out; + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + } + + /** + * Async snapshot of {@link ForStRsMapState#entries()}: returns a fully-materialized list of + * immutable {@link Map.Entry} pairs captured on the worker thread. + * + *

This is the recommended replacement for the sync {@code iterator()} call — see class + * Javadoc for the rationale on returning {@code List} instead of {@code Iterator}. + */ + public CompletableFuture>> entries() { + return runOnKey( + state -> { + try { + List> out = new ArrayList<>(); + for (Map.Entry e : state.entries()) { + out.add( + new AbstractMap.SimpleImmutableEntry<>( + e.getKey(), e.getValue())); + } + return out; + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + } + + /** Async {@link ForStRsMapState#clear()}. */ + public CompletableFuture clear() { + return runOnKey( + state -> { + state.clear(); + return null; + }); + } + + private CompletableFuture runOnKey(Function, R> op) { + final K capturedKey = backend.getCurrentKey(); + if (capturedKey == null) { + CompletableFuture failed = new CompletableFuture<>(); + failed.completeExceptionally( + new IllegalStateException( + "ForStRsAsyncMapState: 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); + ForStRsMapState state = + backend.getMapState( + stateName, userKeySerializer, userValueSerializer); + return op.apply(state); + } + }); + } +} diff --git a/flink-state-backends/flink-statebackend-forst-rs/src/main/java/org/apache/flink/state/forstrs/async/ForStRsAsyncReducingState.java b/flink-state-backends/flink-statebackend-forst-rs/src/main/java/org/apache/flink/state/forstrs/async/ForStRsAsyncReducingState.java new file mode 100644 index 00000000000000..e9609e72c56dc4 --- /dev/null +++ b/flink-state-backends/flink-statebackend-forst-rs/src/main/java/org/apache/flink/state/forstrs/async/ForStRsAsyncReducingState.java @@ -0,0 +1,115 @@ +/* + * 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.ReduceFunction; +import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.state.forstrs.keyed.ForStRsKeyedStateBackend; +import org.apache.flink.state.forstrs.state.ForStRsReducingState; + +import java.util.concurrent.CompletableFuture; +import java.util.function.Function; + +/** + * Async wrapper around {@link ForStRsReducingState} for Flink 2.x async stateful operators (spec + * §6e). + * + * @param backend key type + * @param element / accumulator type + */ +@Internal +public final class ForStRsAsyncReducingState { + + private final ForStRsKeyedStateBackend backend; + private final String stateName; + private final TypeSerializer serializer; + private final ReduceFunction reduceFunction; + private final PerKeyFuturesChain chain; + + public ForStRsAsyncReducingState( + ForStRsKeyedStateBackend backend, + String stateName, + TypeSerializer serializer, + ReduceFunction reduceFunction, + PerKeyFuturesChain chain) { + this.backend = backend; + this.stateName = stateName; + this.serializer = serializer; + this.reduceFunction = reduceFunction; + this.chain = chain; + } + + /** Async {@link ForStRsReducingState#get()}. */ + public CompletableFuture get() { + return runOnKey( + state -> { + try { + return state.get(); + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + } + + /** Async {@link ForStRsReducingState#add(Object)}. */ + public CompletableFuture add(T value) { + return runOnKey( + state -> { + try { + state.add(value); + return null; + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + } + + /** Async {@link ForStRsReducingState#clear()}. */ + public CompletableFuture clear() { + return runOnKey( + state -> { + state.clear(); + return null; + }); + } + + private CompletableFuture runOnKey(Function, R> op) { + final K capturedKey = backend.getCurrentKey(); + if (capturedKey == null) { + CompletableFuture failed = new CompletableFuture<>(); + failed.completeExceptionally( + new IllegalStateException( + "ForStRsAsyncReducingState: 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); + ForStRsReducingState state = + backend.getReducingState(stateName, serializer, reduceFunction); + return op.apply(state); + } + }); + } +} diff --git a/flink-state-backends/flink-statebackend-forst-rs/src/main/java/org/apache/flink/state/forstrs/async/ForStRsAsyncValueState.java b/flink-state-backends/flink-statebackend-forst-rs/src/main/java/org/apache/flink/state/forstrs/async/ForStRsAsyncValueState.java new file mode 100644 index 00000000000000..3e786916199eb6 --- /dev/null +++ b/flink-state-backends/flink-statebackend-forst-rs/src/main/java/org/apache/flink/state/forstrs/async/ForStRsAsyncValueState.java @@ -0,0 +1,170 @@ +/* + * 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.ForStRsValueState; + +import java.util.concurrent.CompletableFuture; +import java.util.function.Function; + +/** + * Async wrapper around {@link ForStRsValueState} for Flink 2.x async stateful operators (spec §6e). + * + *

Two key-binding modes. + * + *

    + *
  1. Implicit (current-key): {@link #value()} / {@link #update(Object)} / {@link + * #clear()} capture {@linkplain ForStRsKeyedStateBackend#getCurrentKey backend.getCurrentKey()} + * at the moment of the API call. Caller is responsible for ensuring no other thread is + * concurrently mutating the backend's current key (i.e., this is the natural mode for + * single-threaded operator code that issues async ops). + *
  2. Explicit-key: {@link #value(Object)} / {@link #update(Object, Object)} / {@link + * #clear(Object)} accept the key directly. This is the recommended mode for chained + * continuations (e.g., {@code state.value(k).thenCompose(v -> state.update(k, v + 1))}) + * and for thread-pool submitters where a stable key is captured ahead of time. + *
+ * + *

Worker-thread safety. The underlying L5 delegate is single-threaded by design — its + * {@code setCurrentKey}/{@code buildPrefix} use shared mutable buffers. To make the async API + * safe under cross-key parallelism, every worker-thread step (set-key + state-fetch + op) + * synchronises on the {@linkplain ForStRsKeyedStateBackend backend}. Per-key serialisation is + * still provided by {@link PerKeyFuturesChain}; the lock just protects the shared-buffer window. + * + * @param backend key type + * @param state value type + */ +@Internal +public final class ForStRsAsyncValueState { + + private final ForStRsKeyedStateBackend backend; + private final String stateName; + private final TypeSerializer valueSerializer; + private final PerKeyFuturesChain chain; + + public ForStRsAsyncValueState( + ForStRsKeyedStateBackend backend, + String stateName, + TypeSerializer valueSerializer, + PerKeyFuturesChain chain) { + this.backend = backend; + this.stateName = stateName; + this.valueSerializer = valueSerializer; + this.chain = chain; + } + + /** Implicit-key {@link ForStRsValueState#value()} — captures {@code backend.getCurrentKey()}. */ + public CompletableFuture value() { + K k = backend.getCurrentKey(); + if (k == null) { + return failedFuture("ForStRsAsyncValueState.value()"); + } + return value(k); + } + + /** Explicit-key {@link ForStRsValueState#value()}. */ + public CompletableFuture value(K key) { + return runOnKey( + key, + state -> { + try { + return state.value(); + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + } + + /** Implicit-key {@link ForStRsValueState#update(Object)}. */ + public CompletableFuture update(T newValue) { + K k = backend.getCurrentKey(); + if (k == null) { + return failedFuture("ForStRsAsyncValueState.update()"); + } + return update(k, newValue); + } + + /** Explicit-key {@link ForStRsValueState#update(Object)}. */ + public CompletableFuture update(K key, T newValue) { + return runOnKey( + key, + state -> { + try { + state.update(newValue); + return null; + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + } + + /** Implicit-key {@link ForStRsValueState#clear()}. */ + public CompletableFuture clear() { + K k = backend.getCurrentKey(); + if (k == null) { + return failedFuture("ForStRsAsyncValueState.clear()"); + } + return clear(k); + } + + /** Explicit-key {@link ForStRsValueState#clear()}. */ + public CompletableFuture clear(K key) { + return runOnKey( + key, + state -> { + state.clear(); + return null; + }); + } + + private CompletableFuture failedFuture(String caller) { + CompletableFuture failed = new CompletableFuture<>(); + failed.completeExceptionally( + new IllegalStateException( + caller + + ": getCurrentKey() returned null — setCurrentKey must be" + + " invoked before implicit-key async ops")); + return failed; + } + + /** + * Enqueues {@code op} on the per-key chain for {@code key}. The worker thread acquires the + * backend monitor before binding the delegate's current key + fetching the sync state + + * invoking the op. + */ + private CompletableFuture runOnKey(K key, Function, R> op) { + if (key == null) { + CompletableFuture failed = new CompletableFuture<>(); + failed.completeExceptionally(new IllegalStateException("key must be non-null")); + return failed; + } + return chain.enqueue( + key, + () -> { + synchronized (backend) { + backend.setCurrentKey(key); + ForStRsValueState state = + backend.getValueState(stateName, valueSerializer); + return op.apply(state); + } + }); + } +} diff --git a/flink-state-backends/flink-statebackend-forst-rs/src/main/java/org/apache/flink/state/forstrs/async/PerKeyFuturesChain.java b/flink-state-backends/flink-statebackend-forst-rs/src/main/java/org/apache/flink/state/forstrs/async/PerKeyFuturesChain.java new file mode 100644 index 00000000000000..c8fc2bbc3b0fae --- /dev/null +++ b/flink-state-backends/flink-statebackend-forst-rs/src/main/java/org/apache/flink/state/forstrs/async/PerKeyFuturesChain.java @@ -0,0 +1,126 @@ +/* + * 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 java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Executor; +import java.util.function.Supplier; + +/** + * Per-key serialization of asynchronous work via a {@link ConcurrentHashMap} of {@link + * CompletableFuture chain heads}. + * + *

Spec §6e (async state API): Flink 2.x async stateful operators may submit multiple async state + * ops for the same key in flight (e.g., {@code value().thenCompose(v -> update(v + 1))}). Per-key + * serialization preserves "happens-before" semantics: ops for the same key execute in submit order + * on a worker thread; ops for distinct keys run in parallel. + * + *

Wire-up. One {@code PerKeyFuturesChain} instance is shared across all async state + * objects of a single keyed-backend (so {@code valueState.update(K1, v).thenCompose(__ -> + * listState.add(K1, x))} respects ordering across state types). The {@code Executor} is typically + * {@link java.util.concurrent.Executors#newVirtualThreadPerTaskExecutor()} (JDK 21+) so the chain + * can scale to many in-flight per-key tails without exhausting platform threads. + * + *

Memory. The chain head is removed from the map when its future completes, so chains + * shrink to empty when there is no in-flight work for a key. Cross-key parallelism is unbounded; + * the {@code Executor} alone caps total parallelism. + * + * @param key type + */ +@Internal +public final class PerKeyFuturesChain { + + private final ConcurrentHashMap> chains = new ConcurrentHashMap<>(); + private final Executor executor; + + /** + * Constructs a new chain with the supplied executor for all enqueued work. + * + * @param executor the executor on which the {@link Supplier work} will run; must be non-{@code + * null}. Typically {@link + * java.util.concurrent.Executors#newVirtualThreadPerTaskExecutor()}. + */ + public PerKeyFuturesChain(Executor executor) { + if (executor == null) { + throw new NullPointerException("executor"); + } + this.executor = executor; + } + + /** + * Atomically enqueues {@code work} on the per-key chain for {@code key}. The returned future + * completes after the work runs on the configured executor; subsequent {@code enqueue(key, …)} + * calls observe the same head and chain off this future, so per-key ordering is preserved. + * + *

If the supplied {@code work} throws, the returned future completes exceptionally with the + * original throwable; the chain still advances so subsequent ops for the same key proceed. + * + * @param key the key under which to serialize the work + * @param work supplier to invoke on the worker thread; its return value becomes the future's + * value + * @param result type + * @return a {@link CompletableFuture} that completes with the work's result (or exception) once + * the chain reaches it + */ + public CompletableFuture enqueue(K key, Supplier work) { + if (key == null) { + throw new NullPointerException("key"); + } + if (work == null) { + throw new NullPointerException("work"); + } + final CompletableFuture result = new CompletableFuture<>(); + chains.compute( + key, + (k, prior) -> { + CompletableFuture base = + prior != null ? prior : CompletableFuture.completedFuture(null); + base.whenCompleteAsync( + (__, ___) -> { + try { + result.complete(work.get()); + } catch (Throwable t) { + result.completeExceptionally(t); + } + }, + executor); + return result; + }); + // Cleanup: when this future completes, remove from the map IF still the head. + // Using compute() preserves atomicity vs. concurrent enqueues that might have already + // installed a successor head. + result.whenComplete( + (__, ___) -> + chains.compute(key, (k, current) -> current == result ? null : current)); + return result; + } + + /** + * Returns the number of distinct keys with at least one in-flight or recently-completed + * (cleanup not yet observed) chain head. Test/diagnostic only — production code must not make + * decisions based on this number, since it has no meaningful happens-before relationship with + * concurrent {@link #enqueue} calls. + */ + public int activeKeyCount() { + return chains.size(); + } +} diff --git a/flink-state-backends/flink-statebackend-forst-rs/src/main/java/org/apache/flink/state/forstrs/keyed/ForStRsAbstractKeyedStateBackend.java b/flink-state-backends/flink-statebackend-forst-rs/src/main/java/org/apache/flink/state/forstrs/keyed/ForStRsAbstractKeyedStateBackend.java index 877fa6cb43c96f..e382d7aa24ab13 100644 --- a/flink-state-backends/flink-statebackend-forst-rs/src/main/java/org/apache/flink/state/forstrs/keyed/ForStRsAbstractKeyedStateBackend.java +++ b/flink-state-backends/flink-statebackend-forst-rs/src/main/java/org/apache/flink/state/forstrs/keyed/ForStRsAbstractKeyedStateBackend.java @@ -20,6 +20,8 @@ import org.apache.flink.annotation.Internal; import org.apache.flink.api.common.ExecutionConfig; +import org.apache.flink.api.common.functions.AggregateFunction; +import org.apache.flink.api.common.functions.ReduceFunction; import org.apache.flink.api.common.state.State; import org.apache.flink.api.common.state.StateDescriptor; import org.apache.flink.api.common.typeutils.TypeSerializer; @@ -46,10 +48,18 @@ import org.apache.flink.runtime.state.metrics.LatencyTrackingStateConfig; import org.apache.flink.runtime.state.metrics.SizeTrackingStateConfig; import org.apache.flink.runtime.state.ttl.TtlTimeProvider; +import org.apache.flink.state.forstrs.async.ForStRsAsyncAggregatingState; +import org.apache.flink.state.forstrs.async.ForStRsAsyncListState; +import org.apache.flink.state.forstrs.async.ForStRsAsyncMapState; +import org.apache.flink.state.forstrs.async.ForStRsAsyncReducingState; +import org.apache.flink.state.forstrs.async.ForStRsAsyncValueState; +import org.apache.flink.state.forstrs.async.PerKeyFuturesChain; import org.apache.flink.state.forstrs.keyed.sst.ForStRsSstRegistry; import java.io.IOException; import java.util.List; +import java.util.concurrent.Executor; +import java.util.concurrent.Executors; import java.util.concurrent.FutureTask; import java.util.concurrent.RunnableFuture; import java.util.stream.Stream; @@ -92,6 +102,19 @@ public class ForStRsAbstractKeyedStateBackend extends AbstractKeyedStateBacke */ private ForStRsSstRegistry sstRegistry; + /** + * Per-key futures chain shared by all {@code getAsync*State} factories on this backend (spec + * §6e). Lazily initialised on first call so backends used purely synchronously do not pay the + * cost of a virtual-thread executor. + */ + private volatile PerKeyFuturesChain asyncChain; + + /** + * Executor backing {@link #asyncChain}. Owned by this backend; closed in {@link #close()}. + * Lazily created via {@link #ensureAsyncChain()}. + */ + private volatile java.util.concurrent.ExecutorService asyncExecutor; + /** * Convenience constructor that wires the smallest-possible Flink runtime context (no metrics * tracking, no compression, no kvState registry, single key-group range [0, 0]) and delegates @@ -284,11 +307,91 @@ public int numKeyValueStateEntries() { return (int) Math.min(Integer.MAX_VALUE, n); } + // ------------------------------------------------------------------ + // Async state SPI (spec §6e — Flink 2.x async stateful operators) + // ------------------------------------------------------------------ + + /** + * Returns an {@link ForStRsAsyncValueState} bound to {@code stateName}. The returned wrapper + * captures {@linkplain ForStRsKeyedStateBackend#getCurrentKey current key} on each method call + * and serialises per-key ops via this backend's {@link PerKeyFuturesChain}. + * + *

Per-key ordering and cross-key parallelism are guaranteed by the shared chain. Callers may + * submit multiple ops for the same key in flight and observe submit-order completion. + */ + public ForStRsAsyncValueState getAsyncValueState( + String stateName, TypeSerializer valueSerializer) { + return new ForStRsAsyncValueState<>( + delegate, stateName, valueSerializer, ensureAsyncChain()); + } + + /** Async counterpart of {@link ForStRsKeyedStateBackend#getListState}. */ + public ForStRsAsyncListState getAsyncListState( + String stateName, TypeSerializer elementSerializer) { + return new ForStRsAsyncListState<>( + delegate, stateName, elementSerializer, ensureAsyncChain()); + } + + /** Async counterpart of {@link ForStRsKeyedStateBackend#getMapState}. */ + public ForStRsAsyncMapState getAsyncMapState( + String stateName, + TypeSerializer userKeySerializer, + TypeSerializer userValueSerializer) { + return new ForStRsAsyncMapState<>( + delegate, stateName, userKeySerializer, userValueSerializer, ensureAsyncChain()); + } + + /** Async counterpart of {@link ForStRsKeyedStateBackend#getReducingState}. */ + public ForStRsAsyncReducingState getAsyncReducingState( + String stateName, TypeSerializer serializer, ReduceFunction reduceFunction) { + return new ForStRsAsyncReducingState<>( + delegate, stateName, serializer, reduceFunction, ensureAsyncChain()); + } + + /** Async counterpart of {@link ForStRsKeyedStateBackend#getAggregatingState}. */ + public ForStRsAsyncAggregatingState getAsyncAggregatingState( + String stateName, + TypeSerializer accSerializer, + AggregateFunction aggregateFunction) { + return new ForStRsAsyncAggregatingState<>( + delegate, stateName, accSerializer, aggregateFunction, ensureAsyncChain()); + } + + /** + * Lazily constructs the per-key futures chain backed by a virtual-thread executor on first + * call. Idempotent and thread-safe via double-checked locking on the {@code asyncChain} field. + */ + private PerKeyFuturesChain ensureAsyncChain() { + PerKeyFuturesChain chain = asyncChain; + if (chain == null) { + synchronized (this) { + chain = asyncChain; + if (chain == null) { + java.util.concurrent.ExecutorService exec = + Executors.newVirtualThreadPerTaskExecutor(); + chain = new PerKeyFuturesChain<>((Executor) exec); + this.asyncExecutor = exec; + this.asyncChain = chain; + } + } + } + return chain; + } + + /** Test accessor for the async chain (or {@code null} if no async state op has been issued). */ + public PerKeyFuturesChain getAsyncChain() { + return asyncChain; + } + @Override public void close() throws IOException { try { super.close(); } finally { + // Best-effort shutdown of the async executor — long-running async ops are interrupted. + if (asyncExecutor != null) { + asyncExecutor.shutdownNow(); + } delegate.close(); } } diff --git a/flink-state-backends/flink-statebackend-forst-rs/src/test/java/org/apache/flink/state/forstrs/async/ForStRsAsyncValueStateTest.java b/flink-state-backends/flink-statebackend-forst-rs/src/test/java/org/apache/flink/state/forstrs/async/ForStRsAsyncValueStateTest.java new file mode 100644 index 00000000000000..113d1b5466e96a --- /dev/null +++ b/flink-state-backends/flink-statebackend-forst-rs/src/test/java/org/apache/flink/state/forstrs/async/ForStRsAsyncValueStateTest.java @@ -0,0 +1,270 @@ +/* + * 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.api.common.ExecutionConfig; +import org.apache.flink.api.common.typeutils.base.IntSerializer; +import org.apache.flink.api.common.typeutils.base.StringSerializer; +import org.apache.flink.core.fs.CloseableRegistry; +import org.apache.flink.state.forstrs.ffm.ForStRsLinker; +import org.apache.flink.state.forstrs.ffm.FrsCfHandle; +import org.apache.flink.state.forstrs.ffm.FrsDb; +import org.apache.flink.state.forstrs.keyed.ForStRsAbstractKeyedStateBackend; +import org.apache.flink.state.forstrs.keyed.ForStRsKeyedStateBackend; + +import org.junit.jupiter.api.Test; + +import java.lang.foreign.Arena; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Tests for {@link ForStRsAsyncValueState} including the spec-mandated 100k-op / 1k-key concurrent + * stress test (Task 8.9). + * + *

For the stress test we run 100 ops per key over 1k distinct keys (total 100k ops). Each op is + * a get-then-put cycle: read current value, write {@code v + 1}. With per-key serialization + * provided by {@link PerKeyFuturesChain}, the final value per key must equal the number of ops + * issued (100), proving no lost updates regardless of cross-key parallelism. + */ +class ForStRsAsyncValueStateTest { + + private static ForStRsAbstractKeyedStateBackend buildBackend( + Arena arena, + ForStRsLinker linker, + FrsDb db, + FrsCfHandle cf, + org.apache.flink.api.common.typeutils.TypeSerializer keySer, + CloseableRegistry registry) { + ForStRsKeyedStateBackend delegate = + new ForStRsKeyedStateBackend<>(arena, linker, db, cf, keySer, false); + return new ForStRsAbstractKeyedStateBackend<>( + keySer, + Thread.currentThread().getContextClassLoader(), + new ExecutionConfig(), + registry, + delegate); + } + + @Test + void singleKeyAsyncRoundTrip() throws Exception { + try (Arena arena = Arena.ofShared()) { + ForStRsLinker linker = new ForStRsLinker(arena); + try (FrsDb db = linker.dbOpenMemory(arena); + FrsCfHandle cf = linker.dbDefaultCf(db, arena); + CloseableRegistry cr = new CloseableRegistry(); + ForStRsAbstractKeyedStateBackend backend = + buildBackend(arena, linker, db, cf, StringSerializer.INSTANCE, cr)) { + + backend.getDelegate().setCurrentKey("kA"); + ForStRsAsyncValueState state = + backend.getAsyncValueState("counter", IntSerializer.INSTANCE); + + assertNull(state.value().get(5, TimeUnit.SECONDS)); + state.update(1).get(5, TimeUnit.SECONDS); + assertEquals(1, state.value().get(5, TimeUnit.SECONDS)); + state.update(42).get(5, TimeUnit.SECONDS); + assertEquals(42, state.value().get(5, TimeUnit.SECONDS)); + state.clear().get(5, TimeUnit.SECONDS); + assertNull(state.value().get(5, TimeUnit.SECONDS)); + } + } + } + + @Test + void asyncSubmitFailsBeforeSetCurrentKey() throws Exception { + try (Arena arena = Arena.ofShared()) { + ForStRsLinker linker = new ForStRsLinker(arena); + try (FrsDb db = linker.dbOpenMemory(arena); + FrsCfHandle cf = linker.dbDefaultCf(db, arena); + CloseableRegistry cr = new CloseableRegistry(); + ForStRsAbstractKeyedStateBackend backend = + buildBackend(arena, linker, db, cf, StringSerializer.INSTANCE, cr)) { + + ForStRsAsyncValueState state = + backend.getAsyncValueState("nokey", IntSerializer.INSTANCE); + CompletableFuture f = state.value(); + assertTrue(f.isCompletedExceptionally()); + assertThrows( + java.util.concurrent.ExecutionException.class, + () -> f.get(1, TimeUnit.SECONDS)); + } + } + } + + /** + * Spec stress: 100 ops × 1k keys = 100k total. Each op is a get + update(get+1); per-key + * serialization must yield final value = 100 for every key. + */ + @Test + void stressGetThenPut100kAcross1kKeys() throws Exception { + final int keys = 1_000; + final int opsPerKey = 100; + + try (Arena arena = Arena.ofShared()) { + ForStRsLinker linker = new ForStRsLinker(arena); + try (FrsDb db = linker.dbOpenMemory(arena); + FrsCfHandle cf = linker.dbDefaultCf(db, arena); + CloseableRegistry cr = new CloseableRegistry(); + ForStRsAbstractKeyedStateBackend backend = + buildBackend(arena, linker, db, cf, IntSerializer.INSTANCE, cr)) { + + ForStRsKeyedStateBackend delegate = backend.getDelegate(); + + // Submitter thread pool — also virtual threads, separate from the chain executor + // owned by the backend. Each submitter sets the current key on the delegate, + // captures it in the async wrapper, then submits both get and update. + List> finals = new ArrayList<>(keys); + AtomicInteger submitterErrors = new AtomicInteger(0); + + // Spawn one virtual thread per key that does opsPerKey serial submits. Each submit + // races with submits for OTHER keys; per-key the chain serializes get→update. + // We synchronize on the delegate around setCurrentKey + state-fetch because the + // delegate's setCurrentKey is not thread-safe — but per-key chain semantics still + // hold because each thread captures its own key value before enqueueing. + List threads = new ArrayList<>(keys); + List>> perKeyOps = new ArrayList<>(keys); + for (int i = 0; i < keys; i++) { + perKeyOps.add(new ArrayList<>(opsPerKey)); + } + for (int kIdx = 0; kIdx < keys; kIdx++) { + final int kFinal = kIdx; + Thread t = + Thread.ofVirtual() + .name("submitter-" + kFinal) + .unstarted( + () -> { + try { + // Use the explicit-key async API so submit + // and continuation both bind to the same + // captured key without depending on the + // shared delegate's currentKey field. + ForStRsAsyncValueState + state = + backend + . + getAsyncValueState( + "ctr", + IntSerializer + .INSTANCE); + // Chain each iteration off the previous so + // get-then-put on the same key is observed + // as a single atomic step before the next + // iteration's get is enqueued. Without this + // the submitter would enqueue 100 gets + // back-to-back and only then start enqueuing + // updates as gets complete, producing + // classic lost-update behaviour. + CompletableFuture chainTail = + CompletableFuture.completedFuture(null); + for (int op = 0; op < opsPerKey; op++) { + chainTail = + chainTail.thenCompose( + ___ -> + state.value(kFinal) + .thenCompose( + v -> + state + .update( + kFinal, + v + == null + ? 1 + : v + + 1))); + } + synchronized (perKeyOps.get(kFinal)) { + perKeyOps.get(kFinal).add(chainTail); + } + } catch (Throwable t1) { + submitterErrors.incrementAndGet(); + } + }); + threads.add(t); + } + for (Thread t : threads) { + t.start(); + } + for (Thread t : threads) { + t.join(120_000); + } + assertEquals(0, submitterErrors.get(), "submitter thread errors"); + + // Wait for every submitted op to complete. + for (int kIdx = 0; kIdx < keys; kIdx++) { + for (CompletableFuture f : perKeyOps.get(kIdx)) { + f.get(60, TimeUnit.SECONDS); + } + } + + // Assert final value per key == opsPerKey (100). Use the explicit-key API for + // a deadlock-free read (the chain executor needs the backend monitor for the + // op; holding it on the caller while awaiting the future would deadlock). + ForStRsAsyncValueState verifier = + backend.getAsyncValueState("ctr", IntSerializer.INSTANCE); + for (int kIdx = 0; kIdx < keys; kIdx++) { + Integer v = verifier.value(kIdx).get(60, TimeUnit.SECONDS); + assertNotNull(v, "key " + kIdx + " missing value"); + assertEquals( + opsPerKey, + v.intValue(), + "key " + kIdx + " lost updates: expected " + opsPerKey + " got " + v); + } + } + } + } + + /** + * Smaller per-key sequence-number test: enqueues 200 update(i) ops on a single key and verifies + * the final stored value equals 199 (last write wins under per-key serialization). + * Distinguishes the chain's ordering guarantee from a plain "atomicity" property. + */ + @Test + void perKeySerializationProducesLastWriteWins() throws Exception { + try (Arena arena = Arena.ofShared()) { + ForStRsLinker linker = new ForStRsLinker(arena); + try (FrsDb db = linker.dbOpenMemory(arena); + FrsCfHandle cf = linker.dbDefaultCf(db, arena); + CloseableRegistry cr = new CloseableRegistry(); + ForStRsAbstractKeyedStateBackend backend = + buildBackend(arena, linker, db, cf, StringSerializer.INSTANCE, cr)) { + backend.getDelegate().setCurrentKey("seq-key"); + ForStRsAsyncValueState state = + backend.getAsyncValueState("seq", IntSerializer.INSTANCE); + List> fs = new ArrayList<>(200); + for (int i = 0; i < 200; i++) { + fs.add(state.update(i)); + } + for (CompletableFuture f : fs) { + f.get(5, TimeUnit.SECONDS); + } + assertEquals(199, state.value().get(5, TimeUnit.SECONDS)); + } + } + } +} diff --git a/flink-state-backends/flink-statebackend-forst-rs/src/test/java/org/apache/flink/state/forstrs/async/PerKeyFuturesChainTest.java b/flink-state-backends/flink-statebackend-forst-rs/src/test/java/org/apache/flink/state/forstrs/async/PerKeyFuturesChainTest.java new file mode 100644 index 00000000000000..661466aac022ed --- /dev/null +++ b/flink-state-backends/flink-statebackend-forst-rs/src/test/java/org/apache/flink/state/forstrs/async/PerKeyFuturesChainTest.java @@ -0,0 +1,330 @@ +/* + * 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.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.IntStream; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * 12 unit tests covering {@link PerKeyFuturesChain}'s ordering, isolation, error-recovery, and + * map-shrinkage invariants. + */ +class PerKeyFuturesChainTest { + + private ExecutorService executor; + + @BeforeEach + void setUp() { + // Virtual-thread executor as called for in spec §6e — the chain is designed to run on + // many in-flight virtual threads without exhausting platform threads. + executor = Executors.newVirtualThreadPerTaskExecutor(); + } + + @AfterEach + void tearDown() { + executor.shutdown(); + } + + /** 1: a single enqueue on a key returns the supplier's value. */ + @Test + void singleEnqueueReturnsValue() throws Exception { + PerKeyFuturesChain chain = new PerKeyFuturesChain<>(executor); + CompletableFuture f = chain.enqueue("k", () -> 42); + assertEquals(42, f.get(5, TimeUnit.SECONDS)); + } + + /** 2: per-key ordering — 100 enqueues for the same key complete in submit order. */ + @Test + void perKeyOrderingPreserved() throws Exception { + PerKeyFuturesChain chain = new PerKeyFuturesChain<>(executor); + List observed = new ArrayList<>(); + List> futures = new ArrayList<>(); + for (int i = 0; i < 100; i++) { + final int v = i; + futures.add( + chain.enqueue( + "k", + () -> { + synchronized (observed) { + observed.add(v); + } + return v; + })); + } + for (CompletableFuture f : futures) { + f.get(5, TimeUnit.SECONDS); + } + for (int i = 0; i < 100; i++) { + assertEquals(i, observed.get(i), "submit-order violated at index " + i); + } + } + + /** 3: cross-key parallelism — two slow ops on different keys overlap in time. */ + @Test + void crossKeyParallelism() throws Exception { + PerKeyFuturesChain chain = new PerKeyFuturesChain<>(executor); + CountDownLatch entered = new CountDownLatch(2); + CountDownLatch release = new CountDownLatch(1); + CompletableFuture a = + chain.enqueue( + "kA", + () -> { + entered.countDown(); + try { + return release.await(5, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return false; + } + }); + CompletableFuture b = + chain.enqueue( + "kB", + () -> { + entered.countDown(); + try { + return release.await(5, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return false; + } + }); + // Both must enter their suppliers without one waiting on the other; if cross-key serialized + // this latch would never reach 0. + assertTrue( + entered.await(5, TimeUnit.SECONDS), + "different-key suppliers were not concurrently entered"); + release.countDown(); + assertTrue(a.get(5, TimeUnit.SECONDS)); + assertTrue(b.get(5, TimeUnit.SECONDS)); + } + + /** 4: chain shrinks on completion — single key returns to empty. */ + @Test + void chainShrinksAfterCompletion() throws Exception { + PerKeyFuturesChain chain = new PerKeyFuturesChain<>(executor); + CompletableFuture f = chain.enqueue("k", () -> 1); + f.get(5, TimeUnit.SECONDS); + // Cleanup runs in the chain itself via whenComplete; give it a brief window to observe. + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (chain.activeKeyCount() != 0 && System.nanoTime() < deadline) { + Thread.sleep(1); + } + assertEquals(0, chain.activeKeyCount()); + } + + /** 5: chain shrinks on completion — multiple keys all clean up. */ + @Test + void multiKeyChainShrinks() throws Exception { + PerKeyFuturesChain chain = new PerKeyFuturesChain<>(executor); + List> fs = new ArrayList<>(); + for (int i = 0; i < 50; i++) { + final int v = i; + fs.add(chain.enqueue(v, () -> v)); + } + for (CompletableFuture f : fs) { + f.get(5, TimeUnit.SECONDS); + } + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (chain.activeKeyCount() != 0 && System.nanoTime() < deadline) { + Thread.sleep(1); + } + assertEquals(0, chain.activeKeyCount()); + } + + /** 6: throwing supplier completes future exceptionally; chain still advances. */ + @Test + void exceptionInSupplierPropagates() { + PerKeyFuturesChain chain = new PerKeyFuturesChain<>(executor); + CompletableFuture f = + chain.enqueue( + "k", + () -> { + throw new IllegalStateException("boom"); + }); + ExecutionException ex = + assertThrows(ExecutionException.class, () -> f.get(5, TimeUnit.SECONDS)); + assertNotNull(ex.getCause()); + assertEquals("boom", ex.getCause().getMessage()); + } + + /** 7: after a throwing supplier, the next enqueue on the same key still runs. */ + @Test + void chainAdvancesPastException() throws Exception { + PerKeyFuturesChain chain = new PerKeyFuturesChain<>(executor); + CompletableFuture failing = + chain.enqueue( + "k", + () -> { + throw new RuntimeException("fail"); + }); + CompletableFuture next = chain.enqueue("k", () -> 99); + assertEquals(99, next.get(5, TimeUnit.SECONDS)); + assertTrue(failing.isCompletedExceptionally()); + } + + /** 8: per-key ordering preserved even under interleaved enqueues across many keys. */ + @Test + void interleavedKeysPreservePerKeyOrder() throws Exception { + PerKeyFuturesChain chain = new PerKeyFuturesChain<>(executor); + final int keys = 8; + final int opsPerKey = 50; + @SuppressWarnings("unchecked") + List[] observed = new List[keys]; + for (int i = 0; i < keys; i++) { + observed[i] = new ArrayList<>(); + } + List> all = new ArrayList<>(); + for (int op = 0; op < opsPerKey; op++) { + for (int k = 0; k < keys; k++) { + final int kFinal = k; + final int opFinal = op; + all.add( + chain.enqueue( + kFinal, + () -> { + synchronized (observed[kFinal]) { + observed[kFinal].add(opFinal); + } + return null; + })); + } + } + for (CompletableFuture f : all) { + f.get(5, TimeUnit.SECONDS); + } + for (int k = 0; k < keys; k++) { + for (int op = 0; op < opsPerKey; op++) { + assertEquals(op, observed[k].get(op), "key=" + k + " op-index=" + op); + } + } + } + + /** 9: null key throws NPE eagerly. */ + @Test + void nullKeyRejected() { + PerKeyFuturesChain chain = new PerKeyFuturesChain<>(executor); + assertThrows(NullPointerException.class, () -> chain.enqueue(null, () -> 1)); + } + + /** 10: null work throws NPE eagerly. */ + @Test + void nullWorkRejected() { + PerKeyFuturesChain chain = new PerKeyFuturesChain<>(executor); + assertThrows(NullPointerException.class, () -> chain.enqueue("k", null)); + } + + /** 11: null executor in constructor rejected. */ + @Test + void nullExecutorRejected() { + assertThrows(NullPointerException.class, () -> new PerKeyFuturesChain<>(null)); + } + + /** + * 12: a supplier returning {@code null} completes the future with {@code null} (not a NPE in + * cleanup). Distinguishes "absent value" from "result not yet produced". + */ + @Test + void nullResultIsLegal() throws Exception { + PerKeyFuturesChain chain = new PerKeyFuturesChain<>(executor); + CompletableFuture f = chain.enqueue("k", () -> null); + assertNull(f.get(5, TimeUnit.SECONDS)); + // chain still shrinks for null-returning suppliers + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (chain.activeKeyCount() != 0 && System.nanoTime() < deadline) { + Thread.sleep(1); + } + assertEquals(0, chain.activeKeyCount()); + } + + /** + * Bonus stress: 1k keys × 50 ops with concurrent submitters; counts each per-key supplier + * invocation and asserts the count == ops. Lives here (not in ForStRsAsyncValueStateTest) + * because it covers the chain itself, not the state wrapper. + */ + @Test + void stressMultiSubmitter() throws Exception { + PerKeyFuturesChain chain = new PerKeyFuturesChain<>(executor); + final int keys = 1_000; + final int opsPerKey = 50; + AtomicInteger[] perKeyCount = new AtomicInteger[keys]; + for (int i = 0; i < keys; i++) { + perKeyCount[i] = new AtomicInteger(); + } + List> all = new ArrayList<>(keys * opsPerKey); + IntStream.range(0, opsPerKey) + .forEach( + op -> + IntStream.range(0, keys) + .forEach( + k -> + all.add( + chain.enqueue( + k, + () -> { + perKeyCount[k] + .incrementAndGet(); + return null; + })))); + for (CompletableFuture f : all) { + f.get(30, TimeUnit.SECONDS); + } + for (int k = 0; k < keys; k++) { + assertEquals(opsPerKey, perKeyCount[k].get(), "per-key op count diverged at key " + k); + } + } + + /** + * Sanity that {@link Objects#hashCode} based key equality (mutable wrapper) still works for the + * chain — if we ever swap the impl from ConcurrentHashMap.compute to a different primitive this + * guards against an accidental identity-comparison regression. + */ + @Test + void distinctKeyInstancesWithEqualValueShareChain() { + PerKeyFuturesChain chain = new PerKeyFuturesChain<>(executor); + // Two `new String("k")` instances are equals() but != by identity. + assertDoesNotThrow(() -> chain.enqueue(new String("k"), () -> 1).get(5, TimeUnit.SECONDS)); + // Another enqueue with an equal-but-distinct String should reuse the chain head, not start + // a new one. + CompletableFuture f = chain.enqueue(new String("k"), () -> 2); + assertFalse(f.isCompletedExceptionally()); + assertDoesNotThrow(() -> f.get(5, TimeUnit.SECONDS)); + } +}