From 24f45699c0f83464356b313d7450fb00dcd1c543 Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Sun, 26 Jul 2026 20:25:14 +0500 Subject: [PATCH 1/2] [GSoC 2026] Kafka Streams runner: windowed GroupByKey via ReduceFnRunner Replaces the global-window-only GroupByKey with a windowed one that drives Beam's ReduceFnRunner on the runner side, the way the Flink and Spark portable runners do, backed by Kafka Streams state and timers. WindowedGroupByKeyProcessor builds a ReduceFnRunner per key (like GroupAlsoByWindowViaWindowSetNewDoFn) over two new backends: KafkaStreamsStateInternals, which stores each Beam state cell as one entry in a KeyValueStore under a composite key of key + namespace + tag (modelled on SparkStateInternals), and KafkaStreamsTimerInternals, which persists timers keyed by identity and is fired by the processor scanning for due event-time timers on each input-watermark advance. GroupByKeyTranslator hydrates the input windowing strategy from the pipeline proto and wires the state and timer stores. Windowing, the default trigger, panes, allowed lateness and timestamp combiners all come from ReduceFnRunner. Also drops KafkaStreamsTestRunner.roundTripInternalTopics: TopologyTestDriver loops an internal repartition topic from sink back to source on its own, so the manual round-trip was delivering every GroupByKey record twice. The old global-window GroupByKey masked this by firing each key once and latching; a correctly windowed GroupByKey emits a late pane for the duplicate, which surfaced the bug. Advancing the wall clock now drives the pipeline to completion. validatesRunner goes from 44 to 49 tests: GroupByKeyTest.BasicTests 7 -> 9 (timestamp combiners un-sickbayed) and GroupByKeyTest.WindowTests 0 -> 3. testGroupByKeyMergingWindows stays sickbayed; session (merging) windows land in a follow-up. Adds FixedWindowGroupByKeyTest. --- runners/kafka-streams/build.gradle | 11 +- .../translation/GroupByKeyProcessor.java | 224 ---------- .../translation/GroupByKeyTranslator.java | 58 ++- .../KafkaStreamsStateInternals.java | 393 ++++++++++++++++++ .../KafkaStreamsTimerInternals.java | 180 ++++++++ .../WindowedGroupByKeyProcessor.java | 284 +++++++++++++ .../kafka/streams/KafkaStreamsTestRunner.java | 101 +---- .../FixedWindowGroupByKeyTest.java | 106 +++++ 8 files changed, 1025 insertions(+), 332 deletions(-) delete mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsStateInternals.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternals.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FixedWindowGroupByKeyTest.java diff --git a/runners/kafka-streams/build.gradle b/runners/kafka-streams/build.gradle index 2d0209eafbf5..bdf7e3be0585 100644 --- a/runners/kafka-streams/build.gradle +++ b/runners/kafka-streams/build.gradle @@ -87,11 +87,11 @@ dependencies { // Known-failing @ValidatesRunner tests, excluded until the feature they need lands. def sickbayTests = [ - // Non-global windowing (FixedWindows, merging windows, timestamp combiners) is not supported - // yet; these apply a window and assert on window-derived output, hitting a GlobalWindow cast. - 'org.apache.beam.sdk.transforms.GroupByKeyTest$WindowTests', - 'org.apache.beam.sdk.transforms.GroupByKeyTest$BasicTests.testTimestampCombinerLatest', - 'org.apache.beam.sdk.transforms.GroupByKeyTest$BasicTests.testTimestampCombinerEarliest', + // Merging (session) windows are not supported yet: ReduceFnRunner drives them through a merging + // window set that moves per-window state as windows merge, which this first windowing pass does + // not implement. Non-merging windows (fixed, sliding), the default trigger and timestamp + // combiners do work. Lands with the follow-up windowing PR. + 'org.apache.beam.sdk.transforms.GroupByKeyTest$WindowTests.testGroupByKeyMergingWindows', // A DoFn whose @StartBundle throws never gets to report its error: SdkHarnessClient.newBundle // sends the ProcessBundleRequest and then blocks in GrpcDataService.createOutboundAggregator // waiting for the SDK harness to open its data stream, which a bundle that failed during setup @@ -100,6 +100,7 @@ def sickbayTests = [ // anything specific to this runner; the Flink runner sickbays all of LifecycleTests and the // Prism runner sickbays each of its three error tests. The @ProcessElement and @FinishBundle // variants do pass here, because by then the data stream is established. + // Tracked by https://github.com/apache/beam/issues/39452. 'org.apache.beam.sdk.transforms.ParDoTest$LifecycleTests.testParDoWithErrorInStartBatch', ] diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java deleted file mode 100644 index 3e82935b807d..000000000000 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java +++ /dev/null @@ -1,224 +0,0 @@ -/* - * 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.beam.runners.kafka.streams.translation; - -import java.util.ArrayList; -import java.util.List; -import java.util.Set; -import org.apache.beam.sdk.coders.Coder; -import org.apache.beam.sdk.coders.CoderException; -import org.apache.beam.sdk.coders.IterableCoder; -import org.apache.beam.sdk.transforms.windowing.BoundedWindow; -import org.apache.beam.sdk.transforms.windowing.GlobalWindow; -import org.apache.beam.sdk.util.CoderUtils; -import org.apache.beam.sdk.values.KV; -import org.apache.beam.sdk.values.WindowedValue; -import org.apache.beam.sdk.values.WindowedValues; -import org.apache.kafka.streams.processor.api.Processor; -import org.apache.kafka.streams.processor.api.ProcessorContext; -import org.apache.kafka.streams.processor.api.Record; -import org.apache.kafka.streams.state.KeyValueIterator; -import org.apache.kafka.streams.state.KeyValueStore; -import org.checkerframework.checker.nullness.qual.Nullable; -import org.joda.time.Instant; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** - * Executes a {@code GroupByKey} (GlobalWindow, default trigger, no allowed lateness). - * - *

