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..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 @@ -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,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 @@ -82,12 +91,18 @@ 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 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); @@ -111,27 +126,70 @@ 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, + holdsIndexStoreName, + timerStoreName, + timerIndexStoreName, 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); + // 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); } + /** 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..c17d7c45ad4d --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsStateInternals.java @@ -0,0 +1,463 @@ +/* + * 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.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.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; + +/** + * 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 { + + /** 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, + KeyValueStore holdsIndexStore) { + this.key = key; + this.encodedKey = encodedKey; + this.store = store; + this.holdsIndexStore = holdsIndexStore; + } + + @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}. + * + *

    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) { + 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; + } + + /** 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 { + 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(); + 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 + 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..dddc28eb4381 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternals.java @@ -0,0 +1,265 @@ +/* + * 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.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 two Kafka Streams stores shared by a GroupByKey. + * + *

    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: + * + *

      + *
    • 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. + *
    + * + *

    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 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 identityStore, + KeyValueStore indexStore, + Coder windowCoder, + Instant inputWatermarkTime, + Instant outputWatermarkTime, + Instant processingTime) { + this.encodedKey = encodedKey; + this.identityStore = identityStore; + this.indexStore = indexStore; + this.timerCoder = TimerInternals.TimerDataCoderV2.of(windowCoder); + this.inputWatermarkTime = inputWatermarkTime; + this.outputWatermarkTime = outputWatermarkTime; + 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) { + 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) { + deleteByIdentity(identityKey(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) { + deleteByIdentity(identityKey(encodedKey, timerKey)); + } + + private void deleteByIdentity(byte[] identityKey) { + byte[] indexKey = identityStore.get(identityKey); + if (indexKey != null) { + indexStore.delete(indexKey); + } + identityStore.delete(identityKey); + } + + @Override + 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; + } + + @Override + public Instant currentInputWatermarkTime() { + return inputWatermarkTime; + } + + /** + * The watermark this GroupByKey has last forwarded downstream, which trails {@link + * #currentInputWatermarkTime} by the pending watermark holds. + */ + @Override + public Instant currentOutputWatermarkTime() { + return outputWatermarkTime; + } + + private byte[] encodeTimer(TimerData timerData) { + try { + return CoderUtils.encodeToByteArray(timerCoder, timerData); + } catch (CoderException e) { + throw new RuntimeException("Failed to encode timer " + timerData, e); + } + } + + /** 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(), + timerData.getDomain(), + timerData.getNamespace()); + } + + static byte[] identityKey( + byte[] encodedKey, + String timerId, + String timerFamilyId, + TimeDomain domain, + StateNamespace namespace) { + 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; + } + + /** 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 new file mode 100644 index 000000000000..1b0ffae23fde --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java @@ -0,0 +1,323 @@ +/* + * 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.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 holdsIndexStoreName; + private final String timerStoreName; + private final String timerIndexStoreName; + 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 holdsIndexStore; + private @Nullable KeyValueStore timerStore; + private @Nullable KeyValueStore timerIndexStore; + + WindowedGroupByKeyProcessor( + String stateStoreName, + String holdsIndexStoreName, + String timerStoreName, + String timerIndexStoreName, + String transformId, + Set upstreamTransformIds, + Coder keyCoder, + Coder valueCoder, + 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; + 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.holdsIndexStore = context.getStateStore(holdsIndexStoreName); + this.timerStore = context.getStateStore(timerStoreName); + this.timerIndexStore = context.getStateStore(timerIndexStoreName); + } + + @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(inputWatermark)) { + inputWatermark = advanced; + fireDueEventTimeTimers(record, advanced); + } + // 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) { + 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. + * + *

    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 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 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 = 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); + firedIndexKeys.add(entry.key); + } + } + // 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( + 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), checkInitialized(holdsIndexStore)); + KafkaStreamsTimerInternals timerInternals = + new KafkaStreamsTimerInternals( + encodedKey, + checkInitialized(timerStore), + checkInitialized(timerIndexStore), + windowCoder, + inputWatermark, + lastForwardedWatermark, + 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 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]")); + } + } +} 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")); + } +}