bbejeck commented on code in PR #22652: URL: https://github.com/apache/kafka/pull/22652#discussion_r3467932276
########## streams/src/main/java/org/apache/kafka/streams/state/internals/InMemorySessionTransactionBuffer.java: ########## @@ -0,0 +1,392 @@ +/* + * 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.kafka.streams.state.internals; + +import org.apache.kafka.common.utils.Bytes; +import org.apache.kafka.streams.KeyValue; +import org.apache.kafka.streams.kstream.Windowed; + +import java.util.Iterator; +import java.util.Map; +import java.util.NoSuchElementException; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.ConcurrentNavigableMap; +import java.util.concurrent.ConcurrentSkipListMap; + +/** + * A {@link TransactionBuffer} implementation for {@link InMemorySessionStore}. + * Uses a composite key of (endTime, key, startTime) to maintain correct sort order in the staging map. + */ +class InMemorySessionTransactionBuffer extends AbstractTransactionBuffer<InMemorySessionTransactionBuffer.SessionEntryKey> { + + private final ConcurrentNavigableMap<Long, ConcurrentNavigableMap<Bytes, ConcurrentNavigableMap<Long, byte[]>>> endTimeMap; + + InMemorySessionTransactionBuffer( + final ConcurrentNavigableMap<Long, ConcurrentNavigableMap<Bytes, ConcurrentNavigableMap<Long, byte[]>>> endTimeMap) { + this.endTimeMap = endTimeMap; + } + + /** + * Composite key for the session store staging map. Sorts by endTime, then key, then startTime. + */ + static final class SessionEntryKey implements Comparable<SessionEntryKey> { + private final long endTime; + private final Bytes key; + private final long startTime; + + SessionEntryKey(final long endTime, final Bytes key, final long startTime) { + this.endTime = endTime; + this.key = key; + this.startTime = startTime; + } + + long endTime() { + return endTime; + } + + Bytes key() { + return key; + } + + long startTime() { + return startTime; + } + + @Override + public int compareTo(final SessionEntryKey other) { + int cmp = Long.compare(this.endTime, other.endTime); + if (cmp != 0) { + return cmp; + } + cmp = this.key.compareTo(other.key); + if (cmp != 0) { + return cmp; + } + return Long.compare(this.startTime, other.startTime); + } + + @Override + public boolean equals(final Object o) { + if (this == o) return true; + if (!(o instanceof SessionEntryKey)) return false; + final SessionEntryKey that = (SessionEntryKey) o; + return endTime == that.endTime && startTime == that.startTime && Objects.equals(key, that.key); + } + + @Override + public int hashCode() { + return Objects.hash(endTime, key, startTime); + } + } + + // -- Convenience methods for the store -- + + void stage(final Windowed<Bytes> sessionKey, final byte[] value) { + super.stage(new SessionEntryKey(sessionKey.window().end(), sessionKey.key(), sessionKey.window().start()), value); + } + + Optional<byte[]> get(final Bytes key, final long startTime, final long endTime) { + return super.get(new SessionEntryKey(endTime, key, startTime)); + } + + // -- AbstractTransactionBuffer implementation -- + + @Override + int estimateKeySize(final SessionEntryKey key) { + return 2 * Long.BYTES + key.key().get().length; Review Comment: why 2 * X ########## streams/src/main/java/org/apache/kafka/streams/state/internals/InMemorySessionStore.java: ########## Review Comment: This goes to the buffer when restoring unlike the I[nMemoryKeyValueStore which bypasses the buffer and uses putInternal](https://github.com/apache/kafka/blob/trunk/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryKeyValueStore.java#L90). I know during one review we went with putInternal to preserve the `position`. For restoring I guess going through the buffer is OK - I'm not considering all the trade-offs, the bigger issue is we should have a standard approach across all stores if possible. If we have a reason for doing this, I don't mind but it should be documented in the code describing why the differences. ########## streams/src/main/java/org/apache/kafka/streams/state/internals/InMemorySessionStore.java: ########## @@ -423,6 +525,86 @@ private InMemorySessionStoreIterator registerNewIterator(final Bytes keyFrom, return iterator; } + /** + * A session iterator backed by a transactional buffer's merge scan. + * Converts SessionEntryKey/byte[] pairs into Windowed<Bytes>/byte[] pairs, + * filtering by latestSessionStartTime. + */ + private static class TransactionalSessionIterator implements KeyValueIterator<Windowed<Bytes>, byte[]> { + private final KeyValueIterator<InMemorySessionTransactionBuffer.SessionEntryKey, byte[]> delegate; + private final long latestSessionStartTime; + private KeyValue<Windowed<Bytes>, byte[]> prefetched; + + TransactionalSessionIterator( + final InMemorySessionTransactionBuffer buffer, + final Bytes keyFrom, + final Bytes keyTo, + final long latestSessionStartTime, + final long earliestSessionEndTime, + final long latestSessionEndTime, + final boolean forward) { + this.latestSessionStartTime = latestSessionStartTime; + + final InMemorySessionTransactionBuffer.SessionEntryKey from = + new InMemorySessionTransactionBuffer.SessionEntryKey( + earliestSessionEndTime, + keyFrom != null ? keyFrom : Bytes.wrap(new byte[0]), + 0); + final InMemorySessionTransactionBuffer.SessionEntryKey to = + new InMemorySessionTransactionBuffer.SessionEntryKey( + latestSessionEndTime, + keyTo != null ? keyTo : Bytes.wrap(new byte[]{(byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF}), Review Comment: why this if `keyTo` is null? Plus it's kinda jarring, worth it to a static final variable instead ########## streams/src/main/java/org/apache/kafka/streams/state/internals/InMemorySessionTransactionBuffer.java: ########## @@ -0,0 +1,392 @@ +/* + * 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.kafka.streams.state.internals; + +import org.apache.kafka.common.utils.Bytes; +import org.apache.kafka.streams.KeyValue; +import org.apache.kafka.streams.kstream.Windowed; + +import java.util.Iterator; +import java.util.Map; +import java.util.NoSuchElementException; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.ConcurrentNavigableMap; +import java.util.concurrent.ConcurrentSkipListMap; + +/** + * A {@link TransactionBuffer} implementation for {@link InMemorySessionStore}. + * Uses a composite key of (endTime, key, startTime) to maintain correct sort order in the staging map. + */ +class InMemorySessionTransactionBuffer extends AbstractTransactionBuffer<InMemorySessionTransactionBuffer.SessionEntryKey> { + + private final ConcurrentNavigableMap<Long, ConcurrentNavigableMap<Bytes, ConcurrentNavigableMap<Long, byte[]>>> endTimeMap; + + InMemorySessionTransactionBuffer( + final ConcurrentNavigableMap<Long, ConcurrentNavigableMap<Bytes, ConcurrentNavigableMap<Long, byte[]>>> endTimeMap) { + this.endTimeMap = endTimeMap; + } + + /** + * Composite key for the session store staging map. Sorts by endTime, then key, then startTime. + */ + static final class SessionEntryKey implements Comparable<SessionEntryKey> { + private final long endTime; + private final Bytes key; + private final long startTime; + + SessionEntryKey(final long endTime, final Bytes key, final long startTime) { + this.endTime = endTime; + this.key = key; + this.startTime = startTime; + } + + long endTime() { + return endTime; + } + + Bytes key() { + return key; + } + + long startTime() { + return startTime; + } + + @Override + public int compareTo(final SessionEntryKey other) { + int cmp = Long.compare(this.endTime, other.endTime); + if (cmp != 0) { + return cmp; + } + cmp = this.key.compareTo(other.key); + if (cmp != 0) { + return cmp; + } + return Long.compare(this.startTime, other.startTime); + } + + @Override + public boolean equals(final Object o) { + if (this == o) return true; + if (!(o instanceof SessionEntryKey)) return false; + final SessionEntryKey that = (SessionEntryKey) o; + return endTime == that.endTime && startTime == that.startTime && Objects.equals(key, that.key); + } + + @Override + public int hashCode() { + return Objects.hash(endTime, key, startTime); + } + } + + // -- Convenience methods for the store -- + + void stage(final Windowed<Bytes> sessionKey, final byte[] value) { + super.stage(new SessionEntryKey(sessionKey.window().end(), sessionKey.key(), sessionKey.window().start()), value); + } + + Optional<byte[]> get(final Bytes key, final long startTime, final long endTime) { + return super.get(new SessionEntryKey(endTime, key, startTime)); + } + + // -- AbstractTransactionBuffer implementation -- + + @Override + int estimateKeySize(final SessionEntryKey key) { + return 2 * Long.BYTES + key.key().get().length; + } + + @Override + void stageToBackend(final SessionEntryKey key, final byte[] value) { + // no-op — staging map is sufficient; no write-batch concept for in-memory + } + + @Override + ManagedKeyValueIterator<SessionEntryKey, byte[]> newBaseIterator(final SessionEntryKey from, final SessionEntryKey to) { + return newBaseIterator(from, to, true, true); + } + + @Override + ManagedKeyValueIterator<SessionEntryKey, byte[]> newBaseIterator(final SessionEntryKey from, final SessionEntryKey to, + final boolean forward, final boolean toInclusive) { + final ConcurrentNavigableMap<Long, ConcurrentNavigableMap<Bytes, ConcurrentNavigableMap<Long, byte[]>>> timeRange; + if (from != null && to != null) { + timeRange = endTimeMap.subMap(from.endTime(), true, to.endTime(), true); + } else if (from != null) { + timeRange = endTimeMap.tailMap(from.endTime(), true); + } else if (to != null) { + timeRange = endTimeMap.headMap(to.endTime(), true); + } else { + timeRange = endTimeMap; + } + + return new FlattenedSessionIterator( + forward ? timeRange : timeRange.descendingMap(), + from, to, forward, toInclusive + ); + } + + /** + * Non-owner (IQ) path: eagerly deep-copies the bounded end-time range while the caller holds + * the snapshot read-lock, providing true point-in-time isolation. The returned iterator never + * touches the live end-time map, so concurrent owner mutation cannot disturb it. + */ + @Override + ManagedKeyValueIterator<SessionEntryKey, byte[]> newBaseSnapshotIterator(final SessionEntryKey from, final SessionEntryKey to, + final boolean forward, final boolean toInclusive) { + final ConcurrentNavigableMap<Long, ConcurrentNavigableMap<Bytes, ConcurrentNavigableMap<Long, byte[]>>> timeRange; + if (from != null && to != null) { + timeRange = endTimeMap.subMap(from.endTime(), true, to.endTime(), true); + } else if (from != null) { + timeRange = endTimeMap.tailMap(from.endTime(), true); + } else if (to != null) { + timeRange = endTimeMap.headMap(to.endTime(), true); + } else { + timeRange = endTimeMap; + } + + final ConcurrentNavigableMap<Long, ConcurrentNavigableMap<Bytes, ConcurrentNavigableMap<Long, byte[]>>> copy = new ConcurrentSkipListMap<>(); + for (final Map.Entry<Long, ConcurrentNavigableMap<Bytes, ConcurrentNavigableMap<Long, byte[]>>> endTimeEntry : timeRange.entrySet()) { + final ConcurrentNavigableMap<Bytes, ConcurrentNavigableMap<Long, byte[]>> keyCopy = new ConcurrentSkipListMap<>(); + for (final Map.Entry<Bytes, ConcurrentNavigableMap<Long, byte[]>> keyEntry : endTimeEntry.getValue().entrySet()) { + keyCopy.put(keyEntry.getKey(), new ConcurrentSkipListMap<>(keyEntry.getValue())); + } + copy.put(endTimeEntry.getKey(), keyCopy); + } + + return new FlattenedSessionIterator( + forward ? copy : copy.descendingMap(), + from, to, forward, toInclusive + ); + } + + @Override + void flushToBase() { + for (final Map.Entry<SessionEntryKey, Optional<byte[]>> entry : pendingWrites.entrySet()) { + final long endTime = entry.getKey().endTime(); + final Bytes key = entry.getKey().key(); + final long startTime = entry.getKey().startTime(); + if (entry.getValue().isPresent()) { + endTimeMap.computeIfAbsent(endTime, t -> new ConcurrentSkipListMap<>()) + .computeIfAbsent(key, k -> new ConcurrentSkipListMap<>()) + .put(startTime, entry.getValue().get()); + } else { + final ConcurrentNavigableMap<Bytes, ConcurrentNavigableMap<Long, byte[]>> keyMap = endTimeMap.get(endTime); + if (keyMap != null) { + final ConcurrentNavigableMap<Long, byte[]> startTimeMap = keyMap.get(key); + if (startTimeMap != null) { + startTimeMap.remove(startTime); + if (startTimeMap.isEmpty()) { + keyMap.remove(key); + if (keyMap.isEmpty()) { + endTimeMap.remove(endTime); + } + } + } + } + } + } + } + + @Override + void discardPendingBatch() { + // no-op — no backend batch to discard + } + + /** + * Iterator that flattens the three-level endTimeMap structure into a stream of + * SessionEntryKey/byte[] pairs, respecting key bounds and direction. + */ + private static class FlattenedSessionIterator implements ManagedKeyValueIterator<SessionEntryKey, byte[]> { Review Comment: Seems we don't have any tests that end up hitting this code. The `FlattenedSessionIterator` is complex so I think it would be good to get some tests hitting it. UPDATE: Mucking around a bit I was able to find a bug (I think) - enable transactional statestores - put 2-3 keys - `a`, `b`, `c` - do a `fetch(a)` will return `b` and `c` This behavior isn't present in the non-transactional mode. Any possibility of reusing the non-transactional iterator for the base layer and have `newBaseIterator/newBaseSnapshotIterator` delegate to the same per-segment-key-filtering logic? Just a suggestion, not trying to dictate implementation. ########## streams/src/main/java/org/apache/kafka/streams/state/internals/InMemorySessionTransactionBuffer.java: ########## @@ -0,0 +1,392 @@ +/* + * 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.kafka.streams.state.internals; + +import org.apache.kafka.common.utils.Bytes; +import org.apache.kafka.streams.KeyValue; +import org.apache.kafka.streams.kstream.Windowed; + +import java.util.Iterator; +import java.util.Map; +import java.util.NoSuchElementException; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.ConcurrentNavigableMap; +import java.util.concurrent.ConcurrentSkipListMap; + +/** + * A {@link TransactionBuffer} implementation for {@link InMemorySessionStore}. + * Uses a composite key of (endTime, key, startTime) to maintain correct sort order in the staging map. + */ +class InMemorySessionTransactionBuffer extends AbstractTransactionBuffer<InMemorySessionTransactionBuffer.SessionEntryKey> { + + private final ConcurrentNavigableMap<Long, ConcurrentNavigableMap<Bytes, ConcurrentNavigableMap<Long, byte[]>>> endTimeMap; + + InMemorySessionTransactionBuffer( + final ConcurrentNavigableMap<Long, ConcurrentNavigableMap<Bytes, ConcurrentNavigableMap<Long, byte[]>>> endTimeMap) { + this.endTimeMap = endTimeMap; + } + + /** + * Composite key for the session store staging map. Sorts by endTime, then key, then startTime. + */ + static final class SessionEntryKey implements Comparable<SessionEntryKey> { + private final long endTime; + private final Bytes key; + private final long startTime; + + SessionEntryKey(final long endTime, final Bytes key, final long startTime) { + this.endTime = endTime; + this.key = key; + this.startTime = startTime; + } + + long endTime() { + return endTime; + } + + Bytes key() { + return key; + } + + long startTime() { + return startTime; + } + + @Override + public int compareTo(final SessionEntryKey other) { + int cmp = Long.compare(this.endTime, other.endTime); + if (cmp != 0) { + return cmp; + } + cmp = this.key.compareTo(other.key); + if (cmp != 0) { + return cmp; + } + return Long.compare(this.startTime, other.startTime); + } + + @Override + public boolean equals(final Object o) { + if (this == o) return true; + if (!(o instanceof SessionEntryKey)) return false; + final SessionEntryKey that = (SessionEntryKey) o; + return endTime == that.endTime && startTime == that.startTime && Objects.equals(key, that.key); + } + + @Override + public int hashCode() { + return Objects.hash(endTime, key, startTime); + } + } + + // -- Convenience methods for the store -- + + void stage(final Windowed<Bytes> sessionKey, final byte[] value) { + super.stage(new SessionEntryKey(sessionKey.window().end(), sessionKey.key(), sessionKey.window().start()), value); + } + + Optional<byte[]> get(final Bytes key, final long startTime, final long endTime) { + return super.get(new SessionEntryKey(endTime, key, startTime)); + } + + // -- AbstractTransactionBuffer implementation -- + + @Override + int estimateKeySize(final SessionEntryKey key) { + return 2 * Long.BYTES + key.key().get().length; + } + + @Override + void stageToBackend(final SessionEntryKey key, final byte[] value) { + // no-op — staging map is sufficient; no write-batch concept for in-memory + } + + @Override + ManagedKeyValueIterator<SessionEntryKey, byte[]> newBaseIterator(final SessionEntryKey from, final SessionEntryKey to) { + return newBaseIterator(from, to, true, true); + } + + @Override + ManagedKeyValueIterator<SessionEntryKey, byte[]> newBaseIterator(final SessionEntryKey from, final SessionEntryKey to, + final boolean forward, final boolean toInclusive) { + final ConcurrentNavigableMap<Long, ConcurrentNavigableMap<Bytes, ConcurrentNavigableMap<Long, byte[]>>> timeRange; + if (from != null && to != null) { + timeRange = endTimeMap.subMap(from.endTime(), true, to.endTime(), true); + } else if (from != null) { + timeRange = endTimeMap.tailMap(from.endTime(), true); + } else if (to != null) { + timeRange = endTimeMap.headMap(to.endTime(), true); + } else { + timeRange = endTimeMap; + } + + return new FlattenedSessionIterator( + forward ? timeRange : timeRange.descendingMap(), + from, to, forward, toInclusive + ); + } + + /** + * Non-owner (IQ) path: eagerly deep-copies the bounded end-time range while the caller holds + * the snapshot read-lock, providing true point-in-time isolation. The returned iterator never + * touches the live end-time map, so concurrent owner mutation cannot disturb it. + */ + @Override + ManagedKeyValueIterator<SessionEntryKey, byte[]> newBaseSnapshotIterator(final SessionEntryKey from, final SessionEntryKey to, + final boolean forward, final boolean toInclusive) { Review Comment: nit: parameters on one line ########## streams/src/main/java/org/apache/kafka/streams/state/internals/InMemorySessionStore.java: ########## @@ -405,6 +494,19 @@ private void removeExpiredSegments() { endTimeMap.headMap(minLiveTime, false).clear(); } + private KeyValueIterator<Windowed<Bytes>, byte[]> newTransactionalSessionIterator( + final Bytes keyFrom, + final Bytes keyTo, + final long latestSessionStartTime, + final long earliestSessionEndTime, + final long latestSessionEndTime, + final boolean forward) { + return new TransactionalSessionIterator( + transactionBuffer, keyFrom, keyTo, latestSessionStartTime, + earliestSessionEndTime, latestSessionEndTime, forward Review Comment: nit: 1 parameter per line -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