Records arrive on the repartition topic keyed by the encoded Beam key, so every value of a key - * is co-located here. Each value is appended to a per-key buffer in a Kafka Streams state store. - * Watermark reports are fed to a {@link WatermarkAggregator}; when the input watermark reaches - * {@link BoundedWindow#TIMESTAMP_MAX_VALUE} (the end of the global window) every buffered key is - * emitted once as {@code KV>} and the buffer cleared, then the watermark is - * forwarded downstream. - * - *

Buffering whole value lists and re-encoding on each append is O(n^2) per key; fine for this - * first GroupByKey, and replaced when this moves to runner-core {@code GroupAlsoByWindow}. - */ -class GroupByKeyProcessor - implements Processor, byte[], KStreamsPayload> { - - private static final Logger LOG = LoggerFactory.getLogger(GroupByKeyProcessor.class); - - private final String stateStoreName; - // This transform's own id, stamped on every watermark it forwards downstream. - private final String transformId; - private final Coder keyCoder; - private final IterableCoder<@Nullable Object> bufferCoder; - - // Aggregates the input watermark from the upstream transform's reports, which arrive through the - // repartition topic with the upstream producer's transform id intact (the shuffle forwards - // watermark payloads unchanged). - private final WatermarkAggregator watermarkAggregator; - private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; - // The global window fires exactly once, when the watermark first reaches its end. Later watermark - // reports (e.g. the same terminal watermark broadcast across repartition partitions) must not - // re-fire. This flag is in-memory only; restart correctness comes from the state store plus - // exactly-once-v2: the buffered values and consumer offsets are committed atomically, and the - // store is empty once a key has fired, so a restart cannot double-emit. Persisting watermark - // holds is part of the separate WatermarkManager persistence work, not this initial GroupByKey. - private boolean fired = false; - - private @Nullable ProcessorContext> context; - private @Nullable KeyValueStore store; - - /** - * @param transformId this transform's own id, stamped on the watermarks it emits - * @param upstreamTransformIds the transform ids feeding this GroupByKey (known from the pipeline - * graph), whose reports the {@link WatermarkAggregator} waits for - */ - GroupByKeyProcessor( - String stateStoreName, - String transformId, - Set upstreamTransformIds, - Coder keyCoder, - Coder<@Nullable Object> valueCoder) { - this.stateStoreName = stateStoreName; - this.transformId = transformId; - this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds); - this.keyCoder = keyCoder; - this.bufferCoder = IterableCoder.of(valueCoder); - } - - @Override - public void init(ProcessorContext> context) { - this.context = context; - this.store = context.getStateStore(stateStoreName); - } - - @Override - public void process(Record> record) { - KStreamsPayload payload = record.value(); - if (payload == null) { - // The repartition topic can be written to from outside the runner (or carry a tombstone), - // so recover from the obvious error instead of crashing the task: warn and drop. - LOG.warn( - "GroupByKey {} dropping record with null payload (external write or tombstone)", - transformId); - return; - } - if (payload.isData()) { - byte[] encodedKey = record.key(); - Object element = payload.getData().getValue(); - if (encodedKey == null || element == null) { - throw new IllegalStateException("GroupByKey data record is missing its key or value"); - } - appendValue(encodedKey, element); - return; - } - watermarkAggregator.observe(payload.asWatermark()); - Instant advanced = watermarkAggregator.advance(); - if (!fired && !advanced.isBefore(BoundedWindow.TIMESTAMP_MAX_VALUE)) { - fireAll(record); - fired = true; - } - if (advanced.isAfter(lastForwardedWatermark)) { - lastForwardedWatermark = advanced; - forwardWatermark(record, advanced.getMillis()); - } - } - - private void appendValue(byte[] encodedKey, Object kvObject) { - KV kv = (KV) kvObject; - KeyValueStore kvStore = checkInitialized(store); - byte[] existing = kvStore.get(encodedKey); - List<@Nullable Object> values = existing == null ? new ArrayList<>() : decodeBuffer(existing); - values.add(kv.getValue()); - kvStore.put(encodedKey, encodeBuffer(values)); - } - - private void fireAll(Record> trigger) { - // NOTE: this emits every buffered key in a single watermark turn. For a very large key space - // that risks memory pressure and exceeding the poll / transaction timeout. Acceptable for this - // initial GlobalWindow GroupByKey (fire once at end of input); incremental, timer-driven output - // via runner-core GroupAlsoByWindow lands with the windowing/timers work. - ProcessorContext> ctx = checkInitialized(context); - KeyValueStore kvStore = checkInitialized(store); - List firedKeys = new ArrayList<>(); - try (KeyValueIterator it = kvStore.all()) { - while (it.hasNext()) { - org.apache.kafka.streams.KeyValue entry = it.next(); - Object key = decodeKey(entry.key); - List<@Nullable Object> values = decodeBuffer(entry.value); - // The pane fires at the end of the global window, so the grouped element carries the - // window's max timestamp (END_OF_GLOBAL_WINDOW). Emitting at TIMESTAMP_MIN_VALUE (the - // default of valueInGlobalWindow) would make the output appear arbitrarily late and be - // dropped downstream once the watermark has advanced. - WindowedValue>> output = - WindowedValues.timestampedValueInGlobalWindow( - KV.of(key, (Iterable<@Nullable Object>) values), - GlobalWindow.INSTANCE.maxTimestamp()); - ctx.forward( - new Record>( - entry.key, KStreamsPayload.data(output), trigger.timestamp())); - firedKeys.add(entry.key); - } - } - for (byte[] key : firedKeys) { - kvStore.delete(key); - } - } - - private void forwardWatermark(Record> trigger, long watermarkMillis) { - ProcessorContext> ctx = checkInitialized(context); - // Stamped with this transform's own id; GroupByKey is a single instance for now, so the report - // is for its only partition (0 of 1). - ctx.forward( - new Record>( - trigger.key(), - KStreamsPayload.watermark(watermarkMillis, transformId, 0, 1), - trigger.timestamp())); - } - - private byte[] encodeBuffer(List<@Nullable Object> values) { - try { - return CoderUtils.encodeToByteArray(bufferCoder, values); - } catch (CoderException e) { - throw new RuntimeException("Failed to encode GroupByKey value buffer", e); - } - } - - private List<@Nullable Object> decodeBuffer(byte[] bytes) { - try { - List<@Nullable Object> values = new ArrayList<>(); - for (@Nullable Object value : CoderUtils.decodeFromByteArray(bufferCoder, bytes)) { - values.add(value); - } - return values; - } catch (CoderException e) { - throw new RuntimeException("Failed to decode GroupByKey value buffer", e); - } - } - - private Object decodeKey(byte[] bytes) { - try { - return CoderUtils.decodeFromByteArray(keyCoder, bytes); - } catch (CoderException e) { - throw new RuntimeException("Failed to decode GroupByKey key", e); - } - } - - private static T checkInitialized(@Nullable T value) { - if (value == null) { - throw new IllegalStateException("GroupByKeyProcessor used before init()"); - } - return value; - } -} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java index 9e23dbb5cfb0..7562f68e7898 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java @@ -22,8 +22,12 @@ import org.apache.beam.model.pipeline.v1.RunnerApi; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.construction.RehydratedComponents; +import org.apache.beam.sdk.util.construction.WindowingStrategyTranslation; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.WindowedValues; +import org.apache.beam.sdk.values.WindowingStrategy; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.kafka.common.serialization.Serdes; @@ -35,10 +39,11 @@ * Translates the {@code beam:transform:group_by_key:v1} URN — the runner's first stateful, * shuffle-bearing transform. * - *

This is the simplest GroupByKey: GlobalWindow, default trigger, no allowed lateness (per the - * plan agreed with the mentor). Each key's values are buffered in a Kafka Streams state store and - * emitted once as {@code KV>} when the watermark reaches {@link - * org.apache.beam.sdk.transforms.windowing.BoundedWindow#TIMESTAMP_MAX_VALUE}. + *

Windowing and triggering are executed by Beam's {@link + * org.apache.beam.runners.core.ReduceFnRunner} inside {@link WindowedGroupByKeyProcessor}, the same + * way the Flink and Spark portable runners do it — so fixed/sliding windows, the default trigger, + * allowed lateness and timestamp combiners all work. The input PCollection's windowing strategy is + * hydrated from the pipeline proto and handed to the processor. * *

Topology added (the Beam key becomes the Kafka record key so Kafka Streams shuffles by it): * @@ -49,7 +54,8 @@ * via {@link KStreamsPayloadSerde} and a {@link GroupByKeyBroadcastPartitioner} that hashes * data by key and fans watermark reports out to every partition; *

  • a {@link Topology#addSource source} reading the repartition topic back; - *
  • the {@link GroupByKeyProcessor} plus a persistent state store, wired to the source. + *
  • the {@link WindowedGroupByKeyProcessor} plus persistent state and timer stores, wired to + * the source. * * *

    The repartition topic is expected to exist on the broker before the job starts (same @@ -62,6 +68,7 @@ class GroupByKeyTranslator implements PTransformTranslator { static final String SINK_SUFFIX = "-repartition-sink"; static final String SOURCE_SUFFIX = "-repartition-source"; static final String STATE_STORE_SUFFIX = "-state"; + static final String TIMER_STORE_SUFFIX = "-timers"; static final String REPARTITION_TOPIC_PREFIX = "__beam_gbk_"; @Override @@ -82,12 +89,16 @@ public void translate( Coder<@Nullable Object> valueCoder = (Coder<@Nullable Object>) (Coder) kvCoder.getValueCoder(); + WindowingStrategy windowingStrategy = + hydrateWindowingStrategy(pipeline, inputPCollectionId); + String parentProcessor = context.getProcessorNameForPCollection(inputPCollectionId); String shuffleName = transformId + SHUFFLE_SUFFIX; String sinkName = transformId + SINK_SUFFIX; String sourceName = transformId + SOURCE_SUFFIX; String stateStoreName = transformId + STATE_STORE_SUFFIX; + String timerStoreName = transformId + TIMER_STORE_SUFFIX; String repartitionTopic = repartitionTopic(transformId); KStreamsPayloadSerde> payloadSerde = new KStreamsPayloadSerde<>(inputCoder); @@ -111,27 +122,54 @@ public void translate( payloadSerde.deserializer(), repartitionTopic); - // Buffer values per key and fire KV> at the terminal watermark. Watermark - // reports cross the repartition topic unchanged, so they still carry the id of the transform - // that produced this GroupByKey's input — the parent the shuffle is attached to. + // Group by key and window through Beam's ReduceFnRunner, backed by the state and timer stores. + // Watermark reports cross the repartition topic unchanged, so they still carry the id of the + // transform that produced this GroupByKey's input — the parent the shuffle is attached to. topology.addProcessor( transformId, () -> - new GroupByKeyProcessor( + new WindowedGroupByKeyProcessor( stateStoreName, + timerStoreName, transformId, ImmutableSet.of(parentProcessor), keyCoder, - valueCoder), + valueCoder, + windowingStrategy, + context.getPipelineOptions()), sourceName); topology.addStateStore( Stores.keyValueStoreBuilder( Stores.persistentKeyValueStore(stateStoreName), Serdes.ByteArray(), Serdes.ByteArray()), transformId); + topology.addStateStore( + Stores.keyValueStoreBuilder( + Stores.persistentKeyValueStore(timerStoreName), Serdes.ByteArray(), Serdes.ByteArray()), + transformId); context.registerPCollectionProducer(outputPCollectionId, transformId); } + /** Hydrates the input PCollection's windowing strategy from the pipeline proto. */ + private static WindowingStrategy hydrateWindowingStrategy( + RunnerApi.Pipeline pipeline, String inputPCollectionId) { + RunnerApi.Components components = pipeline.getComponents(); + String windowingStrategyId = + components.getPcollectionsOrThrow(inputPCollectionId).getWindowingStrategyId(); + try { + @SuppressWarnings("unchecked") + WindowingStrategy strategy = + (WindowingStrategy) + WindowingStrategyTranslation.fromProto( + components.getWindowingStrategiesOrThrow(windowingStrategyId), + RehydratedComponents.forComponents(components)); + return strategy; + } catch (Exception e) { + throw new IllegalStateException( + "Failed to hydrate GroupByKey windowing strategy " + windowingStrategyId, e); + } + } + /** The internal repartition topic name for a GroupByKey transform. */ static String repartitionTopic(String transformId) { return REPARTITION_TOPIC_PREFIX + transformId.replaceAll("[^a-zA-Z0-9._-]", "_"); diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsStateInternals.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsStateInternals.java new file mode 100644 index 000000000000..d5ff37008882 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsStateInternals.java @@ -0,0 +1,393 @@ +/* + * 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.beam.runners.kafka.streams.translation; + +import java.io.ByteArrayOutputStream; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import org.apache.beam.runners.core.StateInternals; +import org.apache.beam.runners.core.StateNamespace; +import org.apache.beam.runners.core.StateTag; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.InstantCoder; +import org.apache.beam.sdk.coders.ListCoder; +import org.apache.beam.sdk.state.BagState; +import org.apache.beam.sdk.state.CombiningState; +import org.apache.beam.sdk.state.MapState; +import org.apache.beam.sdk.state.MultimapState; +import org.apache.beam.sdk.state.OrderedListState; +import org.apache.beam.sdk.state.ReadableState; +import org.apache.beam.sdk.state.SetState; +import org.apache.beam.sdk.state.State; +import org.apache.beam.sdk.state.StateBinder; +import org.apache.beam.sdk.state.StateContext; +import org.apache.beam.sdk.state.StateSpec; +import org.apache.beam.sdk.state.ValueState; +import org.apache.beam.sdk.state.WatermarkHoldState; +import org.apache.beam.sdk.transforms.Combine.CombineFn; +import org.apache.beam.sdk.transforms.CombineWithContext; +import org.apache.beam.sdk.transforms.windowing.TimestampCombiner; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.util.CombineFnUtil; +import org.apache.kafka.streams.state.KeyValueStore; +import org.checkerframework.checker.nullness.qual.NonNull; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; + +/** + * A {@link StateInternals} for one key, backed by a Kafka Streams {@link KeyValueStore}. + * + *

    Beam addresses a state cell by {@code (key, StateNamespace, StateTag)}; a windowed pipeline + * puts each window's state in its own namespace. Every cell is stored as one entry in the shared + * per-transform store under a composite byte key {@code len(key)|key | len(ns)|ns | len(tag)|tag}, + * so all cells for one Beam key share a prefix and a whole key's state can be range-scanned. The + * value is the cell's contents encoded with its Beam {@link Coder}. Writing straight to the store + * (rather than buffering and flushing) keeps this restart-safe for free: the store is changelogged + * and, under exactly-once, its writes commit atomically with the input offsets. + * + *

    Modeled on the Spark runner's {@code SparkStateInternals}; the difference is that each cell + * reads and writes its own store entry instead of an in-memory table, so there is no separate + * persist step. + */ +class KafkaStreamsStateInternals implements StateInternals { + + private final @NonNull K key; + private final byte[] encodedKey; + private final KeyValueStore store; + + KafkaStreamsStateInternals( + @NonNull K key, byte[] encodedKey, KeyValueStore store) { + this.key = key; + this.encodedKey = encodedKey; + this.store = store; + } + + @Override + public Object getKey() { + return key; + } + + @Override + public T state( + StateNamespace namespace, StateTag address, StateContext c) { + return address.getSpec().bind(address.getId(), new KafkaStreamsStateBinder(namespace, c)); + } + + /** The composite store key for one cell: {@code len|key len|namespace len|tagId}. */ + private byte[] compositeKey(StateNamespace namespace, String id) { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + writeSegment(out, encodedKey); + writeSegment(out, namespace.stringKey().getBytes(StandardCharsets.UTF_8)); + writeSegment(out, id.getBytes(StandardCharsets.UTF_8)); + return out.toByteArray(); + } + + private static void writeSegment(ByteArrayOutputStream out, byte[] segment) { + int length = segment.length; + out.write((length >>> 24) & 0xff); + out.write((length >>> 16) & 0xff); + out.write((length >>> 8) & 0xff); + out.write(length & 0xff); + out.write(segment, 0, segment.length); + } + + private class KafkaStreamsStateBinder implements StateBinder { + private final StateNamespace namespace; + private final StateContext stateContext; + + private KafkaStreamsStateBinder(StateNamespace namespace, StateContext stateContext) { + this.namespace = namespace; + this.stateContext = stateContext; + } + + @Override + public ValueState bindValue(String id, StateSpec> spec, Coder coder) { + return new KafkaStreamsValueState<>(namespace, id, coder); + } + + @Override + public BagState bindBag(String id, StateSpec> spec, Coder elemCoder) { + return new KafkaStreamsBagState<>(namespace, id, elemCoder); + } + + @Override + public SetState bindSet(String id, StateSpec> spec, Coder elemCoder) { + throw new UnsupportedOperationException( + SetState.class.getSimpleName() + " is not supported by the Kafka Streams runner yet"); + } + + @Override + public MapState bindMap( + String id, + StateSpec> spec, + Coder mapKeyCoder, + Coder mapValueCoder) { + throw new UnsupportedOperationException( + MapState.class.getSimpleName() + " is not supported by the Kafka Streams runner yet"); + } + + @Override + public MultimapState bindMultimap( + String id, + StateSpec> spec, + Coder keyCoder, + Coder valueCoder) { + throw new UnsupportedOperationException( + MultimapState.class.getSimpleName() + + " is not supported by the Kafka Streams runner yet"); + } + + @Override + public OrderedListState bindOrderedList( + String id, StateSpec> spec, Coder elemCoder) { + throw new UnsupportedOperationException( + OrderedListState.class.getSimpleName() + + " is not supported by the Kafka Streams runner yet"); + } + + @Override + public CombiningState bindCombining( + String id, + StateSpec> spec, + Coder accumCoder, + CombineFn combineFn) { + return new KafkaStreamsCombiningState<>(namespace, id, accumCoder, combineFn); + } + + @Override + public + CombiningState bindCombiningWithContext( + String id, + StateSpec> spec, + Coder accumCoder, + CombineWithContext.CombineFnWithContext combineFn) { + return new KafkaStreamsCombiningState<>( + namespace, id, accumCoder, CombineFnUtil.bindContext(combineFn, stateContext)); + } + + @Override + public WatermarkHoldState bindWatermark( + String id, StateSpec spec, TimestampCombiner timestampCombiner) { + return new KafkaStreamsWatermarkHoldState(namespace, id, timestampCombiner); + } + } + + /** Common read/write/clear against the backing store for one cell. */ + private abstract class AbstractState { + final StateNamespace namespace; + final String id; + final Coder coder; + + AbstractState(StateNamespace namespace, String id, Coder coder) { + this.namespace = namespace; + this.id = id; + this.coder = coder; + } + + @Nullable + T readValue() { + byte[] bytes = store.get(compositeKey(namespace, id)); + if (bytes == null) { + return null; + } + try { + return CoderUtils.decodeFromByteArray(coder, bytes); + } catch (CoderException e) { + throw new RuntimeException("Failed to decode state " + id, e); + } + } + + void writeValue(T input) { + try { + store.put(compositeKey(namespace, id), CoderUtils.encodeToByteArray(coder, input)); + } catch (CoderException e) { + throw new RuntimeException("Failed to encode state " + id, e); + } + } + + public void clear() { + store.delete(compositeKey(namespace, id)); + } + + ReadableState isEmptyState() { + return new ReadableState() { + @Override + public Boolean read() { + return store.get(compositeKey(namespace, id)) == null; + } + + @Override + public ReadableState readLater() { + return this; + } + }; + } + } + + private class KafkaStreamsValueState extends AbstractState implements ValueState { + KafkaStreamsValueState(StateNamespace namespace, String id, Coder coder) { + super(namespace, id, coder); + } + + @Override + public KafkaStreamsValueState readLater() { + return this; + } + + @Override + public @Nullable T read() { + return readValue(); + } + + @Override + public void write(T input) { + writeValue(input); + } + } + + private class KafkaStreamsBagState extends AbstractState> implements BagState { + KafkaStreamsBagState(StateNamespace namespace, String id, Coder elemCoder) { + super(namespace, id, ListCoder.of(elemCoder)); + } + + @Override + public KafkaStreamsBagState readLater() { + return this; + } + + @Override + public Iterable read() { + List value = readValue(); + return value == null ? new ArrayList<>() : value; + } + + @Override + public void add(T input) { + List value = readValue(); + if (value == null) { + value = new ArrayList<>(); + } + value.add(input); + writeValue(value); + } + + @Override + public ReadableState isEmpty() { + return isEmptyState(); + } + } + + private class KafkaStreamsWatermarkHoldState extends AbstractState + implements WatermarkHoldState { + private final TimestampCombiner timestampCombiner; + + KafkaStreamsWatermarkHoldState( + StateNamespace namespace, String id, TimestampCombiner timestampCombiner) { + super(namespace, id, InstantCoder.of()); + this.timestampCombiner = timestampCombiner; + } + + @Override + public KafkaStreamsWatermarkHoldState readLater() { + return this; + } + + // GroupingState.read() is typed non-null, but an empty hold reads back null. Beam's state + // interfaces are under-annotated here (https://github.com/apache/beam/issues/20497), which is + // why the Spark and Flink StateInternals suppress nullness for the whole class; this runner + // narrows the suppression to just this method. + @Override + @SuppressWarnings("nullness") + public Instant read() { + return readValue(); + } + + @Override + public void add(Instant outputTime) { + Instant current = readValue(); + writeValue(current == null ? outputTime : timestampCombiner.combine(current, outputTime)); + } + + @Override + public ReadableState isEmpty() { + return isEmptyState(); + } + + @Override + public TimestampCombiner getTimestampCombiner() { + return timestampCombiner; + } + } + + @SuppressWarnings("TypeParameterShadowing") + private class KafkaStreamsCombiningState extends AbstractState + implements CombiningState { + private final CombineFn combineFn; + + KafkaStreamsCombiningState( + StateNamespace namespace, + String id, + Coder accumCoder, + CombineFn combineFn) { + super(namespace, id, accumCoder); + this.combineFn = combineFn; + } + + @Override + public KafkaStreamsCombiningState readLater() { + return this; + } + + // GroupingState.read() is typed non-null but a CombineFn may extract a null output; the same + // under-annotation as WatermarkHoldState.read() (https://github.com/apache/beam/issues/20497). + @Override + @SuppressWarnings("nullness") + public OutputT read() { + return combineFn.extractOutput(getAccum()); + } + + @Override + public void add(InputT input) { + writeValue(combineFn.addInput(getAccum(), input)); + } + + @Override + public AccumT getAccum() { + AccumT accum = readValue(); + return accum == null ? combineFn.createAccumulator() : accum; + } + + @Override + public void addAccum(AccumT accum) { + writeValue(combineFn.mergeAccumulators(Arrays.asList(getAccum(), accum))); + } + + @Override + public AccumT mergeAccumulators(Iterable accumulators) { + return combineFn.mergeAccumulators(accumulators); + } + + @Override + public ReadableState isEmpty() { + return isEmptyState(); + } + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternals.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternals.java new file mode 100644 index 000000000000..8a5bf910298c --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternals.java @@ -0,0 +1,180 @@ +/* + * 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.beam.runners.kafka.streams.translation; + +import java.io.ByteArrayOutputStream; +import java.nio.charset.StandardCharsets; +import org.apache.beam.runners.core.StateNamespace; +import org.apache.beam.runners.core.TimerInternals; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.state.TimeDomain; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.kafka.streams.state.KeyValueStore; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; + +/** + * A {@link TimerInternals} for one key, backed by a shared Kafka Streams timer {@link + * KeyValueStore} for a GroupByKey. + * + *

    Kafka Streams has no per-key timer service, so timers are persisted like any other state. A + * timer is stored under a key built from its identity — {@code key | domain | timerFamily | timerId + * | namespace} — so setting a timer with the same identity overwrites the previous one and deleting + * it removes exactly that entry, matching {@link TimerInternals}' contract. The value is the {@link + * TimerData} encoded with {@link TimerInternals.TimerDataCoderV2}, which carries the fire time. + * + *

    Firing is driven by {@link WindowedGroupByKeyProcessor}: on a watermark advance it scans the + * store for event-time timers whose fire time has passed and replays them through {@link + * org.apache.beam.runners.core.ReduceFnRunner#onTimers}. Scanning the whole store per advance is + * O(timers); acceptable for this first windowing support and replaceable with a fire-time-ordered + * index later, the same way the initial GroupByKey buffered eagerly. + * + *

    This instance reports the current input-watermark and processing times it was constructed with + * (the caller advances them as reports arrive); it never fires timers itself. + */ +class KafkaStreamsTimerInternals implements TimerInternals { + + private final byte[] encodedKey; + private final KeyValueStore timerStore; + private final TimerInternals.TimerDataCoderV2 timerCoder; + private final Instant inputWatermarkTime; + private final Instant processingTime; + + KafkaStreamsTimerInternals( + byte[] encodedKey, + KeyValueStore timerStore, + Coder windowCoder, + Instant inputWatermarkTime, + Instant processingTime) { + this.encodedKey = encodedKey; + this.timerStore = timerStore; + this.timerCoder = TimerInternals.TimerDataCoderV2.of(windowCoder); + this.inputWatermarkTime = inputWatermarkTime; + this.processingTime = processingTime; + } + + @Override + public void setTimer( + StateNamespace namespace, + String timerId, + String timerFamilyId, + Instant target, + Instant outputTimestamp, + TimeDomain timeDomain) { + setTimer(TimerData.of(timerId, timerFamilyId, namespace, target, outputTimestamp, timeDomain)); + } + + @Override + public void setTimer(TimerData timerData) { + timerStore.put(timerStoreKey(encodedKey, timerData), encodeTimer(timerData)); + } + + @Override + public void deleteTimer( + StateNamespace namespace, String timerId, String timerFamilyId, TimeDomain timeDomain) { + timerStore.delete(timerStoreKey(encodedKey, timerId, timerFamilyId, timeDomain, namespace)); + } + + @Override + public void deleteTimer(StateNamespace namespace, String timerId, String timerFamilyId) { + throw new UnsupportedOperationException( + "Deleting a timer without a time domain is not supported; the domain is part of a timer's" + + " store identity."); + } + + @Override + public void deleteTimer(TimerData timerKey) { + timerStore.delete(timerStoreKey(encodedKey, timerKey)); + } + + @Override + public Instant currentProcessingTime() { + return processingTime; + } + + @Override + public @Nullable Instant currentSynchronizedProcessingTime() { + return null; + } + + @Override + public Instant currentInputWatermarkTime() { + return inputWatermarkTime; + } + + @Override + public @Nullable Instant currentOutputWatermarkTime() { + return null; + } + + private byte[] encodeTimer(TimerData timerData) { + try { + return CoderUtils.encodeToByteArray(timerCoder, timerData); + } catch (CoderException e) { + throw new RuntimeException("Failed to encode timer " + timerData, e); + } + } + + /** The store key for a timer's identity. Package-visible so the processor can build/scan it. */ + static byte[] timerStoreKey(byte[] encodedKey, TimerData timerData) { + return timerStoreKey( + encodedKey, + timerData.getTimerId(), + timerData.getTimerFamilyId(), + timerData.getDomain(), + timerData.getNamespace()); + } + + static byte[] timerStoreKey( + byte[] encodedKey, + String timerId, + String timerFamilyId, + TimeDomain domain, + StateNamespace namespace) { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + writeSegment(out, encodedKey); + writeSegment(out, new byte[] {(byte) domain.ordinal()}); + writeSegment(out, timerFamilyId.getBytes(StandardCharsets.UTF_8)); + writeSegment(out, timerId.getBytes(StandardCharsets.UTF_8)); + writeSegment(out, namespace.stringKey().getBytes(StandardCharsets.UTF_8)); + return out.toByteArray(); + } + + /** Reads the encoded Beam key (the first segment) back out of a timer store key. */ + static byte[] encodedKeyOf(byte[] storeKey) { + int length = + ((storeKey[0] & 0xff) << 24) + | ((storeKey[1] & 0xff) << 16) + | ((storeKey[2] & 0xff) << 8) + | (storeKey[3] & 0xff); + byte[] key = new byte[length]; + System.arraycopy(storeKey, 4, key, 0, length); + return key; + } + + private static void writeSegment(ByteArrayOutputStream out, byte[] segment) { + int length = segment.length; + out.write((length >>> 24) & 0xff); + out.write((length >>> 16) & 0xff); + out.write((length >>> 8) & 0xff); + out.write(length & 0xff); + out.write(segment, 0, segment.length); + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java new file mode 100644 index 000000000000..b9ff98da0724 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java @@ -0,0 +1,284 @@ +/* + * 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.beam.runners.kafka.streams.translation; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.core.NullSideInputReader; +import org.apache.beam.runners.core.ReduceFnRunner; +import org.apache.beam.runners.core.StateInternals; +import org.apache.beam.runners.core.SystemReduceFn; +import org.apache.beam.runners.core.TimerInternals.TimerData; +import org.apache.beam.runners.core.triggers.ExecutableTriggerStateMachine; +import org.apache.beam.runners.core.triggers.TriggerStateMachines; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.state.TimeDomain; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.util.construction.TriggerTranslation; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowingStrategy; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.BaseEncoding; +import org.apache.kafka.streams.processor.api.Processor; +import org.apache.kafka.streams.processor.api.ProcessorContext; +import org.apache.kafka.streams.processor.api.Record; +import org.apache.kafka.streams.state.KeyValueIterator; +import org.apache.kafka.streams.state.KeyValueStore; +import org.checkerframework.checker.nullness.qual.NonNull; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Executes a windowed {@code GroupByKey} by driving Beam's {@link ReduceFnRunner} — the same + * windowing + triggering state machine the Flink and Spark portable runners use — with Kafka + * Streams state and timers behind it. + * + *

    Records arrive on the repartition topic keyed by the encoded Beam key, so every value of a key + * is co-located here. For each data record this builds a {@link ReduceFnRunner} for that key over a + * {@link KafkaStreamsStateInternals} and {@link KafkaStreamsTimerInternals} (both backed by + * persistent stores) and feeds the element in; the runner assigns it to windows, updates the + * trigger state, and sets any timers it needs. When the aggregated input watermark advances, every + * event-time timer whose fire time has passed is replayed through {@code onTimers}, which is what + * makes windows emit their panes. The runner is stateless between records — all durable state lives + * in the stores — so a fresh one per record is correct, mirroring {@code + * GroupAlsoByWindowViaWindowSetNewDoFn}. + * + *

    This first version supports the default trigger and non-merging windows well; richer triggers, + * processing-time timers and session (merging) windows build on the same machinery in a follow-up. + */ +class WindowedGroupByKeyProcessor + implements Processor, byte[], KStreamsPayload> { + + private static final Logger LOG = LoggerFactory.getLogger(WindowedGroupByKeyProcessor.class); + + private final String stateStoreName; + private final String timerStoreName; + private final String transformId; + private final Coder keyCoder; + private final WindowingStrategy windowingStrategy; + private final Coder windowCoder; + private final RunnerApi.Trigger triggerProto; + private final SystemReduceFn, Iterable, W> reduceFn; + private final PipelineOptions options; + + private final WatermarkAggregator watermarkAggregator; + private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; + private Instant inputWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; + + private @Nullable ProcessorContext> context; + private @Nullable KeyValueStore stateStore; + private @Nullable KeyValueStore timerStore; + + WindowedGroupByKeyProcessor( + String stateStoreName, + String timerStoreName, + String transformId, + Set upstreamTransformIds, + Coder keyCoder, + Coder valueCoder, + WindowingStrategy windowingStrategy, + PipelineOptions options) { + this.stateStoreName = stateStoreName; + this.timerStoreName = timerStoreName; + this.transformId = transformId; + this.keyCoder = keyCoder; + this.windowingStrategy = windowingStrategy; + this.windowCoder = windowingStrategy.getWindowFn().windowCoder(); + this.triggerProto = TriggerTranslation.toProto(windowingStrategy.getTrigger()); + this.reduceFn = SystemReduceFn.buffering(valueCoder); + this.options = options; + this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds); + } + + @Override + public void init(ProcessorContext> context) { + this.context = context; + this.stateStore = context.getStateStore(stateStoreName); + this.timerStore = context.getStateStore(timerStoreName); + } + + @Override + public void process(Record> record) { + KStreamsPayload payload = record.value(); + if (payload == null) { + LOG.warn( + "GroupByKey {} dropping record with null payload (external write or tombstone)", + transformId); + return; + } + if (payload.isData()) { + processData(record, payload); + return; + } + watermarkAggregator.observe(payload.asWatermark()); + Instant advanced = watermarkAggregator.advance(); + if (advanced.isAfter(lastForwardedWatermark)) { + inputWatermark = advanced; + fireDueEventTimeTimers(record, advanced); + lastForwardedWatermark = advanced; + forwardWatermark(record, advanced.getMillis()); + } + } + + private void processData(Record> record, KStreamsPayload payload) { + byte[] encodedKey = record.key(); + if (encodedKey == null) { + throw new IllegalStateException("GroupByKey data record is missing its key"); + } + @SuppressWarnings("unchecked") + WindowedValue> element = (WindowedValue>) payload.getData(); + K key = decodeKey(encodedKey); + WindowedValue valueElement = element.withValue(element.getValue().getValue()); + runReduceFn( + record, encodedKey, key, Collections.singletonList(valueElement), Collections.emptyList()); + } + + /** Fires every event-time timer whose fire time is at or before the new input watermark. */ + private void fireDueEventTimeTimers( + Record> record, Instant watermark) { + KeyValueStore timers = checkInitialized(timerStore); + // Group due timers by the Beam key they belong to; a key's timers fire together in one + // ReduceFnRunner turn. Whole-store scan per advance is O(timers) — see the class doc. + Map dueByKey = new LinkedHashMap<>(); + List firedStoreKeys = new ArrayList<>(); + try (KeyValueIterator it = timers.all()) { + while (it.hasNext()) { + org.apache.kafka.streams.KeyValue entry = it.next(); + TimerData timer = decodeTimer(entry.value); + if (timer.getDomain() != TimeDomain.EVENT_TIME || timer.getTimestamp().isAfter(watermark)) { + continue; + } + byte[] encodedKey = KafkaStreamsTimerInternals.encodedKeyOf(entry.key); + dueByKey + .computeIfAbsent( + BaseEncoding.base16().encode(encodedKey), k -> new DueTimers(encodedKey)) + .timers + .add(timer); + firedStoreKeys.add(entry.key); + } + } + // Remove fired timers before replaying them; onTimers may legitimately set new ones. + for (byte[] storeKey : firedStoreKeys) { + timers.delete(storeKey); + } + for (DueTimers due : dueByKey.values()) { + runReduceFn( + record, due.encodedKey, decodeKey(due.encodedKey), Collections.emptyList(), due.timers); + } + } + + private void runReduceFn( + Record> record, + byte[] encodedKey, + @NonNull K key, + List> elements, + List timers) { + StateInternals stateInternals = + new KafkaStreamsStateInternals<>(key, encodedKey, checkInitialized(stateStore)); + KafkaStreamsTimerInternals timerInternals = + new KafkaStreamsTimerInternals( + encodedKey, checkInitialized(timerStore), windowCoder, inputWatermark, Instant.now()); + ReduceFnRunner, W> runner = + new ReduceFnRunner<>( + key, + windowingStrategy, + ExecutableTriggerStateMachine.create( + TriggerStateMachines.stateMachineForTrigger(triggerProto)), + stateInternals, + timerInternals, + output -> forwardData(record, encodedKey, output), + NullSideInputReader.empty(), + reduceFn, + options); + try { + runner.processElements(elements); + runner.onTimers(timers); + runner.persist(); + } catch (Exception e) { + throw new RuntimeException("GroupByKey " + transformId + " failed to run windowing", e); + } + } + + private void forwardData( + Record> trigger, + byte[] encodedKey, + WindowedValue>> output) { + ProcessorContext> ctx = checkInitialized(context); + ctx.forward( + new Record>( + encodedKey, KStreamsPayload.data(output), trigger.timestamp())); + } + + private void forwardWatermark(Record> trigger, long watermarkMillis) { + ProcessorContext> ctx = checkInitialized(context); + // Stamped with this transform's own id; GroupByKey is a single instance for now (0 of 1). + ctx.forward( + new Record>( + trigger.key(), + KStreamsPayload.watermark(watermarkMillis, transformId, 0, 1), + trigger.timestamp())); + } + + private @NonNull K decodeKey(byte[] bytes) { + try { + K key = CoderUtils.decodeFromByteArray(keyCoder, bytes); + if (key == null) { + throw new IllegalStateException("GroupByKey key decoded to null"); + } + return key; + } catch (CoderException e) { + throw new RuntimeException("Failed to decode GroupByKey key", e); + } + } + + private TimerData decodeTimer(byte[] bytes) { + try { + return CoderUtils.decodeFromByteArray( + org.apache.beam.runners.core.TimerInternals.TimerDataCoderV2.of(windowCoder), bytes); + } catch (CoderException e) { + throw new RuntimeException("Failed to decode timer", e); + } + } + + private static T checkInitialized(@Nullable T value) { + if (value == null) { + throw new IllegalStateException("WindowedGroupByKeyProcessor used before init()"); + } + return value; + } + + /** The event-time timers due for one Beam key, plus that key's encoded bytes. */ + private static final class DueTimers { + final byte[] encodedKey; + final List timers = new ArrayList<>(); + + DueTimers(byte[] encodedKey) { + this.encodedKey = encodedKey; + } + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTestRunner.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTestRunner.java index 8ab29182c281..bdbde1db36dd 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTestRunner.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTestRunner.java @@ -18,11 +18,7 @@ package org.apache.beam.runners.kafka.streams; import java.time.Duration; -import java.util.ArrayList; -import java.util.HashSet; -import java.util.List; import java.util.Properties; -import java.util.Set; import java.util.UUID; import org.apache.beam.model.pipeline.v1.RunnerApi; import org.apache.beam.runners.core.metrics.MetricsContainerStepMap; @@ -39,16 +35,11 @@ import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; import org.apache.beam.sdk.util.construction.PipelineTranslation; import org.apache.beam.sdk.util.construction.SplittableParDo; -import org.apache.kafka.common.serialization.ByteArrayDeserializer; -import org.apache.kafka.common.serialization.ByteArraySerializer; import org.apache.kafka.common.serialization.Serdes; import org.apache.kafka.streams.StreamsConfig; -import org.apache.kafka.streams.TestInputTopic; -import org.apache.kafka.streams.TestOutputTopic; import org.apache.kafka.streams.Topology; import org.apache.kafka.streams.TopologyDescription; import org.apache.kafka.streams.TopologyTestDriver; -import org.apache.kafka.streams.test.TestRecord; /** * Test harness that runs a Beam {@link Pipeline} through the Kafka Streams runner's translation and @@ -58,16 +49,13 @@ * effects (e.g. a {@code SharedTestCollector} written by a recording DoFn) have completed when it * returns. * - *

    {@link TopologyTestDriver} does not loop a low-level sink topic back into its source, so an - * internal repartition topic (one that is both a sink and a source in the topology — e.g. the one - * GroupByKey introduces) would otherwise dead-end. {@link #run(Pipeline)} discovers those topics - * from the {@link TopologyDescription} and round-trips them until no more records flow, standing in - * for the broker. + *

    {@link TopologyTestDriver} loops each internal repartition topic (one that is both a sink and + * a source in the topology — e.g. the one GroupByKey introduces) from its sink back to its source + * within a single driver step, so advancing the wall clock is enough to drive the whole pipeline to + * completion; no manual broker simulation is needed. */ public final class KafkaStreamsTestRunner { - private static final int MAX_ROUND_TRIPS = 100; - private KafkaStreamsTestRunner() {} /** Pipeline options for a Kafka Streams runner test: the EMBEDDED harness and a unique app id. */ @@ -118,10 +106,12 @@ public static MetricResults run(Pipeline pipeline) { KafkaStreamsTranslationContext context = translate(pipeline); Topology topology = context.getTopology(); try (TopologyTestDriver driver = new TopologyTestDriver(topology, streamsConfig(pipeline))) { - // Fire the Impulse wall-clock punctuator and let the initial records flow. + // Fire the Impulse wall-clock punctuator; TopologyTestDriver then flows the records through + // the whole topology, including looping each internal repartition topic (a sink that is also + // a source, e.g. the one GroupByKey introduces) back to its source, standing in for the + // broker. A second advance covers punctuators that need a later tick. driver.advanceWallClockTime(Duration.ofSeconds(1)); driver.advanceWallClockTime(Duration.ofSeconds(1)); - roundTripInternalTopics(driver, internalTopics(topology)); } return MetricsContainerStepMap.asAttemptedOnlyMetricResults( context.getMetricsContainerStepMap()); @@ -143,81 +133,6 @@ public static String findAnyLeafProcessorName(Topology topology) { throw new IllegalStateException("no leaf processor found in topology"); } - /** Repartition/internal topics are the ones that appear as both a sink and a source. */ - private static Set internalTopics(Topology topology) { - Set sinkTopics = new HashSet<>(); - Set sourceTopics = new HashSet<>(); - for (TopologyDescription.Subtopology subtopology : topology.describe().subtopologies()) { - for (TopologyDescription.Node node : subtopology.nodes()) { - if (node instanceof TopologyDescription.Sink) { - String topic = ((TopologyDescription.Sink) node).topic(); - if (topic != null) { - sinkTopics.add(topic); - } - } else if (node instanceof TopologyDescription.Source) { - sourceTopics.addAll(((TopologyDescription.Source) node).topicSet()); - } - } - } - sinkTopics.retainAll(sourceTopics); - return sinkTopics; - } - - /** - * Simulates the broker for internal repartition topics. - * - *

    The runner shuffles data (and the watermark) through internal topics that a processor both - * writes to (a sink) and reads back from (a source) — e.g. the topic GroupByKey introduces to - * partition by key. On a real broker those records make the round trip automatically, but {@link - * TopologyTestDriver} does not connect a sink back to a source, so the downstream half of the - * topology would never see them. This drains what each internal topic's sink wrote and pipes it - * into that topic's source, repeating until nothing new flows (a fixpoint), which stands in for - * the broker and lets the pipeline run to completion. - */ - private static void roundTripInternalTopics(TopologyTestDriver driver, Set topics) { - // Create the sink-output and source-input handles once and reuse them across rounds; a single - // TestOutputTopic keeps returning newly produced records on each read. - List roundTrips = new ArrayList<>(); - for (String topic : topics) { - roundTrips.add( - new TopicRoundTrip( - driver.createOutputTopic( - topic, new ByteArrayDeserializer(), new ByteArrayDeserializer()), - driver.createInputTopic( - topic, new ByteArraySerializer(), new ByteArraySerializer()))); - } - - for (int round = 0; round < MAX_ROUND_TRIPS; round++) { - boolean progressed = false; - for (TopicRoundTrip roundTrip : roundTrips) { - List> records = roundTrip.output.readRecordsToList(); - if (records.isEmpty()) { - continue; - } - progressed = true; - for (TestRecord record : records) { - roundTrip.input.pipeInput(record); - } - } - if (!progressed) { - return; - } - } - throw new IllegalStateException( - "Internal topics did not reach quiescence after " + MAX_ROUND_TRIPS + " round trips"); - } - - /** The reusable sink-output and source-input handles for one internal topic. */ - private static final class TopicRoundTrip { - final TestOutputTopic output; - final TestInputTopic input; - - TopicRoundTrip(TestOutputTopic output, TestInputTopic input) { - this.output = output; - this.input = input; - } - } - /** Kafka Streams config for a {@link TopologyTestDriver} built from the pipeline's app id. */ public static Properties streamsConfig(Pipeline pipeline) { String applicationId = diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FixedWindowGroupByKeyTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FixedWindowGroupByKeyTest.java new file mode 100644 index 000000000000..73f4dac601eb --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FixedWindowGroupByKeyTest.java @@ -0,0 +1,106 @@ +/* + * 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.beam.runners.kafka.streams.translation; + +import static org.hamcrest.CoreMatchers.hasItems; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.MatcherAssert.assertThat; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.GroupByKey; +import org.apache.beam.sdk.transforms.Impulse; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.windowing.FixedWindows; +import org.apache.beam.sdk.transforms.windowing.Window; +import org.apache.beam.sdk.values.KV; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.junit.Test; + +/** + * End-to-end test that GroupByKey groups per fixed window, not just per key: {@code Impulse -> emit + * timestamped KVs -> Window.into(FixedWindows) -> GroupByKey -> record groups}. + * + *

    The same key "a" has values in two different windows, so a correct windowed GroupByKey emits + * two groups for it (one per window) rather than one combined group. This exercises the {@link + * WindowedGroupByKeyProcessor} path (ReduceFnRunner over the Kafka Streams state and timer stores) + * that the earlier global-window GroupByKey did not. + */ +public class FixedWindowGroupByKeyTest { + + private static final Duration WINDOW_SIZE = Duration.millis(10); + + /** Emits KVs whose timestamps fall into two adjacent fixed windows. */ + private static class EmitTimestampedKvsFn extends DoFn> { + @ProcessElement + public void processElement(OutputReceiver> out) { + // Window [0, 10): a=1, a=2, b=5. + out.outputWithTimestamp(KV.of("a", 1), new Instant(1)); + out.outputWithTimestamp(KV.of("a", 2), new Instant(2)); + out.outputWithTimestamp(KV.of("b", 5), new Instant(3)); + // Window [10, 20): a=3. + out.outputWithTimestamp(KV.of("a", 3), new Instant(15)); + } + } + + /** Records each grouped result as {@code "key=[sorted values]"}. */ + private static class RecordGroupFn extends DoFn>, Void> { + private final SharedTestCollector collector; + + RecordGroupFn(SharedTestCollector collector) { + this.collector = collector; + } + + @ProcessElement + public void processElement(@Element KV> group) { + List values = new ArrayList<>(); + group.getValue().forEach(values::add); + Collections.sort(values); + collector.record(group.getKey() + "=" + values); + } + } + + @Test + public void groupsValuesPerFixedWindow() { + try (SharedTestCollector collector = SharedTestCollector.create()) { + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); + pipeline + .apply("impulse", Impulse.create()) + .apply("emit", ParDo.of(new EmitTimestampedKvsFn())) + .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())) + .apply("window", Window.into(FixedWindows.of(WINDOW_SIZE))) + .apply("gbk", GroupByKey.create()) + .apply("record", ParDo.of(new RecordGroupFn(collector))); + + KafkaStreamsTestRunner.run(pipeline); + + List groups = collector.recorded(); + // a splits across two windows -> two groups; b has one; three groups total. + assertThat(groups.size(), is(3)); + assertThat(groups, hasItems("a=[1, 2]", "a=[3]", "b=[5]")); + } + } +} From f1811cac85ff124b63d2e9bd30770d7692ba2cdc Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Mon, 27 Jul 2026 16:00:46 +0500 Subject: [PATCH 2/2] Address review: index timers by fire time, apply watermark holds Adds a fire-time index so due timers are found by a range scan instead of a scan of every timer of every key, which was the main review concern. Timers now live in two stores: the identity store, keyed by key/domain/family/id/ namespace, which is how a timer is overwritten or deleted and whose value is the index key; and the index store, keyed by domain, fire timestamp and identity, which the processor range-scans over (-inf, watermark] to find exactly the timers that are due. Timestamps are written sign-flipped big-endian so the unsigned byte order Kafka Streams compares keys by matches numeric order, including for negative timestamps. Applies watermark holds when publishing the watermark. ReduceFnRunner takes a hold for a buffered pane that has not been emitted; the GroupByKey now mirrors those holds into a hold-time-ordered index and publishes min(inputWatermark, minimum hold), so downstream is not told that nothing earlier is coming while panes are still buffered. Holds are read after firing, since firing releases them. Builds the composite state key into one exactly-sized array with arraycopy and caches the key/namespace prefix, which a turn of the windowing runner reuses across the several tags it touches. Keeps the length prefixes rather than a separator byte: an encoded key and a namespace string can both contain any byte, so key=a/b,ns=c and key=a,ns=b/c would otherwise collide, and the timer scan reads the encoded key back out of a key. Documents why currentSynchronizedProcessingTime returns null and what would supply it; currentOutputWatermarkTime now returns the watermark actually published instead of null. Adds KafkaStreamsTimerInternalsTest covering the due-timer range scan, negative timestamps, timer replacement and deletion, timers of different keys sharing the index, and processing-time timers being excluded. --- .../translation/GroupByKeyTranslator.java | 20 ++ .../KafkaStreamsStateInternals.java | 102 +++++++-- .../KafkaStreamsTimerInternals.java | 189 +++++++++++----- .../kafka/streams/translation/StoreKeys.java | 105 +++++++++ .../WindowedGroupByKeyProcessor.java | 99 ++++++--- .../KafkaStreamsTimerInternalsTest.java | 208 ++++++++++++++++++ 6 files changed, 625 insertions(+), 98 deletions(-) create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StoreKeys.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternalsTest.java diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java index 7562f68e7898..c5327e28e069 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java @@ -68,7 +68,9 @@ class GroupByKeyTranslator implements PTransformTranslator { static final String SINK_SUFFIX = "-repartition-sink"; static final String SOURCE_SUFFIX = "-repartition-source"; static final String STATE_STORE_SUFFIX = "-state"; + static final String HOLDS_INDEX_STORE_SUFFIX = "-holds-index"; static final String TIMER_STORE_SUFFIX = "-timers"; + static final String TIMER_INDEX_STORE_SUFFIX = "-timers-index"; static final String REPARTITION_TOPIC_PREFIX = "__beam_gbk_"; @Override @@ -98,7 +100,9 @@ public void translate( String sinkName = transformId + SINK_SUFFIX; String sourceName = transformId + SOURCE_SUFFIX; String stateStoreName = transformId + STATE_STORE_SUFFIX; + String holdsIndexStoreName = transformId + HOLDS_INDEX_STORE_SUFFIX; String timerStoreName = transformId + TIMER_STORE_SUFFIX; + String timerIndexStoreName = transformId + TIMER_INDEX_STORE_SUFFIX; String repartitionTopic = repartitionTopic(transformId); KStreamsPayloadSerde> payloadSerde = new KStreamsPayloadSerde<>(inputCoder); @@ -130,7 +134,9 @@ public void translate( () -> new WindowedGroupByKeyProcessor( stateStoreName, + holdsIndexStoreName, timerStoreName, + timerIndexStoreName, transformId, ImmutableSet.of(parentProcessor), keyCoder, @@ -146,6 +152,20 @@ public void translate( Stores.keyValueStoreBuilder( Stores.persistentKeyValueStore(timerStoreName), Serdes.ByteArray(), Serdes.ByteArray()), transformId); + // Indexes ordered by timestamp, so due timers and the minimum watermark hold are range scans + // rather than scans of every timer or every held window. + topology.addStateStore( + Stores.keyValueStoreBuilder( + Stores.persistentKeyValueStore(timerIndexStoreName), + Serdes.ByteArray(), + Serdes.ByteArray()), + transformId); + topology.addStateStore( + Stores.keyValueStoreBuilder( + Stores.persistentKeyValueStore(holdsIndexStoreName), + Serdes.ByteArray(), + Serdes.ByteArray()), + transformId); context.registerPCollectionProducer(outputPCollectionId, transformId); } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsStateInternals.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsStateInternals.java index d5ff37008882..c17d7c45ad4d 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsStateInternals.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsStateInternals.java @@ -17,7 +17,6 @@ */ package org.apache.beam.runners.kafka.streams.translation; -import java.io.ByteArrayOutputStream; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Arrays; @@ -47,6 +46,7 @@ import org.apache.beam.sdk.transforms.windowing.TimestampCombiner; import org.apache.beam.sdk.util.CoderUtils; import org.apache.beam.sdk.util.CombineFnUtil; +import org.apache.kafka.streams.state.KeyValueIterator; import org.apache.kafka.streams.state.KeyValueStore; import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; @@ -69,15 +69,46 @@ */ class KafkaStreamsStateInternals implements StateInternals { + /** The holds index is a set; only its keys carry information. */ + private static final byte[] EMPTY_VALUE = new byte[0]; + + /** + * Reads the minimum watermark hold held by any key and window, or {@code null} if none is held. + * The index is ordered by hold time, so this is the first entry rather than a scan. + */ + static @Nullable Instant minWatermarkHold(KeyValueStore holdsIndexStore) { + try (KeyValueIterator it = holdsIndexStore.all()) { + if (!it.hasNext()) { + return null; + } + return new Instant(StoreKeys.readTimestamp(it.next().key, 0)); + } + } + private final @NonNull K key; private final byte[] encodedKey; private final KeyValueStore store; + private final KeyValueStore holdsIndexStore; + + /** + * The last namespace a composite key was built for, and the {@code key | namespace} prefix that + * was built for it. One turn of the windowing runner touches several tags in the same namespace + * back to back (read the buffer, read the hold, write both), so caching the prefix removes most + * of the per-access encoding work. + */ + private @Nullable StateNamespace cachedNamespace; + + private byte @Nullable [] cachedPrefix; KafkaStreamsStateInternals( - @NonNull K key, byte[] encodedKey, KeyValueStore store) { + @NonNull K key, + byte[] encodedKey, + KeyValueStore store, + KeyValueStore holdsIndexStore) { this.key = key; this.encodedKey = encodedKey; this.store = store; + this.holdsIndexStore = holdsIndexStore; } @Override @@ -91,22 +122,35 @@ public T state( return address.getSpec().bind(address.getId(), new KafkaStreamsStateBinder(namespace, c)); } - /** The composite store key for one cell: {@code len|key len|namespace len|tagId}. */ + /** + * The composite store key for one cell: {@code len|key len|namespace len|tagId}. + * + *

    Built into one exactly-sized array, reusing the cached {@code key | namespace} prefix. This + * runs on every state access, so it avoids the repeated growth and final copy a stream would do. + */ private byte[] compositeKey(StateNamespace namespace, String id) { - ByteArrayOutputStream out = new ByteArrayOutputStream(); - writeSegment(out, encodedKey); - writeSegment(out, namespace.stringKey().getBytes(StandardCharsets.UTF_8)); - writeSegment(out, id.getBytes(StandardCharsets.UTF_8)); - return out.toByteArray(); + byte[] prefix = prefixFor(namespace); + byte[] idBytes = id.getBytes(StandardCharsets.UTF_8); + byte[] compositeKey = new byte[prefix.length + StoreKeys.segmentLength(idBytes)]; + System.arraycopy(prefix, 0, compositeKey, 0, prefix.length); + StoreKeys.writeSegment(compositeKey, prefix.length, idBytes); + return compositeKey; } - private static void writeSegment(ByteArrayOutputStream out, byte[] segment) { - int length = segment.length; - out.write((length >>> 24) & 0xff); - out.write((length >>> 16) & 0xff); - out.write((length >>> 8) & 0xff); - out.write(length & 0xff); - out.write(segment, 0, segment.length); + /** The {@code key | namespace} prefix every cell in {@code namespace} starts with. */ + private byte[] prefixFor(StateNamespace namespace) { + byte[] cached = cachedPrefix; + if (cached != null && namespace.equals(cachedNamespace)) { + return cached; + } + byte[] namespaceBytes = namespace.stringKey().getBytes(StandardCharsets.UTF_8); + byte[] prefix = + new byte[StoreKeys.segmentLength(encodedKey) + StoreKeys.segmentLength(namespaceBytes)]; + int offset = StoreKeys.writeSegment(prefix, 0, encodedKey); + StoreKeys.writeSegment(prefix, offset, namespaceBytes); + cachedNamespace = namespace; + cachedPrefix = prefix; + return prefix; } private class KafkaStreamsStateBinder implements StateBinder { @@ -323,7 +367,33 @@ public Instant read() { @Override public void add(Instant outputTime) { Instant current = readValue(); - writeValue(current == null ? outputTime : timestampCombiner.combine(current, outputTime)); + Instant combined = + current == null ? outputTime : timestampCombiner.combine(current, outputTime); + writeValue(combined); + // Mirror the hold into the index so the processor can find the minimum hold across every key + // and window with one lookup instead of reading all of them. + if (current != null) { + holdsIndexStore.delete(holdIndexKey(current)); + } + holdsIndexStore.put(holdIndexKey(combined), EMPTY_VALUE); + } + + @Override + public void clear() { + Instant current = readValue(); + if (current != null) { + holdsIndexStore.delete(holdIndexKey(current)); + } + super.clear(); + } + + /** {@code holdTimestamp | cell}, so the index is ordered by hold time. */ + private byte[] holdIndexKey(Instant hold) { + byte[] cellKey = compositeKey(namespace, id); + byte[] indexKey = new byte[StoreKeys.TIMESTAMP_BYTES + cellKey.length]; + int offset = StoreKeys.writeTimestamp(indexKey, 0, hold.getMillis()); + System.arraycopy(cellKey, 0, indexKey, offset, cellKey.length); + return indexKey; } @Override diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternals.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternals.java index 8a5bf910298c..dddc28eb4381 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternals.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternals.java @@ -17,7 +17,6 @@ */ package org.apache.beam.runners.kafka.streams.translation; -import java.io.ByteArrayOutputStream; import java.nio.charset.StandardCharsets; import org.apache.beam.runners.core.StateNamespace; import org.apache.beam.runners.core.TimerInternals; @@ -31,42 +30,54 @@ import org.joda.time.Instant; /** - * A {@link TimerInternals} for one key, backed by a shared Kafka Streams timer {@link - * KeyValueStore} for a GroupByKey. + * A {@link TimerInternals} for one key, backed by two Kafka Streams stores shared by a GroupByKey. * - *

    Kafka Streams has no per-key timer service, so timers are persisted like any other state. A - * timer is stored under a key built from its identity — {@code key | domain | timerFamily | timerId - * | namespace} — so setting a timer with the same identity overwrites the previous one and deleting - * it removes exactly that entry, matching {@link TimerInternals}' contract. The value is the {@link - * TimerData} encoded with {@link TimerInternals.TimerDataCoderV2}, which carries the fire time. + *

    Kafka Streams has no per-key timer service, so timers are persisted like any other state, in + * two stores that serve the two ways a timer is looked up: * - *

    Firing is driven by {@link WindowedGroupByKeyProcessor}: on a watermark advance it scans the - * store for event-time timers whose fire time has passed and replays them through {@link - * org.apache.beam.runners.core.ReduceFnRunner#onTimers}. Scanning the whole store per advance is - * O(timers); acceptable for this first windowing support and replaceable with a fire-time-ordered - * index later, the same way the initial GroupByKey buffered eagerly. + *

      + *
    • the identity store, keyed by {@code key | domain | timerFamily | timerId | + * namespace}, is how {@link #setTimer} overwrites and {@link #deleteTimer} removes exactly + * one timer, as {@link TimerInternals}' contract requires. Its value is the index key below, + * so a timer that is overwritten or deleted can have its index entry removed without knowing + * what time it had been set for. + *
    • the index store, keyed by {@code domain | fireTimestamp | identity}, is how due + * timers are found. Because the timestamp is written in the sortable form described on {@link + * StoreKeys}, all event-time timers due at a watermark are one range scan — {@link + * #dueEventTimeRangeStart} to {@link #dueEventTimeRangeEnd} — rather than a scan of every + * timer of every key. Its value is the {@link TimerData}, so firing needs no second lookup. + *
    * - *

    This instance reports the current input-watermark and processing times it was constructed with - * (the caller advances them as reports arrive); it never fires timers itself. + *

    Firing is driven by {@link WindowedGroupByKeyProcessor}: on a watermark advance it range-scans + * the index for event-time timers that are due and replays them through {@link + * org.apache.beam.runners.core.ReduceFnRunner#onTimers}. + * + *

    This instance reports the times it was constructed with; it never fires timers itself. */ class KafkaStreamsTimerInternals implements TimerInternals { private final byte[] encodedKey; - private final KeyValueStore timerStore; + private final KeyValueStore identityStore; + private final KeyValueStore indexStore; private final TimerInternals.TimerDataCoderV2 timerCoder; private final Instant inputWatermarkTime; + private final Instant outputWatermarkTime; private final Instant processingTime; KafkaStreamsTimerInternals( byte[] encodedKey, - KeyValueStore timerStore, + KeyValueStore identityStore, + KeyValueStore indexStore, Coder windowCoder, Instant inputWatermarkTime, + Instant outputWatermarkTime, Instant processingTime) { this.encodedKey = encodedKey; - this.timerStore = timerStore; + this.identityStore = identityStore; + this.indexStore = indexStore; this.timerCoder = TimerInternals.TimerDataCoderV2.of(windowCoder); this.inputWatermarkTime = inputWatermarkTime; + this.outputWatermarkTime = outputWatermarkTime; this.processingTime = processingTime; } @@ -83,13 +94,23 @@ public void setTimer( @Override public void setTimer(TimerData timerData) { - timerStore.put(timerStoreKey(encodedKey, timerData), encodeTimer(timerData)); + byte[] identityKey = identityKey(encodedKey, timerData); + // Setting a timer that already exists replaces it, so drop the old index entry first — + // otherwise the timer would still be due at the time it was originally set for. + byte[] previousIndexKey = identityStore.get(identityKey); + if (previousIndexKey != null) { + indexStore.delete(previousIndexKey); + } + byte[] indexKey = + indexKey(timerData.getDomain(), timerData.getTimestamp().getMillis(), identityKey); + identityStore.put(identityKey, indexKey); + indexStore.put(indexKey, encodeTimer(timerData)); } @Override public void deleteTimer( StateNamespace namespace, String timerId, String timerFamilyId, TimeDomain timeDomain) { - timerStore.delete(timerStoreKey(encodedKey, timerId, timerFamilyId, timeDomain, namespace)); + deleteByIdentity(identityKey(encodedKey, timerId, timerFamilyId, timeDomain, namespace)); } @Override @@ -101,7 +122,15 @@ public void deleteTimer(StateNamespace namespace, String timerId, String timerFa @Override public void deleteTimer(TimerData timerKey) { - timerStore.delete(timerStoreKey(encodedKey, timerKey)); + deleteByIdentity(identityKey(encodedKey, timerKey)); + } + + private void deleteByIdentity(byte[] identityKey) { + byte[] indexKey = identityStore.get(identityKey); + if (indexKey != null) { + indexStore.delete(indexKey); + } + identityStore.delete(identityKey); } @Override @@ -109,6 +138,13 @@ public Instant currentProcessingTime() { return processingTime; } + /** + * Returns {@code null}: a synchronized processing time is the slowest processing time across the + * job's workers, which needs the cross-instance coordination that the runner's watermark reports + * only carry for event time. {@link TimerInternals} allows null here, and nothing on the paths + * this runner supports today reads it — it is consulted for processing-time triggers, which land + * with the processing-time timer support in a follow-up (the same work that would supply it). + */ @Override public @Nullable Instant currentSynchronizedProcessingTime() { return null; @@ -119,9 +155,13 @@ public Instant currentInputWatermarkTime() { return inputWatermarkTime; } + /** + * The watermark this GroupByKey has last forwarded downstream, which trails {@link + * #currentInputWatermarkTime} by the pending watermark holds. + */ @Override - public @Nullable Instant currentOutputWatermarkTime() { - return null; + public Instant currentOutputWatermarkTime() { + return outputWatermarkTime; } private byte[] encodeTimer(TimerData timerData) { @@ -132,9 +172,20 @@ private byte[] encodeTimer(TimerData timerData) { } } - /** The store key for a timer's identity. Package-visible so the processor can build/scan it. */ - static byte[] timerStoreKey(byte[] encodedKey, TimerData timerData) { - return timerStoreKey( + /** Decodes an index store value back into its timer. */ + static TimerData decodeTimer(Coder windowCoder, byte[] bytes) { + try { + return CoderUtils.decodeFromByteArray(TimerInternals.TimerDataCoderV2.of(windowCoder), bytes); + } catch (CoderException e) { + throw new RuntimeException("Failed to decode timer", e); + } + } + + /** + * The identity store key for a timer: {@code key | domain | timerFamily | timerId | namespace}. + */ + static byte[] identityKey(byte[] encodedKey, TimerData timerData) { + return identityKey( encodedKey, timerData.getTimerId(), timerData.getTimerFamilyId(), @@ -142,39 +193,73 @@ static byte[] timerStoreKey(byte[] encodedKey, TimerData timerData) { timerData.getNamespace()); } - static byte[] timerStoreKey( + static byte[] identityKey( byte[] encodedKey, String timerId, String timerFamilyId, TimeDomain domain, StateNamespace namespace) { - ByteArrayOutputStream out = new ByteArrayOutputStream(); - writeSegment(out, encodedKey); - writeSegment(out, new byte[] {(byte) domain.ordinal()}); - writeSegment(out, timerFamilyId.getBytes(StandardCharsets.UTF_8)); - writeSegment(out, timerId.getBytes(StandardCharsets.UTF_8)); - writeSegment(out, namespace.stringKey().getBytes(StandardCharsets.UTF_8)); - return out.toByteArray(); - } - - /** Reads the encoded Beam key (the first segment) back out of a timer store key. */ - static byte[] encodedKeyOf(byte[] storeKey) { - int length = - ((storeKey[0] & 0xff) << 24) - | ((storeKey[1] & 0xff) << 16) - | ((storeKey[2] & 0xff) << 8) - | (storeKey[3] & 0xff); - byte[] key = new byte[length]; - System.arraycopy(storeKey, 4, key, 0, length); + byte[] domainBytes = {(byte) domain.ordinal()}; + byte[] familyBytes = timerFamilyId.getBytes(StandardCharsets.UTF_8); + byte[] idBytes = timerId.getBytes(StandardCharsets.UTF_8); + byte[] namespaceBytes = namespace.stringKey().getBytes(StandardCharsets.UTF_8); + byte[] key = + new byte + [StoreKeys.segmentLength(encodedKey) + + StoreKeys.segmentLength(domainBytes) + + StoreKeys.segmentLength(familyBytes) + + StoreKeys.segmentLength(idBytes) + + StoreKeys.segmentLength(namespaceBytes)]; + int offset = StoreKeys.writeSegment(key, 0, encodedKey); + offset = StoreKeys.writeSegment(key, offset, domainBytes); + offset = StoreKeys.writeSegment(key, offset, familyBytes); + offset = StoreKeys.writeSegment(key, offset, idBytes); + StoreKeys.writeSegment(key, offset, namespaceBytes); return key; } - private static void writeSegment(ByteArrayOutputStream out, byte[] segment) { - int length = segment.length; - out.write((length >>> 24) & 0xff); - out.write((length >>> 16) & 0xff); - out.write((length >>> 8) & 0xff); - out.write(length & 0xff); - out.write(segment, 0, segment.length); + /** The index store key for a timer: {@code domain | fireTimestamp | identity}. */ + static byte[] indexKey(TimeDomain domain, long fireMillis, byte[] identityKey) { + byte[] key = new byte[1 + StoreKeys.TIMESTAMP_BYTES + identityKey.length]; + key[0] = (byte) domain.ordinal(); + int offset = StoreKeys.writeTimestamp(key, 1, fireMillis); + System.arraycopy(identityKey, 0, key, offset, identityKey.length); + return key; + } + + /** Inclusive lower bound of the range scan for due event-time timers. */ + static byte[] dueEventTimeRangeStart() { + byte[] bound = new byte[1 + StoreKeys.TIMESTAMP_BYTES]; + bound[0] = (byte) TimeDomain.EVENT_TIME.ordinal(); + StoreKeys.writeTimestamp(bound, 1, Long.MIN_VALUE); + return bound; + } + + /** + * Inclusive upper bound of the range scan for event-time timers due at {@code watermarkMillis}. + * + *

    Every index key carries a non-empty identity after its timestamp, so no key is equal to the + * bare {@code domain | watermark + 1} prefix returned here: an inclusive scan up to it yields + * exactly the timers whose fire time is at or before the watermark. Beam's maximum timestamp is + * far below {@link Long#MAX_VALUE}, so the increment cannot overflow. + */ + static byte[] dueEventTimeRangeEnd(long watermarkMillis) { + byte[] bound = new byte[1 + StoreKeys.TIMESTAMP_BYTES]; + bound[0] = (byte) TimeDomain.EVENT_TIME.ordinal(); + StoreKeys.writeTimestamp(bound, 1, watermarkMillis + 1); + return bound; + } + + /** Reads the identity key back out of an index key. */ + static byte[] identityKeyOf(byte[] indexKey) { + int offset = 1 + StoreKeys.TIMESTAMP_BYTES; + byte[] identityKey = new byte[indexKey.length - offset]; + System.arraycopy(indexKey, offset, identityKey, 0, identityKey.length); + return identityKey; + } + + /** Reads the encoded Beam key (the first segment) back out of an identity key. */ + static byte[] encodedKeyOf(byte[] identityKey) { + return StoreKeys.readSegment(identityKey, 0); } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StoreKeys.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StoreKeys.java new file mode 100644 index 000000000000..5cb00ca8b6cc --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StoreKeys.java @@ -0,0 +1,105 @@ +/* + * 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.beam.runners.kafka.streams.translation; + +/** + * Byte helpers for the composite keys the runner's state and timer stores are addressed by. + * + *

    Keys are built into a single pre-sized array rather than through a stream, because they are on + * the hot path: one is built for every state cell read or written, several times per element. + * + *

    Variable-length parts are written length-prefixed. A separator byte would be shorter, but both + * an encoded Beam key (arbitrary user coder output) and a {@link + * org.apache.beam.runners.core.StateNamespace#stringKey} can contain any byte value, so no + * separator is safe from collisions — {@code key="a/b", ns="c"} and {@code key="a", ns="b/c"} would + * produce the same bytes. Length prefixes also let the encoded key be read back out of a timer key, + * which the timer scan needs. + * + *

    Timestamps are written sign-flipped big-endian so that the unsigned lexicographic order Kafka + * Streams compares keys by is the same as numeric order. That is what makes a range scan over a + * timestamp-prefixed store return exactly the entries up to a point in time, which is how due + * timers and the minimum watermark hold are found without scanning everything. + */ +final class StoreKeys { + + /** Bytes taken by a length prefix. */ + static final int LENGTH_BYTES = 4; + + /** Bytes taken by a sortable timestamp. */ + static final int TIMESTAMP_BYTES = 8; + + private StoreKeys() {} + + /** Bytes a length-prefixed segment occupies. */ + static int segmentLength(byte[] segment) { + return LENGTH_BYTES + segment.length; + } + + /** Writes {@code segment} length-prefixed at {@code offset}, returning the offset after it. */ + static int writeSegment(byte[] target, int offset, byte[] segment) { + int next = writeLength(target, offset, segment.length); + System.arraycopy(segment, 0, target, next, segment.length); + return next + segment.length; + } + + private static int writeLength(byte[] target, int offset, int length) { + target[offset] = (byte) ((length >>> 24) & 0xff); + target[offset + 1] = (byte) ((length >>> 16) & 0xff); + target[offset + 2] = (byte) ((length >>> 8) & 0xff); + target[offset + 3] = (byte) (length & 0xff); + return offset + LENGTH_BYTES; + } + + /** Reads the length prefix at {@code offset}. */ + static int readLength(byte[] source, int offset) { + return ((source[offset] & 0xff) << 24) + | ((source[offset + 1] & 0xff) << 16) + | ((source[offset + 2] & 0xff) << 8) + | (source[offset + 3] & 0xff); + } + + /** Reads the length-prefixed segment starting at {@code offset}. */ + static byte[] readSegment(byte[] source, int offset) { + int length = readLength(source, offset); + byte[] segment = new byte[length]; + System.arraycopy(source, offset + LENGTH_BYTES, segment, 0, length); + return segment; + } + + /** + * Writes a timestamp so that unsigned byte order matches numeric order: flipping the sign bit + * maps {@link Long#MIN_VALUE}..{@link Long#MAX_VALUE} onto 0x00.. 0xff.. big-endian, so + * negative timestamps (valid in Beam) sort before positive ones. + */ + static int writeTimestamp(byte[] target, int offset, long millis) { + long sortable = millis ^ Long.MIN_VALUE; + for (int i = 0; i < TIMESTAMP_BYTES; i++) { + target[offset + i] = (byte) ((sortable >>> (8 * (TIMESTAMP_BYTES - 1 - i))) & 0xff); + } + return offset + TIMESTAMP_BYTES; + } + + /** Reads a timestamp written by {@link #writeTimestamp}. */ + static long readTimestamp(byte[] source, int offset) { + long sortable = 0; + for (int i = 0; i < TIMESTAMP_BYTES; i++) { + sortable = (sortable << 8) | (source[offset + i] & 0xffL); + } + return sortable ^ Long.MIN_VALUE; + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java index b9ff98da0724..1b0ffae23fde 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java @@ -34,7 +34,6 @@ import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.CoderException; import org.apache.beam.sdk.options.PipelineOptions; -import org.apache.beam.sdk.state.TimeDomain; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.util.CoderUtils; import org.apache.beam.sdk.util.construction.TriggerTranslation; @@ -77,7 +76,9 @@ class WindowedGroupByKeyProcessor private static final Logger LOG = LoggerFactory.getLogger(WindowedGroupByKeyProcessor.class); private final String stateStoreName; + private final String holdsIndexStoreName; private final String timerStoreName; + private final String timerIndexStoreName; private final String transformId; private final Coder keyCoder; private final WindowingStrategy windowingStrategy; @@ -92,11 +93,15 @@ class WindowedGroupByKeyProcessor private @Nullable ProcessorContext> context; private @Nullable KeyValueStore stateStore; + private @Nullable KeyValueStore holdsIndexStore; private @Nullable KeyValueStore timerStore; + private @Nullable KeyValueStore timerIndexStore; WindowedGroupByKeyProcessor( String stateStoreName, + String holdsIndexStoreName, String timerStoreName, + String timerIndexStoreName, String transformId, Set upstreamTransformIds, Coder keyCoder, @@ -104,7 +109,9 @@ class WindowedGroupByKeyProcessor WindowingStrategy windowingStrategy, PipelineOptions options) { this.stateStoreName = stateStoreName; + this.holdsIndexStoreName = holdsIndexStoreName; this.timerStoreName = timerStoreName; + this.timerIndexStoreName = timerIndexStoreName; this.transformId = transformId; this.keyCoder = keyCoder; this.windowingStrategy = windowingStrategy; @@ -119,7 +126,9 @@ class WindowedGroupByKeyProcessor public void init(ProcessorContext> context) { this.context = context; this.stateStore = context.getStateStore(stateStoreName); + this.holdsIndexStore = context.getStateStore(holdsIndexStoreName); this.timerStore = context.getStateStore(timerStoreName); + this.timerIndexStore = context.getStateStore(timerIndexStoreName); } @Override @@ -137,12 +146,32 @@ public void process(Record> record) { } watermarkAggregator.observe(payload.asWatermark()); Instant advanced = watermarkAggregator.advance(); - if (advanced.isAfter(lastForwardedWatermark)) { + if (advanced.isAfter(inputWatermark)) { inputWatermark = advanced; fireDueEventTimeTimers(record, advanced); - lastForwardedWatermark = advanced; - forwardWatermark(record, advanced.getMillis()); } + // Firing may have emitted panes and released their holds, so the output watermark is computed + // after it. + Instant output = outputWatermark(); + if (output.isAfter(lastForwardedWatermark)) { + lastForwardedWatermark = output; + forwardWatermark(record, output.getMillis()); + } + } + + /** + * The watermark to publish downstream: the input watermark, held back by the earliest watermark + * hold any pending pane has taken. + * + *

    {@link ReduceFnRunner} takes a hold for buffered elements that have not been emitted yet, at + * the timestamp their pane will carry. Forwarding the raw input watermark would tell downstream + * that nothing earlier is coming while those panes are still buffered, and the elements would + * then arrive late against the watermark we had already published. + */ + private Instant outputWatermark() { + Instant minHold = + KafkaStreamsStateInternals.minWatermarkHold(checkInitialized(holdsIndexStore)); + return minHold == null || inputWatermark.isBefore(minHold) ? inputWatermark : minHold; } private void processData(Record> record, KStreamsPayload payload) { @@ -158,33 +187,45 @@ private void processData(Record> record, KStreamsPayl record, encodedKey, key, Collections.singletonList(valueElement), Collections.emptyList()); } - /** Fires every event-time timer whose fire time is at or before the new input watermark. */ + /** + * Fires every event-time timer whose fire time is at or before the new input watermark. + * + *

    The timers to fire are found by range-scanning the fire-time-ordered index over exactly the + * window {@code (-inf, watermark]}, so the cost is proportional to the number of timers that are + * actually due rather than to the number of keys that hold a timer. + */ private void fireDueEventTimeTimers( Record> record, Instant watermark) { - KeyValueStore timers = checkInitialized(timerStore); - // Group due timers by the Beam key they belong to; a key's timers fire together in one - // ReduceFnRunner turn. Whole-store scan per advance is O(timers) — see the class doc. + KeyValueStore identityStore = checkInitialized(timerStore); + KeyValueStore indexStore = checkInitialized(timerIndexStore); + // Group the due timers by the Beam key they belong to; a key's timers fire together in one + // ReduceFnRunner turn. Map dueByKey = new LinkedHashMap<>(); - List firedStoreKeys = new ArrayList<>(); - try (KeyValueIterator it = timers.all()) { + List firedIndexKeys = new ArrayList<>(); + try (KeyValueIterator it = + indexStore.range( + KafkaStreamsTimerInternals.dueEventTimeRangeStart(), + KafkaStreamsTimerInternals.dueEventTimeRangeEnd(watermark.getMillis()))) { while (it.hasNext()) { org.apache.kafka.streams.KeyValue entry = it.next(); - TimerData timer = decodeTimer(entry.value); - if (timer.getDomain() != TimeDomain.EVENT_TIME || timer.getTimestamp().isAfter(watermark)) { - continue; - } - byte[] encodedKey = KafkaStreamsTimerInternals.encodedKeyOf(entry.key); + TimerData timer = KafkaStreamsTimerInternals.decodeTimer(windowCoder, entry.value); + byte[] encodedKey = + KafkaStreamsTimerInternals.encodedKeyOf( + KafkaStreamsTimerInternals.identityKeyOf(entry.key)); dueByKey .computeIfAbsent( BaseEncoding.base16().encode(encodedKey), k -> new DueTimers(encodedKey)) .timers .add(timer); - firedStoreKeys.add(entry.key); + firedIndexKeys.add(entry.key); } } - // Remove fired timers before replaying them; onTimers may legitimately set new ones. - for (byte[] storeKey : firedStoreKeys) { - timers.delete(storeKey); + // Clear the fired timers from both stores before replaying them, since onTimers may + // legitimately + // set new ones — including at the same identity. + for (byte[] indexKey : firedIndexKeys) { + indexStore.delete(indexKey); + identityStore.delete(KafkaStreamsTimerInternals.identityKeyOf(indexKey)); } for (DueTimers due : dueByKey.values()) { runReduceFn( @@ -199,10 +240,17 @@ private void runReduceFn( List> elements, List timers) { StateInternals stateInternals = - new KafkaStreamsStateInternals<>(key, encodedKey, checkInitialized(stateStore)); + new KafkaStreamsStateInternals<>( + key, encodedKey, checkInitialized(stateStore), checkInitialized(holdsIndexStore)); KafkaStreamsTimerInternals timerInternals = new KafkaStreamsTimerInternals( - encodedKey, checkInitialized(timerStore), windowCoder, inputWatermark, Instant.now()); + encodedKey, + checkInitialized(timerStore), + checkInitialized(timerIndexStore), + windowCoder, + inputWatermark, + lastForwardedWatermark, + Instant.now()); ReduceFnRunner, W> runner = new ReduceFnRunner<>( key, @@ -256,15 +304,6 @@ private void forwardWatermark(Record> trigger, long w } } - private TimerData decodeTimer(byte[] bytes) { - try { - return CoderUtils.decodeFromByteArray( - org.apache.beam.runners.core.TimerInternals.TimerDataCoderV2.of(windowCoder), bytes); - } catch (CoderException e) { - throw new RuntimeException("Failed to decode timer", e); - } - } - private static T checkInitialized(@Nullable T value) { if (value == null) { throw new IllegalStateException("WindowedGroupByKeyProcessor used before init()"); diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternalsTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternalsTest.java new file mode 100644 index 000000000000..df5b0eda1b26 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternalsTest.java @@ -0,0 +1,208 @@ +/* + * 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.beam.runners.kafka.streams.translation; + +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.MatcherAssert.assertThat; + +import java.util.ArrayList; +import java.util.List; +import org.apache.beam.runners.core.StateNamespace; +import org.apache.beam.runners.core.StateNamespaces; +import org.apache.beam.runners.core.TimerInternals.TimerData; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.state.TimeDomain; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.kafka.common.serialization.Serdes; +import org.apache.kafka.streams.processor.api.MockProcessorContext; +import org.apache.kafka.streams.state.KeyValueIterator; +import org.apache.kafka.streams.state.KeyValueStore; +import org.apache.kafka.streams.state.Stores; +import org.joda.time.Instant; +import org.junit.Before; +import org.junit.Test; + +/** + * Tests the timer store and its fire-time index: that a timer can be set, replaced and deleted by + * identity, and that the timers due at a watermark are found by a range scan over the index rather + * than by inspecting every timer. + */ +public class KafkaStreamsTimerInternalsTest { + + private static final StateNamespace NAMESPACE = + StateNamespaces.window(GlobalWindow.Coder.INSTANCE, GlobalWindow.INSTANCE); + + private KeyValueStore identityStore; + private KeyValueStore indexStore; + + @Before + public void setUp() { + MockProcessorContext context = new MockProcessorContext<>(); + identityStore = newStore("timers", context); + indexStore = newStore("timers-index", context); + } + + private static KeyValueStore newStore( + String name, MockProcessorContext context) { + KeyValueStore store = + Stores.keyValueStoreBuilder( + Stores.inMemoryKeyValueStore(name), Serdes.ByteArray(), Serdes.ByteArray()) + .withLoggingDisabled() + .build(); + store.init(context.getStateStoreContext(), store); + return store; + } + + private KafkaStreamsTimerInternals timersFor(String key) { + return new KafkaStreamsTimerInternals( + encode(key), + identityStore, + indexStore, + GlobalWindow.Coder.INSTANCE, + BoundedWindow.TIMESTAMP_MIN_VALUE, + BoundedWindow.TIMESTAMP_MIN_VALUE, + new Instant(0)); + } + + private static byte[] encode(String key) { + try { + return CoderUtils.encodeToByteArray(StringUtf8Coder.of(), key); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + private static TimerData eventTimer(String id, long millis) { + return TimerData.of( + id, "", NAMESPACE, new Instant(millis), new Instant(millis), TimeDomain.EVENT_TIME); + } + + /** The timers the processor would fire at {@code watermarkMillis}, in fire-time order. */ + private List dueAt(long watermarkMillis) { + List due = new ArrayList<>(); + try (KeyValueIterator it = + indexStore.range( + KafkaStreamsTimerInternals.dueEventTimeRangeStart(), + KafkaStreamsTimerInternals.dueEventTimeRangeEnd(watermarkMillis))) { + while (it.hasNext()) { + due.add( + KafkaStreamsTimerInternals.decodeTimer(GlobalWindow.Coder.INSTANCE, it.next().value)); + } + } + return due; + } + + private static int storeSize(KeyValueStore store) { + int size = 0; + try (KeyValueIterator it = store.all()) { + while (it.hasNext()) { + it.next(); + size++; + } + } + return size; + } + + @Test + public void dueScanReturnsOnlyTimersAtOrBeforeTheWatermark() { + KafkaStreamsTimerInternals timers = timersFor("key"); + timers.setTimer(eventTimer("early", 100L)); + timers.setTimer(eventTimer("onWatermark", 200L)); + timers.setTimer(eventTimer("late", 300L)); + + List due = dueAt(200L); + + // Ordered by fire time, and the timer set exactly at the watermark is included. + assertThat(due.size(), is(2)); + assertThat(due.get(0).getTimerId(), is("early")); + assertThat(due.get(1).getTimerId(), is("onWatermark")); + } + + @Test + public void negativeTimestampsSortBeforePositiveOnes() { + KafkaStreamsTimerInternals timers = timersFor("key"); + timers.setTimer(eventTimer("negative", -5000L)); + timers.setTimer(eventTimer("zero", 0L)); + timers.setTimer(eventTimer("positive", 5000L)); + + List due = dueAt(0L); + + assertThat(due.size(), is(2)); + assertThat(due.get(0).getTimerId(), is("negative")); + assertThat(due.get(1).getTimerId(), is("zero")); + } + + @Test + public void resettingATimerReplacesItsIndexEntry() { + KafkaStreamsTimerInternals timers = timersFor("key"); + timers.setTimer(eventTimer("timer", 100L)); + // Re-setting the same timer identity for a later time must not leave the old entry behind, + // or the timer would still fire at the time it was first set for. + timers.setTimer(eventTimer("timer", 900L)); + + assertThat(dueAt(100L).isEmpty(), is(true)); + assertThat(dueAt(900L).size(), is(1)); + assertThat(storeSize(indexStore), is(1)); + assertThat(storeSize(identityStore), is(1)); + } + + @Test + public void deletingATimerRemovesItFromBothStores() { + KafkaStreamsTimerInternals timers = timersFor("key"); + timers.setTimer(eventTimer("timer", 100L)); + timers.deleteTimer(NAMESPACE, "timer", "", TimeDomain.EVENT_TIME); + + assertThat(dueAt(1000L).isEmpty(), is(true)); + assertThat(storeSize(indexStore), is(0)); + assertThat(storeSize(identityStore), is(0)); + } + + @Test + public void timersOfDifferentKeysAreIndependentButShareTheIndex() { + timersFor("a").setTimer(eventTimer("timer", 100L)); + timersFor("b").setTimer(eventTimer("timer", 150L)); + + // Same timer id under two Beam keys are two distinct timers, and one scan finds both. + assertThat(storeSize(identityStore), is(2)); + assertThat(dueAt(200L).size(), is(2)); + + timersFor("a").deleteTimer(NAMESPACE, "timer", "", TimeDomain.EVENT_TIME); + assertThat(dueAt(200L).size(), is(1)); + } + + @Test + public void processingTimeTimersAreNotReturnedByTheEventTimeScan() { + KafkaStreamsTimerInternals timers = timersFor("key"); + timers.setTimer( + TimerData.of( + "processing", + "", + NAMESPACE, + new Instant(100L), + new Instant(100L), + TimeDomain.PROCESSING_TIME)); + timers.setTimer(eventTimer("event", 100L)); + + List due = dueAt(1000L); + + assertThat(due.size(), is(1)); + assertThat(due.get(0).getTimerId(), is("event")); + } +}