[TEX] Part 3: TrackingKeyValueService: tracking iterator utils

atlasdb-impl-shared/build.gradle

@@ -117,6 +117,7 @@ dependencies {
     exclude group: 'org.hamcrest'
     exclude group: 'org.ow2.asm'
   }
+  testImplementation 'one.util:streamex'
 
   testRuntimeOnly 'ch.qos.logback:logback-classic'
 }

atlasdb-impl-shared/src/main/java/com/palantir/atlasdb/transaction/impl/expectations/TrackingClosableIterator.java

@@ -0,0 +1,34 @@
+/*
+ * (c) Copyright 2022 Palantir Technologies Inc. All rights reserved.
+ *
+ * Licensed 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 com.palantir.atlasdb.transaction.impl.expectations;
+
+import com.palantir.common.base.ClosableIterator;
+import java.util.function.Consumer;
+import java.util.function.ToLongFunction;
+
+public final class TrackingClosableIterator<T> extends TrackingIterator<T, ClosableIterator<T>>
+        implements ClosableIterator<T> {
+
+    public TrackingClosableIterator(ClosableIterator<T> delegate, Consumer<Long> tracker, ToLongFunction<T> measurer) {
+        super(delegate, tracker, measurer);
+    }
+
+    @Override
+    public void close() {
+        delegate().close();
+    }
+}

atlasdb-impl-shared/src/main/java/com/palantir/atlasdb/transaction/impl/expectations/TrackingIterator.java

@@ -0,0 +1,46 @@
+/*
+ * (c) Copyright 2022 Palantir Technologies Inc. All rights reserved.
+ *
+ * Licensed 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 com.palantir.atlasdb.transaction.impl.expectations;
+
+import com.google.common.collect.ForwardingIterator;
+import java.util.Iterator;
+import java.util.function.Consumer;
+import java.util.function.ToLongFunction;
+
+public class TrackingIterator<T, I extends Iterator<T>> extends ForwardingIterator<T> {
+    private final I delegate;
+    private final ToLongFunction<T> measurer;
+    private final Consumer<Long> tracker;
+
+    public TrackingIterator(I delegate, Consumer<Long> tracker, ToLongFunction<T> measurer) {
+        this.delegate = delegate;
+        this.tracker = tracker;
+        this.measurer = measurer;
+    }
+
+    @Override
+    protected I delegate() {
+        return delegate;
+    }
+
+    @Override
+    public T next() {
+        T result = delegate.next();
+        tracker.accept(measurer.applyAsLong(result));
+        return result;
+    }
+}

atlasdb-impl-shared/src/main/java/com/palantir/atlasdb/transaction/impl/expectations/TrackingRowColumnRangeIterator.java

@@ -0,0 +1,33 @@
+/*
+ * (c) Copyright 2022 Palantir Technologies Inc. All rights reserved.
+ *
+ * Licensed 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 com.palantir.atlasdb.transaction.impl.expectations;
+
+import com.palantir.atlasdb.keyvalue.api.Cell;
+import com.palantir.atlasdb.keyvalue.api.RowColumnRangeIterator;
+import com.palantir.atlasdb.keyvalue.api.Value;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.function.Consumer;
+import java.util.function.ToLongFunction;
+
+public final class TrackingRowColumnRangeIterator
+        extends TrackingIterator<Map.Entry<Cell, Value>, RowColumnRangeIterator> implements RowColumnRangeIterator {
+    public TrackingRowColumnRangeIterator(
+            RowColumnRangeIterator delegate, Consumer<Long> tracker, ToLongFunction<Entry<Cell, Value>> measurer) {
+        super(delegate, tracker, measurer);
+    }
+}

atlasdb-impl-shared/src/test/java/com/palantir/atlasdb/transaction/impl/expectations/TrackingClosableIteratorTest.java

@@ -0,0 +1,101 @@
+/*
+ * (c) Copyright 2022 Palantir Technologies Inc. All rights reserved.
+ *
+ * Licensed 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 com.palantir.atlasdb.transaction.impl.expectations;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+
+import com.google.common.collect.ImmutableList;
+import com.palantir.common.base.ClosableIterator;
+import com.palantir.common.base.ClosableIterators;
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+import java.util.function.Consumer;
+import java.util.function.ToLongFunction;
+import one.util.streamex.StreamEx;
+import org.junit.Test;
+
+public final class TrackingClosableIteratorTest {
+    private static final String STRING = "test";
+    private static final ClosableIterator<String> ONE_ELEMENT_ITERATOR =
+            ClosableIterators.wrapWithEmptyClose(List.of(STRING).iterator());
+    private static final ImmutableList<String> STRINGS = ImmutableList.of(
+            "test4", "test4", "test200", "composite", "", "t", "twentyElementString1", "tt", "twentyElementString2");
+
+    // these have to be anonymous inner classes rather than lambdas in order to spy
+    private static final Consumer<Long> NO_OP = new Consumer<>() {
+        @Override
+        public void accept(Long _unused) {}
+    };
+    private static final ToLongFunction<String> STRING_LENGTH_MEASURER = new ToLongFunction<>() {
+        @Override
+        public long applyAsLong(String value) {
+            return value.length();
+        }
+    };
+
+    @Test
+    public void oneElementTrackingClosableIteratorIsWiredCorrectly() {
+        Consumer<Long> tracker = spy(NO_OP);
+        ToLongFunction<String> measurer = spy(STRING_LENGTH_MEASURER);
+        TrackingClosableIterator<String> trackingIterator =
+                new TrackingClosableIterator<>(ONE_ELEMENT_ITERATOR, tracker, measurer);
+
+        assertThat(trackingIterator).toIterable().containsExactlyElementsOf(List.of(STRING));
+        verify(measurer).applyAsLong(STRING);
+        verify(tracker).accept(STRING_LENGTH_MEASURER.applyAsLong(STRING));
+        verifyNoMoreInteractions(tracker, measurer);
+    }
+
+    @Test
+    public void multiElementTrackingClosableIteratorIsWiredCorrectly() {
+        ArrayList<Long> consumed = new ArrayList<>();
+        TrackingClosableIterator<String> trackingIterator =
+                new TrackingClosableIterator<>(createClosableStringIterator(), consumed::add, STRING_LENGTH_MEASURER);
+
+        assertThat(trackingIterator)
+                .toIterable()
+                .containsExactlyElementsOf(ImmutableList.copyOf(createClosableStringIterator()));
+
+        assertThat(consumed)
+                .containsExactlyElementsOf(StreamEx.of(createClosableStringIterator())
+                        .mapToLong(STRING_LENGTH_MEASURER)
+                        .boxed()
+                        .toList());
+    }
+
+    @Test
+    public void trackingClosableStringIteratorDelegatesClose() {
+        ClosableIterator<String> iterator = spy(createClosableStringIterator());
+        TrackingClosableIterator<String> trackingIterator =
+                new TrackingClosableIterator<>(iterator, NO_OP, STRING_LENGTH_MEASURER);
+        trackingIterator.close();
+        verify(iterator).close();
+        verifyNoMoreInteractions(iterator);
+    }
+
+    private static ClosableIterator<String> createClosableStringIterator() {
+        return ClosableIterators.wrapWithEmptyClose(createStringIterator());
+    }
+
+    private static Iterator<String> createStringIterator() {
+        return STRINGS.stream().iterator();
+    }
+}

atlasdb-impl-shared/src/test/java/com/palantir/atlasdb/transaction/impl/expectations/TrackingIteratorTest.java

@@ -0,0 +1,84 @@
+/*
+ * (c) Copyright 2022 Palantir Technologies Inc. All rights reserved.
+ *
+ * Licensed 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 com.palantir.atlasdb.transaction.impl.expectations;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+
+import com.google.common.collect.ImmutableList;
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+import java.util.function.Consumer;
+import java.util.function.ToLongFunction;
+import one.util.streamex.StreamEx;
+import org.junit.Test;
+
+public final class TrackingIteratorTest {
+    private static final String STRING = "test";
+    private static final Iterator<String> ONE_ELEMENT_ITERATOR = List.of(STRING).iterator();
+    private static final ImmutableList<String> STRINGS = ImmutableList.of(
+            "test4", "test4", "test200", "composite", "", "t", "twentyElementString1", "tt", "twentyElementString2");
+
+    // these have to be anonymous inner classes rather than lambdas in order to spy
+    private static final Consumer<Long> NO_OP = new Consumer<>() {
+        @Override
+        public void accept(Long _unused) {}
+    };
+    private static final ToLongFunction<String> STRING_LENGTH_MEASURER = new ToLongFunction<>() {
+        @Override
+        public long applyAsLong(String value) {
+            return value.length();
+        }
+    };
+
+    @Test
+    public void oneElementTrackingIteratorIsWiredCorrectly() {
+        Consumer<Long> tracker = spy(NO_OP);
+        ToLongFunction<String> measurer = spy(STRING_LENGTH_MEASURER);
+        TrackingIterator<String, Iterator<String>> trackingIterator =
+                new TrackingIterator<>(ONE_ELEMENT_ITERATOR, tracker, measurer);
+
+        assertThat(trackingIterator).toIterable().containsExactlyElementsOf(List.of(STRING));
+        verify(measurer).applyAsLong(STRING);
+        verify(tracker).accept(STRING_LENGTH_MEASURER.applyAsLong(STRING));
+        verifyNoMoreInteractions(tracker, measurer);
+    }
+
+    @Test
+    public void multiElementTrackingIteratorIsWiredCorrectly() {
+        ArrayList<Long> consumed = new ArrayList<>();
+        TrackingIterator<String, Iterator<String>> trackingIterator =
+                new TrackingIterator<>(createStringIterator(), consumed::add, STRING_LENGTH_MEASURER);
+
+        assertThat(trackingIterator)
+                .toIterable()
+                .containsExactlyElementsOf(ImmutableList.copyOf(createStringIterator()));
+
+        assertThat(consumed)
+                .containsExactlyElementsOf(StreamEx.of(createStringIterator())
+                        .mapToLong(STRING_LENGTH_MEASURER)
+                        .boxed()
+                        .toList());
+    }
+
+    public static Iterator<String> createStringIterator() {
+        return STRINGS.stream().iterator();
+    }
+}