runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/GroupingShuffleEntryIteratorTest.java - beam - Git at Google

 /*
  * 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.util.common.worker;

 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.when;

 import java.io.Closeable;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import javax.annotation.Nullable;
 import org.apache.beam.runners.core.metrics.ExecutionStateSampler;
 import org.apache.beam.runners.core.metrics.ExecutionStateTracker;
 import org.apache.beam.runners.dataflow.options.DataflowPipelineDebugOptions;
 import org.apache.beam.runners.dataflow.worker.BatchModeExecutionContext;
 import org.apache.beam.runners.dataflow.worker.DataflowOperationContext.DataflowExecutionState;
 import org.apache.beam.runners.dataflow.worker.ExperimentContext.Experiment;
 import org.apache.beam.runners.dataflow.worker.TestOperationContext.TestDataflowExecutionState;
 import org.apache.beam.runners.dataflow.worker.counters.Counter;
 import org.apache.beam.runners.dataflow.worker.counters.NameContext;
 import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.util.common.Reiterator;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Charsets;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;

 /** Tests for {@link GroupingShuffleEntryIterator}. */
 @RunWith(JUnit4.class)
 public class GroupingShuffleEntryIteratorTest {
   private static final ByteArrayShufflePosition START_POSITION =
       ByteArrayShufflePosition.of("aaa".getBytes(StandardCharsets.UTF_8));
   private static final ByteArrayShufflePosition END_POSITION =
       ByteArrayShufflePosition.of("zzz".getBytes(StandardCharsets.UTF_8));

   private static final String MOCK_STAGE_NAME = "mockStageName";
   private static final String MOCK_ORIGINAL_NAME_FOR_EXECUTING_STEP1 = "mockOriginalName1";
   private static final String MOCK_SYSTEM_NAME = "mockSystemName";
   private static final String MOCK_USER_NAME = "mockUserName";
   private static final String ORIGINAL_SHUFFLE_STEP_NAME = "originalName";

   @Mock private ShuffleEntryReader reader;
   private GroupingShuffleEntryIterator iterator;

   private final ExecutionStateSampler sampler = ExecutionStateSampler.newForTest();
   private final ExecutionStateTracker tracker = new ExecutionStateTracker(sampler);
   private Closeable trackerCleanup;

   @Before
   public void setUp() {
     trackerCleanup = tracker.activate();
   }

   @After
   public void tearDown() throws IOException {
     trackerCleanup.close();
   }

   private static class ListReiterator<T> implements Reiterator<T> {
     protected final List<T> entries;
     protected int nextIndex;

     public ListReiterator(List<T> entries, int nextIndex) {
       this.entries = entries;
       this.nextIndex = nextIndex;
     }

     @Override
     public Reiterator<T> copy() {
       return new ListReiterator<T>(entries, nextIndex);
     }

     @Override
     public boolean hasNext() {
       return nextIndex < entries.size();
     }

     @Override
     public T next() {
       T res = entries.get(nextIndex);
       nextIndex++;
       return res;
     }

     @Override
     public void remove() {
       throw new UnsupportedOperationException();
     }
   }

   private void setCurrentExecutionState(String mockOriginalName) {
     DataflowExecutionState state =
         new TestDataflowExecutionState(
             NameContext.create(MOCK_STAGE_NAME, mockOriginalName, MOCK_SYSTEM_NAME, MOCK_USER_NAME),
             "activity");
     tracker.enterState(state);
   }

   private static ShuffleEntry shuffleEntry(String key, String value) {
     return new ShuffleEntry(
         /* use key itself as position */
         ByteArrayShufflePosition.of(key.getBytes(Charsets.UTF_8)),
         key.getBytes(Charsets.UTF_8),
         new byte[0],
         value.getBytes(Charsets.UTF_8));
   }

   @Test
   public void testCopyValuesIterator() throws Exception {
     setCurrentExecutionState(MOCK_ORIGINAL_NAME_FOR_EXECUTING_STEP1);

     MockitoAnnotations.initMocks(this);
     PipelineOptions options = PipelineOptionsFactory.create();
     options
         .as(DataflowPipelineDebugOptions.class)
         .setExperiments(Lists.newArrayList(Experiment.IntertransformIO.getName()));
     BatchModeExecutionContext spyExecutionContext =
         Mockito.spy(BatchModeExecutionContext.forTesting(options, "STAGE"));

     ArrayList<ShuffleEntry> entries = new ArrayList<>();
     entries.add(shuffleEntry("k1", "v11"));
     entries.add(shuffleEntry("k1", "v12"));
     entries.add(shuffleEntry("k1", "v13"));
     when(reader.read(START_POSITION, END_POSITION)).thenReturn(new ListReiterator<>(entries, 0));

     final ShuffleReadCounter shuffleReadCounter =
         new ShuffleReadCounter(ORIGINAL_SHUFFLE_STEP_NAME, true, null);
     iterator =
         new GroupingShuffleEntryIterator(reader, START_POSITION, END_POSITION) {
           @Override
           protected void notifyElementRead(long byteSize) {
             // nothing
           }

           @Override
           protected void commitBytesRead(long bytes) {
             shuffleReadCounter.addBytesRead(bytes);
           }
         };
     assertTrue(iterator.advance());
     KeyGroupedShuffleEntries k1Entries = iterator.getCurrent();
     Reiterator<ShuffleEntry> values1 = k1Entries.values.iterator();
     assertTrue(values1.hasNext());
     Reiterator<ShuffleEntry> values1Copy1 = values1.copy();
     Reiterator<ShuffleEntry> values1Copy2 = values1.copy();
     ShuffleEntry expectedEntry = values1.next();
     assertFalse(iterator.advance()); // Advance the iterator again to record bytes read.

     // Test that the copy works two ways: 1) if we call hasNext, and 2) if we don't.
     assertTrue(values1Copy1.hasNext());
     assertEquals(expectedEntry, values1Copy1.next());

     assertEquals(expectedEntry, values1Copy2.next());

     Map<String, Long> expectedReadBytesMap = new HashMap<>();
     expectedReadBytesMap.put(MOCK_ORIGINAL_NAME_FOR_EXECUTING_STEP1, 15L);

     // Verify that each executing step used when reading from the GroupingShuffleReader
     // has a counter with a bytes read value.
     assertEquals(expectedReadBytesMap.size(), (long) shuffleReadCounter.counterSet.size());
     Iterator it = expectedReadBytesMap.entrySet().iterator();
     while (it.hasNext()) {
       Map.Entry<String, Long> pair = (Map.Entry) it.next();
       Counter counter =
           shuffleReadCounter.counterSet.getExistingCounter(
               ShuffleReadCounter.generateCounterName(ORIGINAL_SHUFFLE_STEP_NAME, pair.getKey()));
       assertEquals(pair.getValue(), counter.getAggregate());
     }
   }

   /** A ShuffleEntryReader that asserts that its iterators never go backwards ("reiterate"). */
   private static class ForwardOnlyShuffleEntryReader implements ShuffleEntryReader {
     private final List<ShuffleEntry> entries;
     private int minNextIndex = 0; // The smallest index that iterators are allowed to advance.

     public ForwardOnlyShuffleEntryReader(List<ShuffleEntry> entries) {
       this.entries = entries;
     }

     @Override
     public Reiterator<ShuffleEntry> read(
         @Nullable ShufflePosition startPosition, @Nullable ShufflePosition endPosition) {
       return new MinIndexAssertingListReiterator(entries, 0);
     }

     @Override
     public void close() {}

     private class MinIndexAssertingListReiterator<T> extends ListReiterator<T> {
       public MinIndexAssertingListReiterator(List<T> entries, int nextIndex) {
         super(entries, nextIndex);
       }

       @Override
       public Reiterator<T> copy() {
         return new MinIndexAssertingListReiterator<T>(entries, nextIndex);
       }

       @Override
       public T next() {
         assertTrue("Reiteration unexpected.", nextIndex >= minNextIndex);
         minNextIndex = nextIndex;
         return super.next();
       }
     }
   }

   /**
    * Tests that GroupingShuffleEntryIterator does not reiterate the underlying shuffle iterator when
    * the returned value iterators are iterated over (i.e., that fast-forwarding works properly).
    */
   @Test
   public void testNoReiteration() throws Exception {
     ArrayList<ShuffleEntry> entries = new ArrayList<>();
     entries.add(shuffleEntry("k1", "v11"));
     entries.add(shuffleEntry("k1", "v12"));
     entries.add(shuffleEntry("k1", "v13"));
     entries.add(shuffleEntry("k2", "v21"));
     entries.add(shuffleEntry("k2", "v22"));
     entries.add(shuffleEntry("k2", "v23"));
     ForwardOnlyShuffleEntryReader reader = new ForwardOnlyShuffleEntryReader(entries);

     iterator =
         new GroupingShuffleEntryIterator(reader, START_POSITION, END_POSITION) {
           @Override
           protected void notifyElementRead(long byteSize) {
             // nothing
           }

           @Override
           protected void commitBytesRead(long bytes) {
             // nothing
           }
         };
     int totalKeys = 0;
     int totalValues = 0;
     while (iterator.advance()) {
       ++totalKeys;
       Reiterator<ShuffleEntry> values = iterator.getCurrent().values.iterator();
       while (values.hasNext()) {
         values.next();
         ++totalValues;
       }
     }
     assertEquals(2, totalKeys);
     assertEquals(6, totalValues);

     // We expect that AssertionException in MinIndexAssertingListReiterator.next() is not thrown.
   }
 }
	/*
	* 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.util.common.worker;

	import static org.junit.Assert.assertEquals;
	import static org.junit.Assert.assertFalse;
	import static org.junit.Assert.assertTrue;
	import static org.mockito.Mockito.when;

	import java.io.Closeable;
	import java.io.IOException;
	import java.nio.charset.StandardCharsets;
	import java.util.ArrayList;
	import java.util.HashMap;
	import java.util.Iterator;
	import java.util.List;
	import java.util.Map;
	import javax.annotation.Nullable;
	import org.apache.beam.runners.core.metrics.ExecutionStateSampler;
	import org.apache.beam.runners.core.metrics.ExecutionStateTracker;
	import org.apache.beam.runners.dataflow.options.DataflowPipelineDebugOptions;
	import org.apache.beam.runners.dataflow.worker.BatchModeExecutionContext;
	import org.apache.beam.runners.dataflow.worker.DataflowOperationContext.DataflowExecutionState;
	import org.apache.beam.runners.dataflow.worker.ExperimentContext.Experiment;
	import org.apache.beam.runners.dataflow.worker.TestOperationContext.TestDataflowExecutionState;
	import org.apache.beam.runners.dataflow.worker.counters.Counter;
	import org.apache.beam.runners.dataflow.worker.counters.NameContext;
	import org.apache.beam.sdk.options.PipelineOptions;
	import org.apache.beam.sdk.options.PipelineOptionsFactory;
	import org.apache.beam.sdk.util.common.Reiterator;
	import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Charsets;
	import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
	import org.junit.After;
	import org.junit.Before;
	import org.junit.Test;
	import org.junit.runner.RunWith;
	import org.junit.runners.JUnit4;
	import org.mockito.Mock;
	import org.mockito.Mockito;
	import org.mockito.MockitoAnnotations;

	/** Tests for {@link GroupingShuffleEntryIterator}. */
	@RunWith(JUnit4.class)
	public class GroupingShuffleEntryIteratorTest {
	private static final ByteArrayShufflePosition START_POSITION =
	ByteArrayShufflePosition.of("aaa".getBytes(StandardCharsets.UTF_8));
	private static final ByteArrayShufflePosition END_POSITION =
	ByteArrayShufflePosition.of("zzz".getBytes(StandardCharsets.UTF_8));

	private static final String MOCK_STAGE_NAME = "mockStageName";
	private static final String MOCK_ORIGINAL_NAME_FOR_EXECUTING_STEP1 = "mockOriginalName1";
	private static final String MOCK_SYSTEM_NAME = "mockSystemName";
	private static final String MOCK_USER_NAME = "mockUserName";
	private static final String ORIGINAL_SHUFFLE_STEP_NAME = "originalName";

	@Mock private ShuffleEntryReader reader;
	private GroupingShuffleEntryIterator iterator;

	private final ExecutionStateSampler sampler = ExecutionStateSampler.newForTest();
	private final ExecutionStateTracker tracker = new ExecutionStateTracker(sampler);
	private Closeable trackerCleanup;

	@Before
	public void setUp() {
	trackerCleanup = tracker.activate();
	}

	@After
	public void tearDown() throws IOException {
	trackerCleanup.close();
	}

	private static class ListReiterator<T> implements Reiterator<T> {
	protected final List<T> entries;
	protected int nextIndex;

	public ListReiterator(List<T> entries, int nextIndex) {
	this.entries = entries;
	this.nextIndex = nextIndex;
	}

	@Override
	public Reiterator<T> copy() {
	return new ListReiterator<T>(entries, nextIndex);
	}

	@Override
	public boolean hasNext() {
	return nextIndex < entries.size();
	}

	@Override
	public T next() {
	T res = entries.get(nextIndex);
	nextIndex++;
	return res;
	}

	@Override
	public void remove() {
	throw new UnsupportedOperationException();
	}
	}

	private void setCurrentExecutionState(String mockOriginalName) {
	DataflowExecutionState state =
	new TestDataflowExecutionState(
	NameContext.create(MOCK_STAGE_NAME, mockOriginalName, MOCK_SYSTEM_NAME, MOCK_USER_NAME),
	"activity");
	tracker.enterState(state);
	}

	private static ShuffleEntry shuffleEntry(String key, String value) {
	return new ShuffleEntry(
	/* use key itself as position */
	ByteArrayShufflePosition.of(key.getBytes(Charsets.UTF_8)),
	key.getBytes(Charsets.UTF_8),
	new byte[0],
	value.getBytes(Charsets.UTF_8));
	}

	@Test
	public void testCopyValuesIterator() throws Exception {
	setCurrentExecutionState(MOCK_ORIGINAL_NAME_FOR_EXECUTING_STEP1);

	MockitoAnnotations.initMocks(this);
	PipelineOptions options = PipelineOptionsFactory.create();
	options
	.as(DataflowPipelineDebugOptions.class)
	.setExperiments(Lists.newArrayList(Experiment.IntertransformIO.getName()));
	BatchModeExecutionContext spyExecutionContext =
	Mockito.spy(BatchModeExecutionContext.forTesting(options, "STAGE"));

	ArrayList<ShuffleEntry> entries = new ArrayList<>();
	entries.add(shuffleEntry("k1", "v11"));
	entries.add(shuffleEntry("k1", "v12"));
	entries.add(shuffleEntry("k1", "v13"));
	when(reader.read(START_POSITION, END_POSITION)).thenReturn(new ListReiterator<>(entries, 0));

	final ShuffleReadCounter shuffleReadCounter =
	new ShuffleReadCounter(ORIGINAL_SHUFFLE_STEP_NAME, true, null);
	iterator =
	new GroupingShuffleEntryIterator(reader, START_POSITION, END_POSITION) {
	@Override
	protected void notifyElementRead(long byteSize) {
	// nothing
	}

	@Override
	protected void commitBytesRead(long bytes) {
	shuffleReadCounter.addBytesRead(bytes);
	}
	};
	assertTrue(iterator.advance());
	KeyGroupedShuffleEntries k1Entries = iterator.getCurrent();
	Reiterator<ShuffleEntry> values1 = k1Entries.values.iterator();
	assertTrue(values1.hasNext());
	Reiterator<ShuffleEntry> values1Copy1 = values1.copy();
	Reiterator<ShuffleEntry> values1Copy2 = values1.copy();
	ShuffleEntry expectedEntry = values1.next();
	assertFalse(iterator.advance()); // Advance the iterator again to record bytes read.

	// Test that the copy works two ways: 1) if we call hasNext, and 2) if we don't.
	assertTrue(values1Copy1.hasNext());
	assertEquals(expectedEntry, values1Copy1.next());

	assertEquals(expectedEntry, values1Copy2.next());

	Map<String, Long> expectedReadBytesMap = new HashMap<>();
	expectedReadBytesMap.put(MOCK_ORIGINAL_NAME_FOR_EXECUTING_STEP1, 15L);

	// Verify that each executing step used when reading from the GroupingShuffleReader
	// has a counter with a bytes read value.
	assertEquals(expectedReadBytesMap.size(), (long) shuffleReadCounter.counterSet.size());
	Iterator it = expectedReadBytesMap.entrySet().iterator();
	while (it.hasNext()) {
	Map.Entry<String, Long> pair = (Map.Entry) it.next();
	Counter counter =
	shuffleReadCounter.counterSet.getExistingCounter(
	ShuffleReadCounter.generateCounterName(ORIGINAL_SHUFFLE_STEP_NAME, pair.getKey()));
	assertEquals(pair.getValue(), counter.getAggregate());
	}
	}

	/** A ShuffleEntryReader that asserts that its iterators never go backwards ("reiterate"). */
	private static class ForwardOnlyShuffleEntryReader implements ShuffleEntryReader {
	private final List<ShuffleEntry> entries;
	private int minNextIndex = 0; // The smallest index that iterators are allowed to advance.

	public ForwardOnlyShuffleEntryReader(List<ShuffleEntry> entries) {
	this.entries = entries;
	}

	@Override
	public Reiterator<ShuffleEntry> read(
	@Nullable ShufflePosition startPosition, @Nullable ShufflePosition endPosition) {
	return new MinIndexAssertingListReiterator(entries, 0);
	}

	@Override
	public void close() {}

	private class MinIndexAssertingListReiterator<T> extends ListReiterator<T> {
	public MinIndexAssertingListReiterator(List<T> entries, int nextIndex) {
	super(entries, nextIndex);
	}

	@Override
	public Reiterator<T> copy() {
	return new MinIndexAssertingListReiterator<T>(entries, nextIndex);
	}

	@Override
	public T next() {
	assertTrue("Reiteration unexpected.", nextIndex >= minNextIndex);
	minNextIndex = nextIndex;
	return super.next();
	}
	}
	}

	/**
	* Tests that GroupingShuffleEntryIterator does not reiterate the underlying shuffle iterator when
	* the returned value iterators are iterated over (i.e., that fast-forwarding works properly).
	*/
	@Test
	public void testNoReiteration() throws Exception {
	ArrayList<ShuffleEntry> entries = new ArrayList<>();
	entries.add(shuffleEntry("k1", "v11"));
	entries.add(shuffleEntry("k1", "v12"));
	entries.add(shuffleEntry("k1", "v13"));
	entries.add(shuffleEntry("k2", "v21"));
	entries.add(shuffleEntry("k2", "v22"));
	entries.add(shuffleEntry("k2", "v23"));
	ForwardOnlyShuffleEntryReader reader = new ForwardOnlyShuffleEntryReader(entries);

	iterator =
	new GroupingShuffleEntryIterator(reader, START_POSITION, END_POSITION) {
	@Override
	protected void notifyElementRead(long byteSize) {
	// nothing
	}

	@Override
	protected void commitBytesRead(long bytes) {
	// nothing
	}
	};
	int totalKeys = 0;
	int totalValues = 0;
	while (iterator.advance()) {
	++totalKeys;
	Reiterator<ShuffleEntry> values = iterator.getCurrent().values.iterator();
	while (values.hasNext()) {
	values.next();
	++totalValues;
	}
	}
	assertEquals(2, totalKeys);
	assertEquals(6, totalValues);

	// We expect that AssertionException in MinIndexAssertingListReiterator.next() is not thrown.
	}
	}