From 418184da6a630a3b7ebcd6db52e7e1d64c0c4ac7 Mon Sep 17 00:00:00 2001 From: Victoria Xia Date: Tue, 14 Feb 2023 18:22:17 -0800 Subject: [PATCH 1/4] add metered wrapper for versioned stores --- .../state/internals/MeteredKeyValueStore.java | 4 +- .../MeteredVersionedKeyValueStore.java | 227 ++++++++++++++ .../MeteredVersionedKeyValueStoreTest.java | 279 ++++++++++++++++++ 3 files changed, 508 insertions(+), 2 deletions(-) create mode 100644 streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStore.java create mode 100644 streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStoreTest.java diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredKeyValueStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredKeyValueStore.java index 9f8ca1c72a91b..362d08f254fa8 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredKeyValueStore.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredKeyValueStore.java @@ -81,14 +81,14 @@ public class MeteredKeyValueStore protected Sensor putSensor; private Sensor putIfAbsentSensor; protected Sensor getSensor; - private Sensor deleteSensor; + protected Sensor deleteSensor; private Sensor putAllSensor; private Sensor allSensor; private Sensor rangeSensor; private Sensor prefixScanSensor; private Sensor flushSensor; private Sensor e2eLatencySensor; - private InternalProcessorContext context; + protected InternalProcessorContext context; private StreamsMetricsImpl streamsMetrics; private TaskId taskId; diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStore.java new file mode 100644 index 0000000000000..ab14b23c2b634 --- /dev/null +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStore.java @@ -0,0 +1,227 @@ +/* + * 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 static org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl.maybeMeasureLatency; + +import java.util.Objects; +import org.apache.kafka.common.serialization.Serde; +import org.apache.kafka.common.utils.Time; +import org.apache.kafka.streams.errors.ProcessorStateException; +import org.apache.kafka.streams.processor.ProcessorContext; +import org.apache.kafka.streams.processor.StateStore; +import org.apache.kafka.streams.processor.StateStoreContext; +import org.apache.kafka.streams.processor.internals.SerdeGetter; +import org.apache.kafka.streams.query.Position; +import org.apache.kafka.streams.query.PositionBound; +import org.apache.kafka.streams.query.Query; +import org.apache.kafka.streams.query.QueryConfig; +import org.apache.kafka.streams.query.QueryResult; +import org.apache.kafka.streams.state.KeyValueStore; +import org.apache.kafka.streams.state.TimestampedKeyValueStore; +import org.apache.kafka.streams.state.ValueAndTimestamp; +import org.apache.kafka.streams.state.VersionedBytesStore; +import org.apache.kafka.streams.state.VersionedKeyValueStore; +import org.apache.kafka.streams.state.VersionedRecord; + +/** + * A metered {@link VersionedKeyValueStore} wrapper that is used for recording operation + * metrics, and hence its inner {@link VersionedBytesStore} implementation does not need to provide + * its own metrics collecting functionality. The inner {@code VersionedBytesStore} of this class + * is a {@link KeyValueStore} of type <Bytes,byte[]>, so we use {@link Serde}s + * to convert from <K,ValueAndTimestamp<V>> to <Bytes,byte[]>. In particular, + * {@link NullableValueAndTimestampSerde} is used since putting a tombstone to a versioned key-value + * store requires putting a null value associated with a timestamp. + * + * @param The key type + * @param The (raw) value type + */ +public class MeteredVersionedKeyValueStore + extends WrappedStateStore + implements VersionedKeyValueStore { + + private final MeteredVersionedKeyValueStoreInternal internal; + + MeteredVersionedKeyValueStore(final VersionedBytesStore inner, + final String metricScope, + final Time time, + final Serde keySerde, + final Serde> valueSerde) { + super(inner); + internal = new MeteredVersionedKeyValueStoreInternal(inner, metricScope, time, keySerde, valueSerde); + } + + /** + * Private helper class which represents the functionality of a {@link VersionedKeyValueStore} + * as a {@link TimestampedKeyValueStore} so that the bulk of the metering logic may be + * inherited from {@link MeteredKeyValueStore}. As a result, the implementation of + * {@link MeteredVersionedKeyValueStore} is a simple wrapper to translate from this + * {@link TimestampedKeyValueStore} representation of a versioned key-value store into the + * {@link VersionedKeyValueStore} interface itself. + */ + private class MeteredVersionedKeyValueStoreInternal + extends MeteredKeyValueStore> + implements TimestampedKeyValueStore { + + private final VersionedBytesStore inner; + + MeteredVersionedKeyValueStoreInternal(final VersionedBytesStore inner, + final String metricScope, + final Time time, + final Serde keySerde, + final Serde> valueSerde) { + super(inner, metricScope, time, keySerde, valueSerde); + this.inner = inner; + } + + @Override + public void put(final K key, final ValueAndTimestamp value) { + super.put( + key, + // versioned stores require a timestamp associated with all puts, including tombstones/deletes + value == null + ? ValueAndTimestamp.makeAllowNullable(null, context.timestamp()) + : value + ); + } + + public ValueAndTimestamp get(final K key, final long asOfTimestamp) { + Objects.requireNonNull(key, "key cannot be null"); + try { + return maybeMeasureLatency(() -> outerValue(inner.get(keyBytes(key), asOfTimestamp)), time, getSensor); + } catch (final ProcessorStateException e) { + final String message = String.format(e.getMessage(), key); + throw new ProcessorStateException(message, e); + } + } + + public ValueAndTimestamp delete(final K key, final long timestamp) { + Objects.requireNonNull(key, "key cannot be null"); + try { + return maybeMeasureLatency(() -> outerValue(inner.delete(keyBytes(key), timestamp)), time, deleteSensor); + } catch (final ProcessorStateException e) { + final String message = String.format(e.getMessage(), key); + throw new ProcessorStateException(message, e); + } + } + + @Override + public QueryResult query(final Query query, + final PositionBound positionBound, + final QueryConfig config) { + final long start = time.nanoseconds(); + final QueryResult result = wrapped().query(query, positionBound, config); + if (config.isCollectExecutionInfo()) { + result.addExecutionInfo( + "Handled in " + getClass() + " in " + (time.nanoseconds() - start) + "ns"); + } + // do not convert query or return types to/from inner bytes store to user-friendly types + // at this time, since we'll want a kip to align on what the types should be, and + // because we'll likely want to introduce a new return type which includes key, + // value, and timestamp for range queries + return result; + } + + @SuppressWarnings("unchecked") + @Override + protected Serde> prepareValueSerdeForStore( + final Serde> valueSerde, + final SerdeGetter getter + ) { + if (valueSerde == null) { + return new NullableValueAndTimestampSerde<>((Serde) getter.valueSerde()); + } else { + return super.prepareValueSerdeForStore(valueSerde, getter); + } + } + } + + @Override + public void put(final K key, final V value, final long timestamp) { + internal.put(key, ValueAndTimestamp.makeAllowNullable(value, timestamp)); + } + + @Override + public VersionedRecord delete(final K key, final long timestamp) { + final ValueAndTimestamp valueAndTimestamp = internal.delete(key, timestamp); + return valueAndTimestamp == null + ? null + : new VersionedRecord<>(valueAndTimestamp.value(), valueAndTimestamp.timestamp()); + } + + @Override + public VersionedRecord get(final K key) { + final ValueAndTimestamp valueAndTimestamp = internal.get(key); + return valueAndTimestamp == null + ? null + : new VersionedRecord<>(valueAndTimestamp.value(), valueAndTimestamp.timestamp()); + } + + @Override + public VersionedRecord get(final K key, final long asOfTimestamp) { + final ValueAndTimestamp valueAndTimestamp = internal.get(key, asOfTimestamp); + return valueAndTimestamp == null + ? null + : new VersionedRecord<>(valueAndTimestamp.value(), valueAndTimestamp.timestamp()); + } + + @Override + public String name() { + return internal.name(); + } + + @Deprecated + @Override + public void init(final ProcessorContext context, final StateStore root) { + internal.init(context, root); + } + + @Override + public void init(final StateStoreContext context, final StateStore root) { + internal.init(context, root); + } + + @Override + public void flush() { + internal.flush(); + } + + @Override + public void close() { + internal.close(); + } + + @Override + public boolean persistent() { + return internal.persistent(); + } + + @Override + public boolean isOpen() { + return internal.isOpen(); + } + + @Override + public QueryResult query(final Query query, final PositionBound positionBound, final QueryConfig config) { + return internal.query(query, positionBound, config); + } + + @Override + public Position getPosition() { + return internal.getPosition(); + } +} \ No newline at end of file diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStoreTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStoreTest.java new file mode 100644 index 0000000000000..1cc733762505d --- /dev/null +++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStoreTest.java @@ -0,0 +1,279 @@ +/* + * 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 static org.apache.kafka.common.utils.Utils.mkEntry; +import static org.apache.kafka.common.utils.Utils.mkMap; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.empty; +import static org.hamcrest.Matchers.greaterThan; +import static org.hamcrest.Matchers.is; +import static org.hamcrest.Matchers.not; +import static org.junit.Assert.assertThrows; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import org.apache.kafka.common.MetricName; +import org.apache.kafka.common.metrics.KafkaMetric; +import org.apache.kafka.common.metrics.Metrics; +import org.apache.kafka.common.metrics.Sensor; +import org.apache.kafka.common.serialization.Deserializer; +import org.apache.kafka.common.serialization.Serde; +import org.apache.kafka.common.serialization.Serdes.StringSerde; +import org.apache.kafka.common.serialization.Serializer; +import org.apache.kafka.common.utils.Bytes; +import org.apache.kafka.common.utils.MockTime; +import org.apache.kafka.common.utils.Time; +import org.apache.kafka.streams.StreamsConfig; +import org.apache.kafka.streams.processor.ProcessorContext; +import org.apache.kafka.streams.processor.StateStoreContext; +import org.apache.kafka.streams.processor.TaskId; +import org.apache.kafka.streams.processor.internals.InternalProcessorContext; +import org.apache.kafka.streams.processor.internals.ProcessorStateManager; +import org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl; +import org.apache.kafka.streams.state.ValueAndTimestamp; +import org.apache.kafka.streams.state.VersionedBytesStore; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.junit.MockitoJUnitRunner; + +@RunWith(MockitoJUnitRunner.StrictStubs.class) +public class MeteredVersionedKeyValueStoreTest { + + private static final String STORE_NAME = "versioned_store"; + private static final Serde STRING_SERDE = new StringSerde(); + private static final Serde> VALUE_AND_TIMESTAMP_SERDE = new NullableValueAndTimestampSerde<>(STRING_SERDE); + private static final String METRICS_SCOPE = "scope"; + private static final String STORE_LEVEL_GROUP = "stream-state-metrics"; + private static final String APPLICATION_ID = "test-app"; + private static final TaskId TASK_ID = new TaskId(0, 0, "My-Topology"); + + private static final String KEY = "k"; + private static final String VALUE = "v"; + private static final long TIMESTAMP = 10L; + private static final Bytes RAW_KEY = new Bytes(STRING_SERDE.serializer().serialize(null, KEY)); + private static final byte[] RAW_VALUE_AND_TIMESTAMP = VALUE_AND_TIMESTAMP_SERDE.serializer() + .serialize(null, ValueAndTimestamp.make(VALUE, TIMESTAMP)); + + private final VersionedBytesStore inner = mock(VersionedBytesStore.class); + private final Metrics metrics = new Metrics(); + private final Time mockTime = new MockTime(); + private final String threadId = Thread.currentThread().getName(); + private InternalProcessorContext context = mock(InternalProcessorContext.class); + private Map tags; + + private MeteredVersionedKeyValueStore store; + + @Before + public void setUp() { + when(inner.name()).thenReturn(STORE_NAME); + when(context.metrics()).thenReturn(new StreamsMetricsImpl(metrics, "test", StreamsConfig.METRICS_LATEST, mockTime)); + when(context.applicationId()).thenReturn(APPLICATION_ID); + when(context.taskId()).thenReturn(TASK_ID); + + metrics.config().recordLevel(Sensor.RecordingLevel.DEBUG); + tags = mkMap( + mkEntry("thread-id", threadId), + mkEntry("task-id", TASK_ID.toString()), + mkEntry(METRICS_SCOPE + "-state-id", STORE_NAME) + ); + + store = newMeteredStore(inner); + store.init((StateStoreContext) context, store); + } + + private MeteredVersionedKeyValueStore newMeteredStore(final VersionedBytesStore inner) { + return new MeteredVersionedKeyValueStore<>( + inner, + METRICS_SCOPE, + mockTime, + STRING_SERDE, + VALUE_AND_TIMESTAMP_SERDE + ); + } + + @SuppressWarnings("deprecation") + @Test + public void shouldDelegateDeprecatedInit() { + // recreate store in order to re-init + store.close(); + final VersionedBytesStore mockInner = mock(VersionedBytesStore.class); + store = newMeteredStore(mockInner); + + store.init((ProcessorContext) context, store); + + verify(mockInner).init((ProcessorContext) context, store); + } + + @Test + public void shouldDelegateInit() { + // init is already called in setUp() + verify(inner).init((StateStoreContext) context, store); + } + + @Test + public void shouldPassChangelogTopicNameToStateStoreSerde() { + final String changelogTopicName = "changelog-topic"; + when(context.changelogFor(STORE_NAME)).thenReturn(changelogTopicName); + doShouldPassChangelogTopicNameToStateStoreSerde(changelogTopicName); + } + + @Test + public void shouldPassDefaultChangelogTopicNameToStateStoreSerdeIfLoggingDisabled() { + final String defaultChangelogTopicName = ProcessorStateManager.storeChangelogTopic(APPLICATION_ID, STORE_NAME, TASK_ID.topologyName()); + when(context.changelogFor(STORE_NAME)).thenReturn(null); + doShouldPassChangelogTopicNameToStateStoreSerde(defaultChangelogTopicName); + } + + @SuppressWarnings("unchecked") + private void doShouldPassChangelogTopicNameToStateStoreSerde(final String changelogTopicName) { + // recreate store with mock serdes + final Serde keySerde = mock(Serde.class); + final Serializer keySerializer = mock(Serializer.class); + final Serde valueSerde = mock(Serde.class); + final Serializer valueSerializer = mock(Serializer.class); + final Deserializer valueDeserializer = mock(Deserializer.class); + when(keySerde.serializer()).thenReturn(keySerializer); + when(valueSerde.serializer()).thenReturn(valueSerializer); + when(valueSerde.deserializer()).thenReturn(valueDeserializer); + + store.close(); + store = new MeteredVersionedKeyValueStore<>( + inner, + METRICS_SCOPE, + mockTime, + keySerde, + new NullableValueAndTimestampSerde<>(valueSerde) + ); + store.init((StateStoreContext) context, store); + + store.put(KEY, VALUE, TIMESTAMP); + + verify(keySerializer).serialize(changelogTopicName, KEY); + verify(valueSerializer).serialize(changelogTopicName, VALUE); + } + + @Test + public void shouldRecordMetricsOnInit() { + // init is called in setUp(). it suffices to verify one restore metric since all restore + // metrics are recorded by the same sensor, and the sensor is tested elsewhere. + assertThat((Double) getMetric("restore-rate").metricValue(), greaterThan(0.0)); + } + + @Test + public void shouldDelegateAndRecordMetricsOnPut() { + store.put(KEY, VALUE, TIMESTAMP); + + verify(inner).put(RAW_KEY, RAW_VALUE_AND_TIMESTAMP); + assertThat((Double) getMetric("put-rate").metricValue(), greaterThan(0.0)); + } + + @Test + public void shouldDelegateAndRecordMetricsOnDelete() { + store.delete(KEY, TIMESTAMP); + + verify(inner).delete(RAW_KEY, TIMESTAMP); + assertThat((Double) getMetric("delete-rate").metricValue(), greaterThan(0.0)); + } + + @Test + public void shouldDelegateAndRecordMetricsOnGet() { + store.get(KEY); + + verify(inner).get(RAW_KEY); + assertThat((Double) getMetric("get-rate").metricValue(), greaterThan(0.0)); + } + + @Test + public void shouldDelegateAndRecordMetricsOnGetWithTimestamp() { + store.get(KEY, TIMESTAMP); + + verify(inner).get(RAW_KEY, TIMESTAMP); + assertThat((Double) getMetric("get-rate").metricValue(), greaterThan(0.0)); + } + + @Test + public void shouldDelegateAndRecordMetricsOnFlush() { + store.flush(); + + verify(inner).flush(); + assertThat((Double) getMetric("flush-rate").metricValue(), greaterThan(0.0)); + } + + @Test + public void shouldDelegateAndRemoveMetricsOnClose() { + assertThat(storeMetrics(), not(empty())); + + store.close(); + + verify(inner).close(); + assertThat(storeMetrics(), empty()); + } + + @Test + public void shouldRemoveMetricsOnCloseEvenIfInnerThrows() { + doThrow(new RuntimeException("uh oh")).when(inner).close(); + assertThat(storeMetrics(), not(empty())); + + assertThrows(RuntimeException.class, () -> store.close()); + + assertThat(storeMetrics(), empty()); + } + + @Test + public void shouldNotSetFlushListenerIfInnerIsNotCaching() { + assertThat(store.setFlushListener(null, false), is(false)); + } + + @Test + public void shouldThrowNullPointerOnPutIfKeyIsNull() { + assertThrows(NullPointerException.class, () -> store.put(null, VALUE, TIMESTAMP)); + } + + @Test + public void shouldThrowNullPointerOnDeleteIfKeyIsNull() { + assertThrows(NullPointerException.class, () -> store.delete(null, TIMESTAMP)); + } + + @Test + public void shouldThrowNullPointerOnGetIfKeyIsNull() { + assertThrows(NullPointerException.class, () -> store.get(null)); + } + + @Test + public void shouldThrowNullPointerOnGetWithTimestampIfKeyIsNull() { + assertThrows(NullPointerException.class, () -> store.get(null, TIMESTAMP)); + } + + private KafkaMetric getMetric(final String name) { + return metrics.metric(new MetricName(name, STORE_LEVEL_GROUP, "", tags)); + } + + private List storeMetrics() { + return metrics.metrics() + .keySet() + .stream() + .filter(name -> name.group().equals(STORE_LEVEL_GROUP) && name.tags().equals(tags)) + .collect(Collectors.toList()); + } +} \ No newline at end of file From 41ce134facd8afce486e5888b4a64f2812317725 Mon Sep 17 00:00:00 2001 From: Victoria Xia Date: Wed, 22 Feb 2023 11:01:40 -0800 Subject: [PATCH 2/4] review feedback --- .../MeteredVersionedKeyValueStore.java | 33 ++++++++++--------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStore.java index ab14b23c2b634..cfcb3ac29b48f 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStore.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStore.java @@ -66,16 +66,22 @@ public class MeteredVersionedKeyValueStore } /** - * Private helper class which represents the functionality of a {@link VersionedKeyValueStore} - * as a {@link TimestampedKeyValueStore} so that the bulk of the metering logic may be - * inherited from {@link MeteredKeyValueStore}. As a result, the implementation of - * {@link MeteredVersionedKeyValueStore} is a simple wrapper to translate from this - * {@link TimestampedKeyValueStore} representation of a versioned key-value store into the - * {@link VersionedKeyValueStore} interface itself. + * Conceptually, {@link MeteredVersionedKeyValueStore} should {@code extend} + * {@link MeteredKeyValueStore}, but due to type conflicts, we cannot do this. (Specifically, + * the first needs to be {@link VersionedKeyValueStore} while the second is {@link KeyValueStore} + * and the two interfaces conflict.) Thus, we use an internal instance of + * {@code MeteredKeyValueStore} to mimic inheritance instead. + *

+ * It's not ideal because it requires an extra step to translate between the APIs of + * {@link VersionedKeyValueStore} in {@link MeteredVersionedKeyValueStore} and + * the APIs of {@link TimestampedKeyValueStore} in {@link MeteredVersionedKeyValueStoreInternal}. + * This extra step is all that the methods of {@code MeteredVersionedKeyValueStoreInternal} do. + *

+ * Note that the addition of {@link #get(Object, long)} and {@link #delete(Object, long)} in + * this class are to match the interface of {@link VersionedKeyValueStore}. */ private class MeteredVersionedKeyValueStoreInternal - extends MeteredKeyValueStore> - implements TimestampedKeyValueStore { + extends MeteredKeyValueStore> { private final VersionedBytesStore inner; @@ -90,13 +96,10 @@ private class MeteredVersionedKeyValueStoreInternal @Override public void put(final K key, final ValueAndTimestamp value) { - super.put( - key, - // versioned stores require a timestamp associated with all puts, including tombstones/deletes - value == null - ? ValueAndTimestamp.makeAllowNullable(null, context.timestamp()) - : value - ); + if (value == null) { + throw new IllegalStateException("Versioned store requires timestamp associated with all puts, including tombstones/deletes"); + } + super.put(key, value); } public ValueAndTimestamp get(final K key, final long asOfTimestamp) { From 4f2f6b6a3e898d26e37955ffda15acab430d0026 Mon Sep 17 00:00:00 2001 From: Victoria Xia Date: Thu, 23 Feb 2023 13:50:59 -0800 Subject: [PATCH 3/4] reserve KeyQuery and RangeQuery query types --- .../state/internals/MeteredKeyValueStore.java | 12 +++---- .../MeteredVersionedKeyValueStore.java | 29 ++++++++-------- .../MeteredVersionedKeyValueStoreTest.java | 33 +++++++++++++++++++ 3 files changed, 54 insertions(+), 20 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredKeyValueStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredKeyValueStore.java index 362d08f254fa8..7144edb3e73cc 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredKeyValueStore.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredKeyValueStore.java @@ -246,9 +246,9 @@ public Position getPosition() { } @SuppressWarnings("unchecked") - private QueryResult runRangeQuery(final Query query, - final PositionBound positionBound, - final QueryConfig config) { + protected QueryResult runRangeQuery(final Query query, + final PositionBound positionBound, + final QueryConfig config) { final QueryResult result; final RangeQuery typedQuery = (RangeQuery) query; @@ -289,9 +289,9 @@ private QueryResult runRangeQuery(final Query query, @SuppressWarnings("unchecked") - private QueryResult runKeyQuery(final Query query, - final PositionBound positionBound, - final QueryConfig config) { + protected QueryResult runKeyQuery(final Query query, + final PositionBound positionBound, + final QueryConfig config) { final QueryResult result; final KeyQuery typedKeyQuery = (KeyQuery) query; final KeyQuery rawKeyQuery = diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStore.java index cfcb3ac29b48f..f6c132cf09c3a 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStore.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStore.java @@ -123,20 +123,21 @@ public ValueAndTimestamp delete(final K key, final long timestamp) { } @Override - public QueryResult query(final Query query, - final PositionBound positionBound, - final QueryConfig config) { - final long start = time.nanoseconds(); - final QueryResult result = wrapped().query(query, positionBound, config); - if (config.isCollectExecutionInfo()) { - result.addExecutionInfo( - "Handled in " + getClass() + " in " + (time.nanoseconds() - start) + "ns"); - } - // do not convert query or return types to/from inner bytes store to user-friendly types - // at this time, since we'll want a kip to align on what the types should be, and - // because we'll likely want to introduce a new return type which includes key, - // value, and timestamp for range queries - return result; + protected QueryResult runRangeQuery(final Query query, + final PositionBound positionBound, + final QueryConfig config) { + // throw exception for now to reserve the ability to implement this in the future + // without clashing with users' custom implementations in the meantime + throw new UnsupportedOperationException("Versioned stores do not support RangeQuery queries at this time."); + } + + @Override + protected QueryResult runKeyQuery(final Query query, + final PositionBound positionBound, + final QueryConfig config) { + // throw exception for now to reserve the ability to implement this in the future + // without clashing with users' custom implementations in the meantime + throw new UnsupportedOperationException("Versioned stores do not support KeyQuery queries at this time."); } @SuppressWarnings("unchecked") diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStoreTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStoreTest.java index 1cc733762505d..eddef6df9ba76 100644 --- a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStoreTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStoreTest.java @@ -24,6 +24,7 @@ import static org.hamcrest.Matchers.is; import static org.hamcrest.Matchers.not; import static org.junit.Assert.assertThrows; +import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; @@ -50,6 +51,12 @@ import org.apache.kafka.streams.processor.internals.InternalProcessorContext; import org.apache.kafka.streams.processor.internals.ProcessorStateManager; import org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl; +import org.apache.kafka.streams.query.KeyQuery; +import org.apache.kafka.streams.query.PositionBound; +import org.apache.kafka.streams.query.Query; +import org.apache.kafka.streams.query.QueryConfig; +import org.apache.kafka.streams.query.QueryResult; +import org.apache.kafka.streams.query.RangeQuery; import org.apache.kafka.streams.state.ValueAndTimestamp; import org.apache.kafka.streams.state.VersionedBytesStore; import org.junit.Before; @@ -265,6 +272,32 @@ public void shouldThrowNullPointerOnGetWithTimestampIfKeyIsNull() { assertThrows(NullPointerException.class, () -> store.get(null, TIMESTAMP)); } + @SuppressWarnings("unchecked") + @Test + public void shouldThrowOnIQv2RangeQuery() { + assertThrows(UnsupportedOperationException.class, () -> store.query(mock(RangeQuery.class), null, null)); + } + + @SuppressWarnings("unchecked") + @Test + public void shouldThrowOnIQv2KeyQuery() { + assertThrows(UnsupportedOperationException.class, () -> store.query(mock(KeyQuery.class), null, null)); + } + + @SuppressWarnings("unchecked") + @Test + public void shouldDelegateAndAddExecutionInfoOnCustomQuery() { + final Query query = mock(Query.class); + final PositionBound positionBound = mock(PositionBound.class); + final QueryConfig queryConfig = mock(QueryConfig.class); + final QueryResult result = mock(QueryResult.class); + when(inner.query(query, positionBound, queryConfig)).thenReturn(result); + when(queryConfig.isCollectExecutionInfo()).thenReturn(true); + + assertThat(store.query(query, positionBound, queryConfig), is(result)); + verify(result).addExecutionInfo(anyString()); + } + private KafkaMetric getMetric(final String name) { return metrics.metric(new MetricName(name, STORE_LEVEL_GROUP, "", tags)); } From b511994573524f2843907ceb02e4006dd139a066 Mon Sep 17 00:00:00 2001 From: Victoria Xia Date: Thu, 23 Feb 2023 21:36:49 -0800 Subject: [PATCH 4/4] review feedback --- .../MeteredVersionedKeyValueStoreTest.java | 57 +++++++++++++++++-- 1 file changed, 51 insertions(+), 6 deletions(-) diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStoreTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStoreTest.java index eddef6df9ba76..e850e1ee7deb6 100644 --- a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStoreTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStoreTest.java @@ -52,6 +52,7 @@ import org.apache.kafka.streams.processor.internals.ProcessorStateManager; import org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl; import org.apache.kafka.streams.query.KeyQuery; +import org.apache.kafka.streams.query.Position; import org.apache.kafka.streams.query.PositionBound; import org.apache.kafka.streams.query.Query; import org.apache.kafka.streams.query.QueryConfig; @@ -59,6 +60,7 @@ import org.apache.kafka.streams.query.RangeQuery; import org.apache.kafka.streams.state.ValueAndTimestamp; import org.apache.kafka.streams.state.VersionedBytesStore; +import org.apache.kafka.streams.state.VersionedRecord; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; @@ -197,25 +199,31 @@ public void shouldDelegateAndRecordMetricsOnPut() { @Test public void shouldDelegateAndRecordMetricsOnDelete() { - store.delete(KEY, TIMESTAMP); + when(inner.delete(RAW_KEY, TIMESTAMP)).thenReturn(RAW_VALUE_AND_TIMESTAMP); - verify(inner).delete(RAW_KEY, TIMESTAMP); + final VersionedRecord result = store.delete(KEY, TIMESTAMP); + + assertThat(result, is(new VersionedRecord<>(VALUE, TIMESTAMP))); assertThat((Double) getMetric("delete-rate").metricValue(), greaterThan(0.0)); } @Test public void shouldDelegateAndRecordMetricsOnGet() { - store.get(KEY); + when(inner.get(RAW_KEY)).thenReturn(RAW_VALUE_AND_TIMESTAMP); + + final VersionedRecord result = store.get(KEY); - verify(inner).get(RAW_KEY); + assertThat(result, is(new VersionedRecord<>(VALUE, TIMESTAMP))); assertThat((Double) getMetric("get-rate").metricValue(), greaterThan(0.0)); } @Test public void shouldDelegateAndRecordMetricsOnGetWithTimestamp() { - store.get(KEY, TIMESTAMP); + when(inner.get(RAW_KEY, TIMESTAMP)).thenReturn(RAW_VALUE_AND_TIMESTAMP); + + final VersionedRecord result = store.get(KEY, TIMESTAMP); - verify(inner).get(RAW_KEY, TIMESTAMP); + assertThat(result, is(new VersionedRecord<>(VALUE, TIMESTAMP))); assertThat((Double) getMetric("get-rate").metricValue(), greaterThan(0.0)); } @@ -298,6 +306,43 @@ public void shouldDelegateAndAddExecutionInfoOnCustomQuery() { verify(result).addExecutionInfo(anyString()); } + @Test + public void shouldDelegateName() { + when(inner.name()).thenReturn(STORE_NAME); + + assertThat(store.name(), is(STORE_NAME)); + } + + @Test + public void shouldDelegatePersistent() { + // `persistent = true` case + when(inner.persistent()).thenReturn(true); + assertThat(store.persistent(), is(true)); + + // `persistent = false` case + when(inner.persistent()).thenReturn(false); + assertThat(store.persistent(), is(false)); + } + + @Test + public void shouldDelegateIsOpen() { + // `isOpen = true` case + when(inner.isOpen()).thenReturn(true); + assertThat(store.isOpen(), is(true)); + + // `isOpen = false` case + when(inner.isOpen()).thenReturn(false); + assertThat(store.isOpen(), is(false)); + } + + @Test + public void shouldDelegateGetPosition() { + final Position position = mock(Position.class); + when(inner.getPosition()).thenReturn(position); + + assertThat(store.getPosition(), is(position)); + } + private KafkaMetric getMetric(final String name) { return metrics.metric(new MetricName(name, STORE_LEVEL_GROUP, "", tags)); }