diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMap.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMap.java index 1a4ab843c516..fabaff6014d4 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMap.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMap.java @@ -19,7 +19,13 @@ import java.io.Closeable; import java.io.IOException; -import java.util.*; +import java.util.AbstractMap; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; import java.util.concurrent.ExecutionException; import java.util.concurrent.Future; import java.util.function.Function; @@ -34,8 +40,6 @@ import org.apache.beam.sdk.util.Weighted; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; -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.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets; @@ -53,13 +57,53 @@ public class WindmillMap extends AbstractWindmillMap { private final String stateFamily; private final Coder keyCoder; private final Coder valueCoder; + + private static class StructuralKey { + private final K originalKey; + private final @Nullable Object structuralValue; + + private StructuralKey(K originalKey, @Nullable Object structuralValue) { + this.originalKey = originalKey; + this.structuralValue = structuralValue; + } + + public static StructuralKey of(K key, Coder coder) { + return new StructuralKey<>(key, key == null ? null : coder.structuralValue(key)); + } + + public K getOriginalKey() { + return originalKey; + } + + @Override + public boolean equals(@Nullable Object o) { + if (this == o) { + return true; + } + if (!(o instanceof StructuralKey)) { + return false; + } + StructuralKey that = (StructuralKey) o; + return Objects.equals(structuralValue, that.structuralValue); + } + + @Override + public int hashCode() { + return Objects.hashCode(structuralValue); + } + } + + private StructuralKey asStructuralKey(K key) { + return StructuralKey.of(key, keyCoder); + } + // TODO(reuvenlax): Should we evict items from the cache? We would have to make sure // that anything in the cache that is not committed is not evicted. negativeCache could be // evicted whenever we want. - private final Map cachedValues = Maps.newHashMap(); - private final Set negativeCache = Sets.newHashSet(); - private final Set localAdditions = Sets.newHashSet(); - private final Set localRemovals = Sets.newHashSet(); + private final Map, V> cachedValues = Maps.newHashMap(); + private final Set> negativeCache = Sets.newHashSet(); + private final Set> localAdditions = Sets.newHashSet(); + private final Set> localRemovals = Sets.newHashSet(); private boolean complete; private boolean cleared = false; @@ -110,10 +154,10 @@ protected Windmill.WorkItemCommitRequest persistDirectly(WindmillStateCache.ForK } cleared = false; - for (K key : localAdditions) { - ByteString keyBytes = protoKeyFromUserKey(key); + for (StructuralKey sk : localAdditions) { + ByteString keyBytes = protoKeyFromUserKey(sk.getOriginalKey()); ByteStringOutputStream valueStream = new ByteStringOutputStream(); - valueCoder.encode(cachedValues.get(key), valueStream, Coder.Context.OUTER); + valueCoder.encode(cachedValues.get(sk), valueStream, Coder.Context.OUTER); ByteString valueBytes = valueStream.toByteString(); commitBuilder @@ -126,19 +170,13 @@ protected Windmill.WorkItemCommitRequest persistDirectly(WindmillStateCache.ForK } localAdditions.clear(); - for (K key : localRemovals) { + for (StructuralKey sk : localRemovals) { ByteStringOutputStream keyStream = new ByteStringOutputStream(); stateKeyPrefix.byteString().writeTo(keyStream); - keyCoder.encode(key, keyStream, Coder.Context.OUTER); + keyCoder.encode(sk.getOriginalKey(), keyStream, Coder.Context.OUTER); ByteString keyBytes = keyStream.toByteString(); // Leaving data blank means that we delete the tag. commitBuilder.addValueUpdatesBuilder().setTag(keyBytes).setStateFamily(stateFamily); - - V cachedValue = cachedValues.remove(key); - if (cachedValue != null) { - ByteStringOutputStream valueStream = new ByteStringOutputStream(); - valueCoder.encode(cachedValues.get(key), valueStream, Coder.Context.OUTER); - } } negativeCache.addAll(localRemovals); localRemovals.clear(); @@ -196,36 +234,50 @@ public ReadableState isEmpty() { @Override public void put(K key, V value) { - V oldValue = cachedValues.put(key, value); - if (valueCoder.consistentWithEquals() && value.equals(oldValue)) { + StructuralKey structKey = asStructuralKey(key); + boolean wasPresent = cachedValues.containsKey(structKey); + V oldValue = cachedValues.put(structKey, value); + if (wasPresent + && valueCoder.consistentWithEquals() + && Objects.equals(value, oldValue) + && !localRemovals.contains(structKey) + && !negativeCache.contains(structKey)) { return; } - localAdditions.add(key); - localRemovals.remove(key); - negativeCache.remove(key); + localAdditions.remove(structKey); + localAdditions.add(structKey); + localRemovals.remove(structKey); + negativeCache.remove(structKey); + } + + private ReadableState computeAndPut(K key, Function mappingFunction) { + V computed = mappingFunction.apply(key); + put(key, computed); + return ReadableStates.immediate(null); } @Override public @UnknownKeyFor @NonNull @Initialized ReadableState computeIfAbsent( K key, Function mappingFunction) { + StructuralKey structKey = asStructuralKey(key); + if (localRemovals.contains(structKey) || negativeCache.contains(structKey)) { + return computeAndPut(key, mappingFunction); + } + @Nullable V cachedValue = cachedValues.get(structKey); + if (cachedValue != null) { + return ReadableStates.immediate(cachedValue); + } + if (complete) { + return computeAndPut(key, mappingFunction); + } + Future persistedData = getFutureForKey(key); try (Closeable scope = scopedReadState()) { - if (localRemovals.contains(key) || negativeCache.contains(key)) { - return ReadableStates.immediate(null); - } - @Nullable V cachedValue = cachedValues.get(key); - if (cachedValue != null || complete) { - return ReadableStates.immediate(cachedValue); - } - V persistedValue = persistedData.get(); if (persistedValue == null) { - // This is a new value. Add it to the map and return null. - put(key, mappingFunction.apply(key)); - return ReadableStates.immediate(null); + return computeAndPut(key, mappingFunction); } - // TODO: Don't do this if it was already in cache. - cachedValues.put(key, persistedValue); + cachedValues.put(structKey, persistedValue); return ReadableStates.immediate(persistedValue); } catch (InterruptedException | ExecutionException | IOException e) { if (e instanceof InterruptedException) { @@ -237,9 +289,10 @@ public void put(K key, V value) { @Override public void remove(K key) { - if (localRemovals.add(key)) { - cachedValues.remove(key); - localAdditions.remove(key); + StructuralKey structKey = asStructuralKey(key); + if (localRemovals.add(structKey)) { + cachedValues.remove(structKey); + localAdditions.remove(structKey); } } @@ -311,12 +364,20 @@ public WindmillMapValuesReadableState(ReadableState>> e } } + private Iterable> userEntries() { + List> result = new ArrayList<>(cachedValues.size()); + for (Map.Entry, V> entry : cachedValues.entrySet()) { + result.add(new AbstractMap.SimpleEntry<>(entry.getKey().getOriginalKey(), entry.getValue())); + } + return Collections.unmodifiableList(result); + } + private class WindmillMapEntriesReadableState implements ReadableState>> { @Override public Iterable> read() { if (complete) { - return ImmutableMap.copyOf(cachedValues).entrySet(); + return userEntries(); } Future>> persistedData = getFuture(); try (Closeable scope = scopedReadState()) { @@ -337,26 +398,31 @@ public Iterable> read() { // This is a known amount of data. Cache it all. transformedData.forEach( e -> { - // The cached data overrides what is read from state, so call putIfAbsent. - cachedValues.putIfAbsent(e.getKey(), e.getValue()); + StructuralKey structKey = asStructuralKey(e.getKey()); + if (!localRemovals.contains(structKey)) { + // The cached data overrides what is read from state, so call putIfAbsent. + cachedValues.putIfAbsent(structKey, e.getValue()); + } }); complete = true; - return ImmutableMap.copyOf(cachedValues).entrySet(); + return userEntries(); } else { - ImmutableMap cachedCopy = ImmutableMap.copyOf(cachedValues); - ImmutableSet removalCopy = ImmutableSet.copyOf(localRemovals); + Map, V> cachedCopy = Maps.newHashMap(cachedValues); + Set> removalCopy = Sets.newHashSet(localRemovals); // This means that the result might be too large to cache, so don't add it to the - // local cache. Instead merge the iterables, giving priority to any local additions + // local cache. Instead, merge the iterables, giving priority to any local additions // (represented in cachedCopy and removalCopy) that may not have been committed // yet. return Iterables.unmodifiableIterable( Iterables.concat( - cachedCopy.entrySet(), + userEntries(), Iterables.filter( transformedData, - e -> - !cachedCopy.containsKey(e.getKey()) - && !removalCopy.contains(e.getKey())))); + e -> { + StructuralKey structKey = asStructuralKey(e.getKey()); + return !cachedCopy.containsKey(structKey) + && !removalCopy.contains(structKey); + }))); } } catch (InterruptedException | ExecutionException | IOException e) { @@ -383,12 +449,21 @@ private class WindmillMapIsEmptyReadableState implements ReadableState @Override public @Nullable Boolean read() { + // If there are any cached values, we know the map is not empty. + if (!cachedValues.isEmpty()) { + return false; + } + if (complete) { + return true; + } return Iterables.isEmpty(keys.read()); } @Override public @UnknownKeyFor @NonNull @Initialized ReadableState readLater() { - keys.readLater(); + if (cachedValues.isEmpty() && !complete) { + keys.readLater(); + } return this; } } @@ -404,22 +479,25 @@ public WindmillMapReadResultReadableState(K key, @Nullable V defaultValue) { @Override public @Nullable V read() { + StructuralKey structKey = asStructuralKey(key); + if (localRemovals.contains(structKey) || negativeCache.contains(structKey)) { + return defaultValue; + } + if (cachedValues.containsKey(structKey)) { + return cachedValues.get(structKey); + } + if (complete) { + return defaultValue; + } + Future persistedData = getFutureForKey(key); try (Closeable scope = scopedReadState()) { - if (localRemovals.contains(key) || negativeCache.contains(key)) { - return null; - } - @Nullable V cachedValue = cachedValues.get(key); - if (cachedValue != null || complete) { - return cachedValue; - } - V persistedValue = persistedData.get(); if (persistedValue == null) { - negativeCache.add(key); + negativeCache.add(structKey); return defaultValue; } - cachedValues.put(key, persistedValue); + cachedValues.put(structKey, persistedValue); return persistedValue; } catch (InterruptedException | ExecutionException | IOException e) { if (e instanceof InterruptedException) { @@ -432,7 +510,13 @@ public WindmillMapReadResultReadableState(K key, @Nullable V defaultValue) { @Override @SuppressWarnings("FutureReturnValueIgnored") public @UnknownKeyFor @NonNull @Initialized ReadableState readLater() { - WindmillMap.this.getFutureForKey(key); + StructuralKey structKey = asStructuralKey(key); + if (!localRemovals.contains(structKey) + && !negativeCache.contains(structKey) + && !cachedValues.containsKey(structKey) + && !complete) { + WindmillMap.this.getFutureForKey(key); + } return this; } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimap.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimap.java index e24e449d2598..c9e796720767 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimap.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimap.java @@ -62,17 +62,18 @@ public void put(KeyT key, ValueT value) { multimap.put(key, value); } + @SuppressWarnings("nullness") @Override public ReadableState computeIfAbsent( KeyT key, Function mappingFunction) { // Note that computeIfAbsent comments indicate that the read is lazy but this matches the - // existing eager - // behavior of WindmillMap. + // existing eager behavior of WindmillMap. Iterable existingValues = multimap.get(key).read(); - if (Iterables.isEmpty(existingValues)) { + Iterator iterator = existingValues.iterator(); + if (!iterator.hasNext()) { ValueT inserted = mappingFunction.apply(key); multimap.put(key, inserted); - return ReadableStates.immediate(inserted); + return ReadableStates.immediate(null); } else { return ReadableStates.immediate(Iterables.getOnlyElement(existingValues)); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillSet.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillSet.java index ee7e6862c7a1..06a0b5292cb7 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillSet.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillSet.java @@ -19,7 +19,6 @@ import java.io.Closeable; import java.io.IOException; -import java.util.Optional; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.sdk.state.ReadableState; import org.apache.beam.sdk.state.SetState; @@ -110,7 +109,7 @@ public WindmillSetAddIfAbsentReadableState(K k) { @Override public Boolean read() { - return Optional.ofNullable(putState.read()).orElse(false); + return putState.read() == null; } @Override diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapTest.java new file mode 100644 index 000000000000..f04007bfa96b --- /dev/null +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapTest.java @@ -0,0 +1,533 @@ +/* + * 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.dataflow.worker.windmill.state; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.when; + +import java.io.Closeable; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.AbstractMap; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicBoolean; +import javax.annotation.Nullable; +import org.apache.beam.runners.core.StateNamespace; +import org.apache.beam.runners.core.StateNamespaceForTest; +import org.apache.beam.runners.core.StateTag; +import org.apache.beam.runners.core.StateTags; +import org.apache.beam.runners.dataflow.options.DataflowWorkerHarnessOptions; +import org.apache.beam.runners.dataflow.worker.WindmillComputationKey; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill; +import org.apache.beam.sdk.coders.ByteArrayCoder; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.Coder.Context; +import org.apache.beam.sdk.coders.NullableCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.state.MapState; +import org.apache.beam.sdk.state.ReadableState; +import org.apache.beam.sdk.util.ByteStringOutputStream; +import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Supplier; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Futures; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.SettableFuture; +import org.hamcrest.Matchers; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.Timeout; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.MockitoAnnotations; + +/** Unit tests demonstrating correctness issues and bugs in {@link WindmillMap}. */ +@RunWith(JUnit4.class) +public class WindmillMapTest { + @Rule public transient Timeout globalTimeout = Timeout.seconds(600); + + private static final String STATE_FAMILY = "family"; + private static final StateNamespace NAMESPACE = new StateNamespaceForTest("ns"); + + @Mock private WindmillStateReader mockReader; + @Mock private Supplier readStateSupplier; + + private WindmillStateCache cache; + private WindmillStateInternals underTest; + + private static ByteString key(StateNamespace namespace, String addrId) { + return ByteString.copyFromUtf8(namespace.stringKey() + "+u" + addrId); + } + + private ByteString protoKey(String tagId, @Nullable K tag, Coder keyCoder) + throws IOException { + ByteStringOutputStream keyStream = new ByteStringOutputStream(); + key(NAMESPACE, tagId).writeTo(keyStream); + if (tag != null) { + keyCoder.encode(tag, keyStream, Context.OUTER); + } + return keyStream.toByteString(); + } + + private ByteString protoKey(@Nullable K tag, Coder keyCoder) throws IOException { + return protoKey("map", tag, keyCoder); + } + + private static ByteString encodeValue(T value, Coder coder) throws IOException { + ByteStringOutputStream out = new ByteStringOutputStream(); + coder.encode(value, out, Context.OUTER); + return out.toByteString(); + } + + private static byte[] dup(byte[] key) { + byte[] res = new byte[key.length]; + System.arraycopy(key, 0, res, 0, key.length); + return res; + } + + @SafeVarargs + private static List weightedList(T... entries) { + WeightedList list = new WeightedList<>(new ArrayList<>()); + for (T entry : entries) { + list.addWeighted(entry, 1); + } + return list; + } + + @Before + public void setUp() { + MockitoAnnotations.initMocks(this); + DataflowWorkerHarnessOptions options = + PipelineOptionsFactory.as(DataflowWorkerHarnessOptions.class); + cache = WindmillStateCache.builder().setSizeMb(options.getWorkerCacheMb()).build(); + + WindmillTagEncoding tagEncoding = WindmillTagEncodingV1.instance(); + underTest = + new WindmillStateInternals<>( + "key", + STATE_FAMILY, + mockReader, + false, + cache + .forComputation("comp") + .forKey( + WindmillComputationKey.create( + "comp", ByteString.copyFrom("key", StandardCharsets.UTF_8), 123), + 17L, + 1L) + .forFamily(STATE_FAMILY), + tagEncoding, + readStateSupplier); + + when(mockReader.valuePrefixFuture(any(), any(), any())) + .thenReturn(Futures.immediateFuture(Collections.emptyList())); + } + + /** + * Tests that {@link WindmillMap#computeIfAbsent} correctly computes and inserts when the key was + * previously read as absent (and therefore placed in {@code negativeCache}). + */ + @Test + public void testComputeIfAbsentInsertsWhenNegativelyCached() throws Exception { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTest.state(NAMESPACE, addr); + + final String key = "tag1"; + SettableFuture future = SettableFuture.create(); + when(mockReader.valueFuture( + protoKey(key, StringUtf8Coder.of()), STATE_FAMILY, VarIntCoder.of())) + .thenReturn(future); + future.set(null); + + // First, read the key -> returns null and populates negativeCache + assertNull(mapState.get(key).read()); + + // Now, call computeIfAbsent on the absent key + AtomicBoolean mappingFunctionCalled = new AtomicBoolean(false); + ReadableState result = + mapState.computeIfAbsent( + key, + k -> { + mappingFunctionCalled.set(true); + return 42; + }); + + assertNull("computeIfAbsent returns null on new insertion", result.read()); + assertTrue( + "computeIfAbsent must invoke mappingFunction when key was in negativeCache", + mappingFunctionCalled.get()); + assertEquals(Integer.valueOf(42), mapState.get(key).read()); + } + + /** + * Tests that {@link WindmillMap#computeIfAbsent} correctly computes and inserts when the key was + * previously removed in the same work item (and therefore placed in {@code localRemovals}). + */ + @Test + public void testComputeIfAbsentInsertsWhenLocallyRemoved() { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTest.state(NAMESPACE, addr); + + final String key = "tag1"; + mapState.put(key, 10); + mapState.remove(key); + + AtomicBoolean mappingFunctionCalled = new AtomicBoolean(false); + ReadableState result = + mapState.computeIfAbsent( + key, + k -> { + mappingFunctionCalled.set(true); + return 42; + }); + + assertNull("computeIfAbsent returns null on new insertion", result.read()); + assertTrue( + "computeIfAbsent must invoke mappingFunction when key was in localRemovals", + mappingFunctionCalled.get()); + assertEquals(Integer.valueOf(42), mapState.get(key).read()); + } + + /** + * Tests that {@link WindmillMap#computeIfAbsent} correctly computes and inserts when {@code + * complete == true} (e.g. after {@link MapState#clear}). + */ + @Test + public void testComputeIfAbsentInsertsWhenComplete() { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTest.state(NAMESPACE, addr); + + mapState.clear(); + + AtomicBoolean mappingFunctionCalled = new AtomicBoolean(false); + ReadableState result = + mapState.computeIfAbsent( + "newKey", + k -> { + mappingFunctionCalled.set(true); + return 42; + }); + + assertNull("computeIfAbsent returns null on new insertion", result.read()); + assertTrue( + "computeIfAbsent must invoke mappingFunction when complete == true", + mappingFunctionCalled.get()); + assertEquals(Integer.valueOf(42), mapState.get("newKey").read()); + } + + /** + * Tests that {@link WindmillMap#getOrDefault} returns {@code defaultValue} on repeated lookups. + */ + @Test + public void testGetOrDefaultReturnsDefaultValueOnRepeatedRead() throws Exception { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTest.state(NAMESPACE, addr); + + final String key1 = "key1"; + SettableFuture future = SettableFuture.create(); + when(mockReader.valueFuture( + protoKey(key1, StringUtf8Coder.of()), STATE_FAMILY, VarIntCoder.of())) + .thenReturn(future); + future.set(null); + + // 1st read: returns defaultValue 99 + assertEquals(Integer.valueOf(99), mapState.getOrDefault(key1, 99).read()); + + // 2nd read: key is now in negativeCache, must still return defaultValue 99 + assertEquals(Integer.valueOf(99), mapState.getOrDefault(key1, 99).read()); + } + + /** + * Tests that {@link WindmillMap#getOrDefault} returns {@code defaultValue} after remove and + * clear. + */ + @Test + public void testGetOrDefaultReturnsDefaultValueAfterRemoveAndClear() { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTest.state(NAMESPACE, addr); + + // After remove: + mapState.put("key2", 1); + mapState.remove("key2"); + assertEquals(Integer.valueOf(99), mapState.getOrDefault("key2", 99).read()); + + // After clear: + mapState.clear(); + assertEquals(Integer.valueOf(99), mapState.getOrDefault("key3", 99).read()); + } + + /** + * Tests that {@link WindmillMap#entries} does not resurrect locally deleted entries when the + * response from Windmill storage is {@link org.apache.beam.sdk.util.Weighted}. + */ + @Test + public void testEntriesDoesNotResurrectDeletedEntryWhenWeighted() throws Exception { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTest.state(NAMESPACE, addr); + + final String tag1 = "tag1"; + final String tag2 = "tag2"; + SettableFuture>> prefixFuture = SettableFuture.create(); + when(mockReader.valuePrefixFuture( + protoKey(null, StringUtf8Coder.of()), STATE_FAMILY, VarIntCoder.of())) + .thenReturn(prefixFuture); + + // Delete tag1 locally before reading entries + mapState.remove(tag1); + + // Windmill storage returns weighted list containing tag1 and tag2 + prefixFuture.set( + weightedList( + new AbstractMap.SimpleEntry<>(protoKey(tag1, StringUtf8Coder.of()), 1), + new AbstractMap.SimpleEntry<>(protoKey(tag2, StringUtf8Coder.of()), 2))); + + Iterable> entries = mapState.entries().read(); + + assertThat(entries, Matchers.not(Matchers.hasItem(new AbstractMap.SimpleEntry<>(tag1, 1)))); + assertThat(entries, Matchers.hasItem(new AbstractMap.SimpleEntry<>(tag2, 2))); + } + + /** + * Tests that {@link WindmillMap#persistDirectly} succeeds and commits deletion of a removed key + * without exception. + */ + @Test + public void testPersistSucceedsAfterRemovalAndWeightedEntriesRead() throws Exception { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTest.state(NAMESPACE, addr); + + final String tag1 = "tag1"; + SettableFuture>> prefixFuture = SettableFuture.create(); + when(mockReader.valuePrefixFuture( + protoKey(null, StringUtf8Coder.of()), STATE_FAMILY, VarIntCoder.of())) + .thenReturn(prefixFuture); + + // 1. Remove tag1 locally (adds to localRemovals) + mapState.remove(tag1); + + // 2. Read entries (weighted) -> does not resurrect tag1 into cachedValues + prefixFuture.set( + weightedList(new AbstractMap.SimpleEntry<>(protoKey(tag1, StringUtf8Coder.of()), 1))); + mapState.entries().read(); + + // 3. Persist state: persistDirectly commits tag1 deletion successfully + Windmill.WorkItemCommitRequest.Builder commitBuilder = + Windmill.WorkItemCommitRequest.newBuilder(); + underTest.persist(commitBuilder); + + assertEquals(1, commitBuilder.getValueUpdatesCount()); + assertEquals(protoKey(tag1, StringUtf8Coder.of()), commitBuilder.getValueUpdates(0).getTag()); + assertTrue(commitBuilder.getValueUpdates(0).getValue().getData().isEmpty()); + } + + /** + * Tests that {@link WindmillMap} matches {@code byte[]} keys using {@link Coder#structuralValue}. + */ + @Test + public void testByteArrayKeysMatchCorrectly() throws Exception { + StateTag> addr = + StateTags.map("map", ByteArrayCoder.of(), VarIntCoder.of()); + MapState mapState = underTest.state(NAMESPACE, addr); + + byte[] key = new byte[] {1, 2, 3}; + mapState.put(key, 42); + + // Lookup with duplicate byte array (same bytes, distinct object instance) + byte[] lookupKey = dup(key); + + assertEquals(Integer.valueOf(42), mapState.get(lookupKey).read()); + Mockito.verifyNoInteractions(mockReader); + } + + /** Tests that {@link WindmillMap#put} supports null values when using {@link NullableCoder}. */ + @Test + public void testPutNullValueSupported() { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), NullableCoder.of(StringUtf8Coder.of())); + MapState mapState = underTest.state(NAMESPACE, addr); + + mapState.put("key", null); + assertNull(mapState.get("key").read()); + } + + /** Tests {@link WindmillMap#isEmpty()}. */ + @Test + public void testIsEmpty() { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTest.state(NAMESPACE, addr); + + mapState.put("key1", 10); + assertFalse(mapState.isEmpty().read()); + + mapState.clear(); + assertTrue(mapState.isEmpty().read()); + } + + /** Tests that persistDirectly with no modifications produces no updates. */ + @Test + public void testPersistDirectlyNoChangesReturnsEmpty() { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + underTest.state(NAMESPACE, addr); + + Windmill.WorkItemCommitRequest.Builder commitBuilder = + Windmill.WorkItemCommitRequest.newBuilder(); + underTest.persist(commitBuilder); + + assertEquals(0, commitBuilder.getValueUpdatesCount()); + assertEquals(0, commitBuilder.getTagValuePrefixDeletesCount()); + } + + /** Tests that persistDirectly commits local additions correctly and clears localAdditions. */ + @Test + public void testPersistDirectlyLocalAdditions() throws Exception { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTest.state(NAMESPACE, addr); + + mapState.put("k1", 100); + mapState.put("k2", 200); + + Windmill.WorkItemCommitRequest.Builder commitBuilder = + Windmill.WorkItemCommitRequest.newBuilder(); + underTest.persist(commitBuilder); + + assertEquals(2, commitBuilder.getValueUpdatesCount()); + ByteString k1Key = protoKey("k1", StringUtf8Coder.of()); + ByteString k2Key = protoKey("k2", StringUtf8Coder.of()); + ByteString k1Val = encodeValue(100, VarIntCoder.of()); + ByteString k2Val = encodeValue(200, VarIntCoder.of()); + + Map updates = new java.util.HashMap<>(); + for (Windmill.TagValue valUpdate : commitBuilder.getValueUpdatesList()) { + assertEquals(STATE_FAMILY, valUpdate.getStateFamily()); + assertEquals(Long.MAX_VALUE, valUpdate.getValue().getTimestamp()); + updates.put(valUpdate.getTag(), valUpdate.getValue().getData()); + } + assertEquals(k1Val, updates.get(k1Key)); + assertEquals(k2Val, updates.get(k2Key)); + + // A second persist should produce no new changes since localAdditions was cleared + Windmill.WorkItemCommitRequest.Builder secondCommit = + Windmill.WorkItemCommitRequest.newBuilder(); + underTest.persist(secondCommit); + assertEquals(0, secondCommit.getValueUpdatesCount()); + } + + /** Tests that persistDirectly commits local removals with empty data (tag delete). */ + @Test + public void testPersistDirectlyLocalRemovals() throws Exception { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTest.state(NAMESPACE, addr); + + mapState.remove("k1"); + + Windmill.WorkItemCommitRequest.Builder commitBuilder = + Windmill.WorkItemCommitRequest.newBuilder(); + underTest.persist(commitBuilder); + + assertEquals(1, commitBuilder.getValueUpdatesCount()); + Windmill.TagValue valUpdate = commitBuilder.getValueUpdates(0); + assertEquals(protoKey("k1", StringUtf8Coder.of()), valUpdate.getTag()); + assertEquals(STATE_FAMILY, valUpdate.getStateFamily()); + assertTrue("Empty data indicates deletion", valUpdate.getValue().getData().isEmpty()); + } + + /** Tests that persistDirectly with clear() commits tag value prefix deletes and additions. */ + @Test + public void testPersistDirectlyClearedAndLocalAdditions() throws Exception { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTest.state(NAMESPACE, addr); + + mapState.clear(); + mapState.put("k1", 42); + + Windmill.WorkItemCommitRequest.Builder commitBuilder = + Windmill.WorkItemCommitRequest.newBuilder(); + underTest.persist(commitBuilder); + + assertEquals(1, commitBuilder.getTagValuePrefixDeletesCount()); + assertEquals(STATE_FAMILY, commitBuilder.getTagValuePrefixDeletes(0).getStateFamily()); + assertEquals(key(NAMESPACE, "map"), commitBuilder.getTagValuePrefixDeletes(0).getTagPrefix()); + + assertEquals(1, commitBuilder.getValueUpdatesCount()); + Windmill.TagValue valUpdate = commitBuilder.getValueUpdates(0); + assertEquals(protoKey("k1", StringUtf8Coder.of()), valUpdate.getTag()); + assertEquals(encodeValue(42, VarIntCoder.of()), valUpdate.getValue().getData()); + } + + /** Tests basic put, get, overwrite, and remove on {@link WindmillMap}. */ + @Test + public void testBasicPutGetOverwriteAndRemove() { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTest.state(NAMESPACE, addr); + + mapState.put("a", 1); + assertEquals(Integer.valueOf(1), mapState.get("a").read()); + + mapState.put("a", 2); + assertEquals(Integer.valueOf(2), mapState.get("a").read()); + + mapState.remove("a"); + assertNull(mapState.get("a").read()); + } + + /** Tests keys, values, and entries on {@link WindmillMap}. */ + @Test + public void testKeysValuesAndEntries() throws Exception { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTest.state(NAMESPACE, addr); + + when(mockReader.valuePrefixFuture( + protoKey(null, StringUtf8Coder.of()), STATE_FAMILY, VarIntCoder.of())) + .thenReturn(Futures.immediateFuture(Collections.emptyList())); + + mapState.put("k1", 10); + mapState.put("k2", 20); + + assertThat(mapState.keys().read(), Matchers.containsInAnyOrder("k1", "k2")); + assertThat(mapState.values().read(), Matchers.containsInAnyOrder(10, 20)); + assertThat( + mapState.entries().read(), + Matchers.containsInAnyOrder( + new AbstractMap.SimpleEntry<>("k1", 10), new AbstractMap.SimpleEntry<>("k2", 20))); + } +} diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimapTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimapTest.java new file mode 100644 index 000000000000..e76f4e228e35 --- /dev/null +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimapTest.java @@ -0,0 +1,285 @@ +/* + * 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.dataflow.worker.windmill.state; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.Mockito.when; + +import java.io.Closeable; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.AbstractMap; +import java.util.Arrays; +import java.util.Collections; +import java.util.concurrent.atomic.AtomicBoolean; +import org.apache.beam.runners.core.StateNamespace; +import org.apache.beam.runners.core.StateNamespaceForTest; +import org.apache.beam.runners.core.StateTag; +import org.apache.beam.runners.core.StateTags; +import org.apache.beam.runners.dataflow.options.DataflowWorkerHarnessOptions; +import org.apache.beam.runners.dataflow.worker.WindmillComputationKey; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.Coder.Context; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.state.MapState; +import org.apache.beam.sdk.state.ReadableState; +import org.apache.beam.sdk.util.ByteStringOutputStream; +import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Supplier; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Futures; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.SettableFuture; +import org.hamcrest.Matchers; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.Timeout; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +/** + * Unit tests demonstrating correctness issues and contract violations in {@link + * WindmillMapViaMultimap}. + */ +@RunWith(JUnit4.class) +public class WindmillMapViaMultimapTest { + @Rule public transient Timeout globalTimeout = Timeout.seconds(600); + + private static final String STATE_FAMILY = "family"; + private static final StateNamespace NAMESPACE = new StateNamespaceForTest("ns"); + + @Mock private WindmillStateReader mockReader; + @Mock private Supplier readStateSupplier; + + private WindmillStateCache cacheViaMultimap; + private WindmillStateInternals underTestMapViaMultimap; + + private static ByteString key(StateNamespace namespace, String addrId) { + return ByteString.copyFromUtf8(namespace.stringKey() + "+u" + addrId); + } + + private static ByteString encodeWithCoder(T value, Coder coder) { + ByteStringOutputStream out = new ByteStringOutputStream(); + try { + coder.encode(value, out, Context.OUTER); + } catch (IOException e) { + throw new RuntimeException(e); + } + return out.toByteString(); + } + + @Before + public void setUp() { + MockitoAnnotations.initMocks(this); + DataflowWorkerHarnessOptions options = + PipelineOptionsFactory.as(DataflowWorkerHarnessOptions.class); + cacheViaMultimap = + WindmillStateCache.builder() + .setSizeMb(options.getWorkerCacheMb()) + .setSupportMapViaMultimap(true) + .build(); + + WindmillTagEncoding tagEncoding = WindmillTagEncodingV1.instance(); + underTestMapViaMultimap = + new WindmillStateInternals<>( + "key", + STATE_FAMILY, + mockReader, + false, + cacheViaMultimap + .forComputation("comp") + .forKey( + WindmillComputationKey.create( + "comp", ByteString.copyFrom("key", StandardCharsets.UTF_8), 123), + 17L, + 1L) + .forFamily(STATE_FAMILY), + tagEncoding, + readStateSupplier); + + when(mockReader.multimapFetchAllFuture(anyBoolean(), any(), any(), any())) + .thenReturn(Futures.immediateFuture(Collections.emptyList())); + when(mockReader.multimapFetchSingleEntryFuture(any(), any(), any(), any())) + .thenReturn(Futures.immediateFuture(Collections.emptyList())); + } + + /** + * Tests that {@link WindmillMapViaMultimap#computeIfAbsent} returns {@code null} on insertion, + * conforming to the {@link MapState#computeIfAbsent} specification. + */ + @Test + public void testComputeIfAbsentReturnsNullOnInsert() { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTestMapViaMultimap.state(NAMESPACE, addr); + + final String key = "tag1"; + SettableFuture> future = SettableFuture.create(); + when(mockReader.multimapFetchSingleEntryFuture( + encodeWithCoder(key, StringUtf8Coder.of()), + key(NAMESPACE, "map"), + STATE_FAMILY, + VarIntCoder.of())) + .thenReturn(future); + future.set(Collections.emptyList()); + + ReadableState result = mapState.computeIfAbsent(key, k -> 42); + Integer readValue = result.read(); + + assertNull( + "computeIfAbsent returns null when key was absent and computed value is inserted", + readValue); + assertEquals(Integer.valueOf(42), mapState.get(key).read()); + } + + /** + * Tests that {@link WindmillMapViaMultimap} throws {@link IllegalArgumentException} if multiple + * values exist for a key, enforcing at-most-one-value map invariants. + */ + @Test(expected = IllegalArgumentException.class) + public void testGetThrowsOnMultiValuedKey() { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTestMapViaMultimap.state(NAMESPACE, addr); + + final String key = "tag1"; + SettableFuture> future = SettableFuture.create(); + when(mockReader.multimapFetchSingleEntryFuture( + encodeWithCoder(key, StringUtf8Coder.of()), + key(NAMESPACE, "map"), + STATE_FAMILY, + VarIntCoder.of())) + .thenReturn(future); + // Storage returns multiple values for key + future.set(Arrays.asList(10, 20)); + + mapState.get(key).read(); + } + + /** + * Demonstrates Issue 3.3: Every {@link WindmillMapViaMultimap#put} on an unread key executes + * {@code multimap.remove(key)} followed by {@code multimap.put(key, value)}, which marks the key + * as locally removed and forces a redundant {@code DeleteAll = true} update request to Windmill + * on commit. + */ + @Test + public void testPutEmitsDeleteAllAndAddValues() throws Exception { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTestMapViaMultimap.state(NAMESPACE, addr); + + // Call put without prior read + mapState.put("key1", 42); + + Windmill.WorkItemCommitRequest.Builder commitBuilder = + Windmill.WorkItemCommitRequest.newBuilder(); + underTestMapViaMultimap.persist(commitBuilder); + + assertEquals(1, commitBuilder.getMultimapUpdatesCount()); + Windmill.TagMultimapUpdateRequest update = commitBuilder.getMultimapUpdates(0); + assertEquals(1, update.getUpdatesCount()); + Windmill.TagMultimapEntry entry = update.getUpdates(0); + + // DEMONSTRATION OF OVERHEAD: entry has DeleteAll = true AND AddValues, because put() called + // remove() first! + assertTrue( + "Demonstrates Issue 3.3: put() without prior read forced deleteAll=true update request", + entry.getDeleteAll()); + assertEquals(1, entry.getValuesCount()); + } + + /** Tests basic put, get, overwrite, and remove on {@link WindmillMapViaMultimap}. */ + @Test + public void testBasicPutGetOverwriteAndRemove() { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTestMapViaMultimap.state(NAMESPACE, addr); + + mapState.put("a", 1); + assertEquals(Integer.valueOf(1), mapState.get("a").read()); + + mapState.put("a", 2); + assertEquals(Integer.valueOf(2), mapState.get("a").read()); + + mapState.remove("a"); + assertNull(mapState.get("a").read()); + } + + /** Tests keys, values, and entries on {@link WindmillMapViaMultimap}. */ + @Test + public void testKeysValuesAndEntries() { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTestMapViaMultimap.state(NAMESPACE, addr); + + mapState.put("k1", 10); + mapState.put("k2", 20); + + assertThat(mapState.keys().read(), Matchers.containsInAnyOrder("k1", "k2")); + assertThat(mapState.values().read(), Matchers.containsInAnyOrder(10, 20)); + assertThat( + mapState.entries().read(), + Matchers.containsInAnyOrder( + new AbstractMap.SimpleEntry<>("k1", 10), new AbstractMap.SimpleEntry<>("k2", 20))); + } + + /** Tests clear() and isEmpty() on {@link WindmillMapViaMultimap}. */ + @Test + public void testClearAndIsEmpty() { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTestMapViaMultimap.state(NAMESPACE, addr); + + mapState.put("key1", 10); + assertFalse(mapState.isEmpty().read()); + + mapState.clear(); + assertTrue(mapState.isEmpty().read()); + } + + /** Tests that computeIfAbsent on an existing key returns the existing value. */ + @Test + public void testComputeIfAbsentExistingKeyReturnsExistingValue() { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTestMapViaMultimap.state(NAMESPACE, addr); + + mapState.put("key1", 10); + AtomicBoolean called = new AtomicBoolean(false); + ReadableState result = + mapState.computeIfAbsent( + "key1", + k -> { + called.set(true); + return 99; + }); + + assertEquals(Integer.valueOf(10), result.read()); + assertFalse("Mapping function should not be called for existing key", called.get()); + } +} diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillSetTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillSetTest.java new file mode 100644 index 000000000000..a932a9958955 --- /dev/null +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillSetTest.java @@ -0,0 +1,319 @@ +/* + * 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.dataflow.worker.windmill.state; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.when; + +import java.io.Closeable; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.Collections; +import javax.annotation.Nullable; +import org.apache.beam.runners.core.StateNamespace; +import org.apache.beam.runners.core.StateNamespaceForTest; +import org.apache.beam.runners.core.StateTag; +import org.apache.beam.runners.core.StateTags; +import org.apache.beam.runners.dataflow.options.DataflowWorkerHarnessOptions; +import org.apache.beam.runners.dataflow.worker.WindmillComputationKey; +import org.apache.beam.sdk.coders.BooleanCoder; +import org.apache.beam.sdk.coders.ByteArrayCoder; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.Coder.Context; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.state.ReadableState; +import org.apache.beam.sdk.state.SetState; +import org.apache.beam.sdk.util.ByteStringOutputStream; +import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Supplier; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Futures; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.SettableFuture; +import org.hamcrest.Matchers; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.Timeout; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +/** Unit tests demonstrating correctness issues and contract violations in {@link WindmillSet}. */ +@RunWith(JUnit4.class) +public class WindmillSetTest { + @Rule public transient Timeout globalTimeout = Timeout.seconds(600); + + private static final String STATE_FAMILY = "family"; + private static final StateNamespace NAMESPACE = new StateNamespaceForTest("ns"); + + @Mock private WindmillStateReader mockReader; + @Mock private Supplier readStateSupplier; + + private WindmillStateCache cache; + private WindmillStateCache cacheViaMultimap; + private WindmillStateInternals underTest; + private WindmillStateInternals underTestMapViaMultimap; + + private static ByteString key(StateNamespace namespace, String addrId) { + return ByteString.copyFromUtf8(namespace.stringKey() + "+u" + addrId); + } + + private static ByteString encodeWithCoder(T value, Coder coder) { + ByteStringOutputStream out = new ByteStringOutputStream(); + try { + coder.encode(value, out, Context.OUTER); + } catch (IOException e) { + throw new RuntimeException(e); + } + return out.toByteString(); + } + + private ByteString protoKey(String tagId, @Nullable K tag, Coder keyCoder) + throws IOException { + ByteStringOutputStream keyStream = new ByteStringOutputStream(); + key(NAMESPACE, tagId).writeTo(keyStream); + if (tag != null) { + keyCoder.encode(tag, keyStream, Context.OUTER); + } + return keyStream.toByteString(); + } + + private static byte[] dup(byte[] key) { + byte[] res = new byte[key.length]; + System.arraycopy(key, 0, res, 0, key.length); + return res; + } + + @Before + public void setUp() { + MockitoAnnotations.initMocks(this); + DataflowWorkerHarnessOptions options = + PipelineOptionsFactory.as(DataflowWorkerHarnessOptions.class); + cache = WindmillStateCache.builder().setSizeMb(options.getWorkerCacheMb()).build(); + cacheViaMultimap = + WindmillStateCache.builder() + .setSizeMb(options.getWorkerCacheMb()) + .setSupportMapViaMultimap(true) + .build(); + + WindmillTagEncoding tagEncoding = WindmillTagEncodingV1.instance(); + underTest = + new WindmillStateInternals<>( + "key", + STATE_FAMILY, + mockReader, + false, + cache + .forComputation("comp") + .forKey( + WindmillComputationKey.create( + "comp", ByteString.copyFrom("key", StandardCharsets.UTF_8), 123), + 17L, + 1L) + .forFamily(STATE_FAMILY), + tagEncoding, + readStateSupplier); + + underTestMapViaMultimap = + new WindmillStateInternals<>( + "key", + STATE_FAMILY, + mockReader, + false, + cacheViaMultimap + .forComputation("comp") + .forKey( + WindmillComputationKey.create( + "comp", ByteString.copyFrom("key", StandardCharsets.UTF_8), 123), + 17L, + 1L) + .forFamily(STATE_FAMILY), + tagEncoding, + readStateSupplier); + + when(mockReader.valuePrefixFuture(any(), any(), any())) + .thenReturn(Futures.immediateFuture(Collections.emptyList())); + } + + /** + * Tests that {@link WindmillSet#addIfAbsent} returns {@code true} when the element was absent + * (and newly added), and returns {@code false} when the element was already present. + */ + @Test + public void testAddIfAbsentReturnWithWindmillMap() throws Exception { + StateTag> addr = StateTags.set("set", StringUtf8Coder.of()); + SetState setState = underTest.state(NAMESPACE, addr); + + final String elem = "elem1"; + SettableFuture future = SettableFuture.create(); + when(mockReader.valueFuture( + protoKey("set", elem, StringUtf8Coder.of()), STATE_FAMILY, BooleanCoder.of())) + .thenReturn(future); + future.set(null); // null indicates not in Windmill storage + + // 1st add: element is absent from storage. It should be added and return true. + ReadableState add1 = setState.addIfAbsent(elem); + Boolean result1 = add1.read(); + assertTrue("addIfAbsent should return true when element was absent", result1); + + // 2nd add: element is now already present in the set. It should return false. + ReadableState add2 = setState.addIfAbsent(elem); + Boolean result2 = add2.read(); + assertFalse("addIfAbsent should return false when element was already present", result2); + } + + /** + * Tests that when backed by {@link WindmillMapViaMultimap}, {@link WindmillSet#addIfAbsent} + * returns {@code true} on first add and {@code false} when already present. + */ + @Test + public void testAddIfAbsentWithMapViaMultimap() throws Exception { + StateTag> addr = StateTags.set("set", StringUtf8Coder.of()); + SetState setState = underTestMapViaMultimap.state(NAMESPACE, addr); + + final String elem = "elem1"; + SettableFuture> future = SettableFuture.create(); + when(mockReader.multimapFetchSingleEntryFuture( + encodeWithCoder(elem, StringUtf8Coder.of()), + key(NAMESPACE, "set"), + STATE_FAMILY, + BooleanCoder.of())) + .thenReturn(future); + future.set(Collections.emptyList()); + + // 1st add: returns true + assertTrue(setState.addIfAbsent(elem).read()); + + // 2nd add: element is already present, returns false + assertFalse(setState.addIfAbsent(elem).read()); + } + + /** + * Tests that {@link WindmillSet#contains} returns {@code false} (not {@code null}) on repeated + * lookups. + */ + @Test + public void testContainsReturnsFalseOnRepeatedRead() throws Exception { + StateTag> addr = StateTags.set("set", StringUtf8Coder.of()); + SetState setState = underTest.state(NAMESPACE, addr); + + final String elem = "elem1"; + SettableFuture future = SettableFuture.create(); + when(mockReader.valueFuture( + protoKey("set", elem, StringUtf8Coder.of()), STATE_FAMILY, BooleanCoder.of())) + .thenReturn(future); + future.set(null); + + // 1st call returns Boolean.FALSE (and adds to negativeCache) + assertEquals(Boolean.FALSE, setState.contains(elem).read()); + + // 2nd call: key is in negativeCache, must return Boolean.FALSE + assertEquals(Boolean.FALSE, setState.contains(elem).read()); + } + + /** Tests that {@link WindmillSet#contains} returns {@code false} after an element is removed. */ + @Test + public void testContainsReturnsFalseAfterRemove() { + StateTag> addr = StateTags.set("set", StringUtf8Coder.of()); + SetState setState = underTest.state(NAMESPACE, addr); + + setState.add("elem1"); + setState.remove("elem1"); + + assertEquals(Boolean.FALSE, setState.contains("elem1").read()); + } + + /** Tests that {@link WindmillSet#contains} returns {@code false} after {@link SetState#clear}. */ + @Test + public void testContainsReturnsFalseAfterClear() { + StateTag> addr = StateTags.set("set", StringUtf8Coder.of()); + SetState setState = underTest.state(NAMESPACE, addr); + + setState.add("elem1"); + setState.clear(); + + assertEquals(Boolean.FALSE, setState.contains("elem1").read()); + } + + /** + * Tests that {@link WindmillSet} correctly matches {@code byte[]} elements using {@link + * Coder#structuralValue}. + */ + @Test + public void testByteArrayKeyEquality() throws Exception { + StateTag> addr = StateTags.set("set", ByteArrayCoder.of()); + SetState setState = underTest.state(NAMESPACE, addr); + + byte[] elem = new byte[] {1, 2, 3}; + setState.add(elem); + + byte[] lookupElem = dup(elem); + assertTrue( + "WindmillSet matches byte[] elements via structuralValue", + setState.contains(lookupElem).read()); + } + + /** Tests basic add and contains on {@link WindmillSet}. */ + @Test + public void testBasicAddAndContains() throws Exception { + StateTag> addr = StateTags.set("set", StringUtf8Coder.of()); + SetState setState = underTest.state(NAMESPACE, addr); + + when(mockReader.valueFuture( + protoKey("set", "elem1", StringUtf8Coder.of()), STATE_FAMILY, BooleanCoder.of())) + .thenReturn(Futures.immediateFuture(null)); + + assertFalse(setState.contains("elem1").read()); + setState.add("elem1"); + assertTrue(setState.contains("elem1").read()); + } + + /** Tests read() iteration on {@link WindmillSet}. */ + @Test + public void testReadIteration() throws Exception { + StateTag> addr = StateTags.set("set", StringUtf8Coder.of()); + SetState setState = underTest.state(NAMESPACE, addr); + + when(mockReader.valuePrefixFuture( + protoKey("set", null, StringUtf8Coder.of()), STATE_FAMILY, BooleanCoder.of())) + .thenReturn(Futures.immediateFuture(Collections.emptyList())); + + setState.add("a"); + setState.add("b"); + + assertThat(setState.read(), Matchers.containsInAnyOrder("a", "b")); + } + + /** Tests isEmpty on {@link WindmillSet}. */ + @Test + public void testIsEmpty() { + StateTag> addr = StateTags.set("set", StringUtf8Coder.of()); + SetState setState = underTest.state(NAMESPACE, addr); + + setState.add("elem1"); + assertFalse(setState.isEmpty().read()); + + setState.clear(); + assertTrue(setState.isEmpty().read()); + } +}