runners/direct-java/src/test/java/org/apache/beam/runners/direct/EvaluationContextTest.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.direct;

 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.apache.beam.sdk.testing.PCollectionViewTesting.materializeValuesFor;
 import static org.hamcrest.Matchers.contains;
 import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.emptyIterable;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.not;
 import static org.junit.Assert.assertThat;

 import java.util.Collection;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
 import org.apache.beam.runners.core.SideInputReader;
 import org.apache.beam.runners.core.StateNamespaces;
 import org.apache.beam.runners.core.StateTag;
 import org.apache.beam.runners.core.StateTags;
 import org.apache.beam.runners.core.TimerInternals.TimerData;
 import org.apache.beam.runners.direct.DirectExecutionContext.DirectStepContext;
 import org.apache.beam.runners.direct.WatermarkManager.FiredTimers;
 import org.apache.beam.runners.direct.WatermarkManager.TimerUpdate;
 import org.apache.beam.runners.local.StructuralKey;
 import org.apache.beam.sdk.coders.ByteArrayCoder;
 import org.apache.beam.sdk.coders.IterableCoder;
 import org.apache.beam.sdk.coders.KvCoder;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
 import org.apache.beam.sdk.coders.VarIntCoder;
 import org.apache.beam.sdk.coders.VoidCoder;
 import org.apache.beam.sdk.io.GenerateSequence;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.runners.AppliedPTransform;
 import org.apache.beam.sdk.state.BagState;
 import org.apache.beam.sdk.state.TimeDomain;
 import org.apache.beam.sdk.testing.TestPipeline;
 import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.View;
 import org.apache.beam.sdk.transforms.WithKeys;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo.Timing;
 import org.apache.beam.sdk.util.WindowedValue;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.PCollection.IsBounded;
 import org.apache.beam.sdk.values.PCollectionView;
 import org.apache.beam.sdk.values.WindowingStrategy;
 import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableList;
 import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables;
 import org.joda.time.Instant;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;

 /** Tests for {@link EvaluationContext}. */
 @RunWith(JUnit4.class)
 public class EvaluationContextTest {
   private EvaluationContext context;

   private PCollection<Integer> created;
   private PCollection<KV<String, Integer>> downstream;
   private PCollectionView<Iterable<Integer>> view;
   private PCollection<Long> unbounded;

   private DirectGraph graph;

   private AppliedPTransform<?, ?, ?> createdProducer;
   private AppliedPTransform<?, ?, ?> downstreamProducer;
   private AppliedPTransform<?, ?, ?> viewProducer;
   private AppliedPTransform<?, ?, ?> unboundedProducer;

   @Rule public TestPipeline p = TestPipeline.create().enableAbandonedNodeEnforcement(false);

   @Before
   public void setup() {
     DirectRunner runner = DirectRunner.fromOptions(PipelineOptionsFactory.create());

     created = p.apply(Create.of(1, 2, 3));
     downstream = created.apply(WithKeys.of("foo"));
     view = created.apply(View.asIterable());
     unbounded = p.apply(GenerateSequence.from(0));

     p.replaceAll(runner.defaultTransformOverrides());

     KeyedPValueTrackingVisitor keyedPValueTrackingVisitor = KeyedPValueTrackingVisitor.create();
     p.traverseTopologically(keyedPValueTrackingVisitor);

     BundleFactory bundleFactory = ImmutableListBundleFactory.create();
     DirectGraphs.performDirectOverrides(p);
     graph = DirectGraphs.getGraph(p);
     context =
         EvaluationContext.create(
             NanosOffsetClock.create(),
             bundleFactory,
             graph,
             keyedPValueTrackingVisitor.getKeyedPValues(),
             Executors.newSingleThreadExecutor());

     createdProducer = graph.getProducer(created);
     downstreamProducer = graph.getProducer(downstream);
     viewProducer = graph.getProducer(view);
     unboundedProducer = graph.getProducer(unbounded);
   }

   @Test
   public void writeToViewWriterThenReadReads() {
     PCollectionViewWriter<?, Iterable<Integer>> viewWriter =
         context.createPCollectionViewWriter(
             PCollection.createPrimitiveOutputInternal(
                 p,
                 WindowingStrategy.globalDefault(),
                 IsBounded.BOUNDED,
                 IterableCoder.of(KvCoder.of(VoidCoder.of(), VarIntCoder.of()))),
             view);
     BoundedWindow window = new TestBoundedWindow(new Instant(1024L));
     BoundedWindow second = new TestBoundedWindow(new Instant(899999L));
     ImmutableList.Builder<WindowedValue<?>> valuesBuilder = ImmutableList.builder();
     for (Object materializedValue : materializeValuesFor(View.asIterable(), 1)) {
       valuesBuilder.add(
           WindowedValue.of(
               materializedValue, new Instant(1222), window, PaneInfo.ON_TIME_AND_ONLY_FIRING));
     }
     for (Object materializedValue : materializeValuesFor(View.asIterable(), 2)) {
       valuesBuilder.add(
           WindowedValue.of(
               materializedValue,
               new Instant(8766L),
               second,
               PaneInfo.createPane(true, false, Timing.ON_TIME, 0, 0)));
     }
     viewWriter.add((Iterable) valuesBuilder.build());

     SideInputReader reader = context.createSideInputReader(ImmutableList.of(view));
     assertThat(reader.get(view, window), containsInAnyOrder(1));
     assertThat(reader.get(view, second), containsInAnyOrder(2));

     ImmutableList.Builder<WindowedValue<?>> overwrittenValuesBuilder = ImmutableList.builder();
     for (Object materializedValue : materializeValuesFor(View.asIterable(), 4444)) {
       overwrittenValuesBuilder.add(
           WindowedValue.of(
               materializedValue,
               new Instant(8677L),
               second,
               PaneInfo.createPane(false, true, Timing.LATE, 1, 1)));
     }
     viewWriter.add((Iterable) overwrittenValuesBuilder.build());
     assertThat(reader.get(view, second), containsInAnyOrder(2));
     // The cached value is served in the earlier reader
     reader = context.createSideInputReader(ImmutableList.of(view));
     assertThat(reader.get(view, second), containsInAnyOrder(4444));
   }

   @Test
   public void getExecutionContextSameStepSameKeyState() {
     DirectExecutionContext fooContext =
         context.getExecutionContext(createdProducer, StructuralKey.of("foo", StringUtf8Coder.of()));

     StateTag<BagState<Integer>> intBag = StateTags.bag("myBag", VarIntCoder.of());

     DirectStepContext stepContext = fooContext.getStepContext("s1");
     stepContext.stateInternals().state(StateNamespaces.global(), intBag).add(1);

     context.handleResult(
         ImmutableListBundleFactory.create()
             .createKeyedBundle(StructuralKey.of("foo", StringUtf8Coder.of()), created)
             .commit(Instant.now()),
         ImmutableList.of(),
         StepTransformResult.withoutHold(createdProducer)
             .withState(stepContext.commitState())
             .build());

     DirectExecutionContext secondFooContext =
         context.getExecutionContext(createdProducer, StructuralKey.of("foo", StringUtf8Coder.of()));
     assertThat(
         secondFooContext
             .getStepContext("s1")
             .stateInternals()
             .state(StateNamespaces.global(), intBag)
             .read(),
         contains(1));
   }

   @Test
   public void getExecutionContextDifferentKeysIndependentState() {
     DirectExecutionContext fooContext =
         context.getExecutionContext(createdProducer, StructuralKey.of("foo", StringUtf8Coder.of()));

     StateTag<BagState<Integer>> intBag = StateTags.bag("myBag", VarIntCoder.of());

     fooContext.getStepContext("s1").stateInternals().state(StateNamespaces.global(), intBag).add(1);

     DirectExecutionContext barContext =
         context.getExecutionContext(createdProducer, StructuralKey.of("bar", StringUtf8Coder.of()));
     assertThat(barContext, not(equalTo(fooContext)));
     assertThat(
         barContext
             .getStepContext("s1")
             .stateInternals()
             .state(StateNamespaces.global(), intBag)
             .read(),
         emptyIterable());
   }

   @Test
   public void getExecutionContextDifferentStepsIndependentState() {
     StructuralKey<?> myKey = StructuralKey.of("foo", StringUtf8Coder.of());
     DirectExecutionContext fooContext = context.getExecutionContext(createdProducer, myKey);

     StateTag<BagState<Integer>> intBag = StateTags.bag("myBag", VarIntCoder.of());

     fooContext.getStepContext("s1").stateInternals().state(StateNamespaces.global(), intBag).add(1);

     DirectExecutionContext barContext = context.getExecutionContext(downstreamProducer, myKey);
     assertThat(
         barContext
             .getStepContext("s1")
             .stateInternals()
             .state(StateNamespaces.global(), intBag)
             .read(),
         emptyIterable());
   }

   @Test
   public void handleResultStoresState() {
     StructuralKey<?> myKey = StructuralKey.of("foo".getBytes(UTF_8), ByteArrayCoder.of());
     DirectExecutionContext fooContext = context.getExecutionContext(downstreamProducer, myKey);

     StateTag<BagState<Integer>> intBag = StateTags.bag("myBag", VarIntCoder.of());

     CopyOnAccessInMemoryStateInternals<?> state = fooContext.getStepContext("s1").stateInternals();
     BagState<Integer> bag = state.state(StateNamespaces.global(), intBag);
     bag.add(1);
     bag.add(2);
     bag.add(4);

     TransformResult<?> stateResult =
         StepTransformResult.withoutHold(downstreamProducer).withState(state).build();

     context.handleResult(
         context.createKeyedBundle(myKey, created).commit(Instant.now()),
         ImmutableList.of(),
         stateResult);

     DirectExecutionContext afterResultContext =
         context.getExecutionContext(downstreamProducer, myKey);

     CopyOnAccessInMemoryStateInternals<?> afterResultState =
         afterResultContext.getStepContext("s1").stateInternals();
     assertThat(afterResultState.state(StateNamespaces.global(), intBag).read(), contains(1, 2, 4));
   }

   @Test
   public void callAfterOutputMustHaveBeenProducedAfterEndOfWatermarkCallsback() throws Exception {
     final CountDownLatch callLatch = new CountDownLatch(1);
     Runnable callback = callLatch::countDown;

     // Should call back after the end of the global window
     context.scheduleAfterOutputWouldBeProduced(
         downstream, GlobalWindow.INSTANCE, WindowingStrategy.globalDefault(), callback);

     TransformResult<?> result =
         StepTransformResult.withHold(createdProducer, new Instant(0)).build();

     context.handleResult(null, ImmutableList.of(), result);
     // Difficult to demonstrate that we took no action in a multithreaded world; poll for a bit
     // will likely be flaky if this logic is broken
     assertThat(callLatch.await(500L, TimeUnit.MILLISECONDS), is(false));

     TransformResult<?> finishedResult = StepTransformResult.withoutHold(createdProducer).build();
     context.handleResult(null, ImmutableList.of(), finishedResult);
     context.forceRefresh();
     // Obtain the value via blocking call
     assertThat(callLatch.await(1, TimeUnit.SECONDS), is(true));
   }

   @Test
   public void callAfterOutputMustHaveBeenProducedAlreadyAfterCallsImmediately() throws Exception {
     TransformResult<?> finishedResult = StepTransformResult.withoutHold(createdProducer).build();
     context.handleResult(null, ImmutableList.of(), finishedResult);

     final CountDownLatch callLatch = new CountDownLatch(1);
     context.extractFiredTimers();
     Runnable callback = callLatch::countDown;
     context.scheduleAfterOutputWouldBeProduced(
         downstream, GlobalWindow.INSTANCE, WindowingStrategy.globalDefault(), callback);
     assertThat(callLatch.await(1, TimeUnit.SECONDS), is(true));
   }

   @Test
   public void extractFiredTimersExtractsTimers() {
     TransformResult<?> holdResult =
         StepTransformResult.withHold(createdProducer, new Instant(0)).build();
     context.handleResult(null, ImmutableList.of(), holdResult);

     StructuralKey<?> key = StructuralKey.of("foo".length(), VarIntCoder.of());
     TimerData toFire =
         TimerData.of(StateNamespaces.global(), new Instant(100L), TimeDomain.EVENT_TIME);
     TransformResult<?> timerResult =
         StepTransformResult.withoutHold(downstreamProducer)
             .withState(CopyOnAccessInMemoryStateInternals.withUnderlying(key, null))
             .withTimerUpdate(TimerUpdate.builder(key).setTimer(toFire).build())
             .build();

     // haven't added any timers, must be empty
     assertThat(context.extractFiredTimers(), emptyIterable());
     context.handleResult(
         context.createKeyedBundle(key, created).commit(Instant.now()),
         ImmutableList.of(),
         timerResult);

     // timer hasn't fired
     assertThat(context.extractFiredTimers(), emptyIterable());

     TransformResult<?> advanceResult = StepTransformResult.withoutHold(createdProducer).build();
     // Should cause the downstream timer to fire
     context.handleResult(null, ImmutableList.of(), advanceResult);

     Collection<FiredTimers<AppliedPTransform<?, ?, ?>>> fired = context.extractFiredTimers();
     assertThat(Iterables.getOnlyElement(fired).getKey(), equalTo(key));

     FiredTimers<AppliedPTransform<?, ?, ?>> firedForKey = Iterables.getOnlyElement(fired);
     // Contains exclusively the fired timer
     assertThat(firedForKey.getTimers(), contains(toFire));

     // Don't reextract timers
     assertThat(context.extractFiredTimers(), emptyIterable());
   }

   @Test
   public void createKeyedBundleKeyed() {
     StructuralKey<String> key = StructuralKey.of("foo", StringUtf8Coder.of());
     CommittedBundle<KV<String, Integer>> keyedBundle =
         context.createKeyedBundle(key, downstream).commit(Instant.now());
     assertThat(keyedBundle.getKey(), equalTo(key));
   }

   @Test
   public void isDoneWithUnboundedPCollection() {
     assertThat(context.isDone(unboundedProducer), is(false));

     context.handleResult(
         null, ImmutableList.of(), StepTransformResult.withoutHold(unboundedProducer).build());
     context.extractFiredTimers();
     assertThat(context.isDone(unboundedProducer), is(true));
   }

   @Test
   public void isDoneWithPartiallyDone() {
     assertThat(context.isDone(), is(false));

     UncommittedBundle<Integer> rootBundle = context.createBundle(created);
     rootBundle.add(WindowedValue.valueInGlobalWindow(1));
     CommittedResult handleResult =
         context.handleResult(
             null,
             ImmutableList.of(),
             StepTransformResult.<Integer>withoutHold(createdProducer)
                 .addOutput(rootBundle)
                 .build());
     @SuppressWarnings("unchecked")
     CommittedBundle<Integer> committedBundle =
         (CommittedBundle<Integer>) Iterables.getOnlyElement(handleResult.getOutputs());
     context.handleResult(
         null, ImmutableList.of(), StepTransformResult.withoutHold(unboundedProducer).build());
     assertThat(context.isDone(), is(false));

     for (AppliedPTransform<?, ?, ?> consumers : graph.getPerElementConsumers(created)) {
       context.handleResult(
           committedBundle, ImmutableList.of(), StepTransformResult.withoutHold(consumers).build());
     }
     context.extractFiredTimers();
     assertThat(context.isDone(), is(true));
   }

   private static class TestBoundedWindow extends BoundedWindow {
     private final Instant ts;

     public TestBoundedWindow(Instant ts) {
       this.ts = ts;
     }

     @Override
     public Instant maxTimestamp() {
       return ts;
     }
   }
 }
	/*
	* 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.direct;

	import static java.nio.charset.StandardCharsets.UTF_8;
	import static org.apache.beam.sdk.testing.PCollectionViewTesting.materializeValuesFor;
	import static org.hamcrest.Matchers.contains;
	import static org.hamcrest.Matchers.containsInAnyOrder;
	import static org.hamcrest.Matchers.emptyIterable;
	import static org.hamcrest.Matchers.equalTo;
	import static org.hamcrest.Matchers.is;
	import static org.hamcrest.Matchers.not;
	import static org.junit.Assert.assertThat;

	import java.util.Collection;
	import java.util.concurrent.CountDownLatch;
	import java.util.concurrent.Executors;
	import java.util.concurrent.TimeUnit;
	import org.apache.beam.runners.core.SideInputReader;
	import org.apache.beam.runners.core.StateNamespaces;
	import org.apache.beam.runners.core.StateTag;
	import org.apache.beam.runners.core.StateTags;
	import org.apache.beam.runners.core.TimerInternals.TimerData;
	import org.apache.beam.runners.direct.DirectExecutionContext.DirectStepContext;
	import org.apache.beam.runners.direct.WatermarkManager.FiredTimers;
	import org.apache.beam.runners.direct.WatermarkManager.TimerUpdate;
	import org.apache.beam.runners.local.StructuralKey;
	import org.apache.beam.sdk.coders.ByteArrayCoder;
	import org.apache.beam.sdk.coders.IterableCoder;
	import org.apache.beam.sdk.coders.KvCoder;
	import org.apache.beam.sdk.coders.StringUtf8Coder;
	import org.apache.beam.sdk.coders.VarIntCoder;
	import org.apache.beam.sdk.coders.VoidCoder;
	import org.apache.beam.sdk.io.GenerateSequence;
	import org.apache.beam.sdk.options.PipelineOptionsFactory;
	import org.apache.beam.sdk.runners.AppliedPTransform;
	import org.apache.beam.sdk.state.BagState;
	import org.apache.beam.sdk.state.TimeDomain;
	import org.apache.beam.sdk.testing.TestPipeline;
	import org.apache.beam.sdk.transforms.Create;
	import org.apache.beam.sdk.transforms.View;
	import org.apache.beam.sdk.transforms.WithKeys;
	import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
	import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
	import org.apache.beam.sdk.transforms.windowing.PaneInfo;
	import org.apache.beam.sdk.transforms.windowing.PaneInfo.Timing;
	import org.apache.beam.sdk.util.WindowedValue;
	import org.apache.beam.sdk.values.KV;
	import org.apache.beam.sdk.values.PCollection;
	import org.apache.beam.sdk.values.PCollection.IsBounded;
	import org.apache.beam.sdk.values.PCollectionView;
	import org.apache.beam.sdk.values.WindowingStrategy;
	import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableList;
	import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables;
	import org.joda.time.Instant;
	import org.junit.Before;
	import org.junit.Rule;
	import org.junit.Test;
	import org.junit.runner.RunWith;
	import org.junit.runners.JUnit4;

	/** Tests for {@link EvaluationContext}. */
	@RunWith(JUnit4.class)
	public class EvaluationContextTest {
	private EvaluationContext context;

	private PCollection<Integer> created;
	private PCollection<KV<String, Integer>> downstream;
	private PCollectionView<Iterable<Integer>> view;
	private PCollection<Long> unbounded;

	private DirectGraph graph;

	private AppliedPTransform<?, ?, ?> createdProducer;
	private AppliedPTransform<?, ?, ?> downstreamProducer;
	private AppliedPTransform<?, ?, ?> viewProducer;
	private AppliedPTransform<?, ?, ?> unboundedProducer;

	@Rule public TestPipeline p = TestPipeline.create().enableAbandonedNodeEnforcement(false);

	@Before
	public void setup() {
	DirectRunner runner = DirectRunner.fromOptions(PipelineOptionsFactory.create());

	created = p.apply(Create.of(1, 2, 3));
	downstream = created.apply(WithKeys.of("foo"));
	view = created.apply(View.asIterable());
	unbounded = p.apply(GenerateSequence.from(0));

	p.replaceAll(runner.defaultTransformOverrides());

	KeyedPValueTrackingVisitor keyedPValueTrackingVisitor = KeyedPValueTrackingVisitor.create();
	p.traverseTopologically(keyedPValueTrackingVisitor);

	BundleFactory bundleFactory = ImmutableListBundleFactory.create();
	DirectGraphs.performDirectOverrides(p);
	graph = DirectGraphs.getGraph(p);
	context =
	EvaluationContext.create(
	NanosOffsetClock.create(),
	bundleFactory,
	graph,
	keyedPValueTrackingVisitor.getKeyedPValues(),
	Executors.newSingleThreadExecutor());

	createdProducer = graph.getProducer(created);
	downstreamProducer = graph.getProducer(downstream);
	viewProducer = graph.getProducer(view);
	unboundedProducer = graph.getProducer(unbounded);
	}

	@Test
	public void writeToViewWriterThenReadReads() {
	PCollectionViewWriter<?, Iterable<Integer>> viewWriter =
	context.createPCollectionViewWriter(
	PCollection.createPrimitiveOutputInternal(
	p,
	WindowingStrategy.globalDefault(),
	IsBounded.BOUNDED,
	IterableCoder.of(KvCoder.of(VoidCoder.of(), VarIntCoder.of()))),
	view);
	BoundedWindow window = new TestBoundedWindow(new Instant(1024L));
	BoundedWindow second = new TestBoundedWindow(new Instant(899999L));
	ImmutableList.Builder<WindowedValue<?>> valuesBuilder = ImmutableList.builder();
	for (Object materializedValue : materializeValuesFor(View.asIterable(), 1)) {
	valuesBuilder.add(
	WindowedValue.of(
	materializedValue, new Instant(1222), window, PaneInfo.ON_TIME_AND_ONLY_FIRING));
	}
	for (Object materializedValue : materializeValuesFor(View.asIterable(), 2)) {
	valuesBuilder.add(
	WindowedValue.of(
	materializedValue,
	new Instant(8766L),
	second,
	PaneInfo.createPane(true, false, Timing.ON_TIME, 0, 0)));
	}
	viewWriter.add((Iterable) valuesBuilder.build());

	SideInputReader reader = context.createSideInputReader(ImmutableList.of(view));
	assertThat(reader.get(view, window), containsInAnyOrder(1));
	assertThat(reader.get(view, second), containsInAnyOrder(2));

	ImmutableList.Builder<WindowedValue<?>> overwrittenValuesBuilder = ImmutableList.builder();
	for (Object materializedValue : materializeValuesFor(View.asIterable(), 4444)) {
	overwrittenValuesBuilder.add(
	WindowedValue.of(
	materializedValue,
	new Instant(8677L),
	second,
	PaneInfo.createPane(false, true, Timing.LATE, 1, 1)));
	}
	viewWriter.add((Iterable) overwrittenValuesBuilder.build());
	assertThat(reader.get(view, second), containsInAnyOrder(2));
	// The cached value is served in the earlier reader
	reader = context.createSideInputReader(ImmutableList.of(view));
	assertThat(reader.get(view, second), containsInAnyOrder(4444));
	}

	@Test
	public void getExecutionContextSameStepSameKeyState() {
	DirectExecutionContext fooContext =
	context.getExecutionContext(createdProducer, StructuralKey.of("foo", StringUtf8Coder.of()));

	StateTag<BagState<Integer>> intBag = StateTags.bag("myBag", VarIntCoder.of());

	DirectStepContext stepContext = fooContext.getStepContext("s1");
	stepContext.stateInternals().state(StateNamespaces.global(), intBag).add(1);

	context.handleResult(
	ImmutableListBundleFactory.create()
	.createKeyedBundle(StructuralKey.of("foo", StringUtf8Coder.of()), created)
	.commit(Instant.now()),
	ImmutableList.of(),
	StepTransformResult.withoutHold(createdProducer)
	.withState(stepContext.commitState())
	.build());

	DirectExecutionContext secondFooContext =
	context.getExecutionContext(createdProducer, StructuralKey.of("foo", StringUtf8Coder.of()));
	assertThat(
	secondFooContext
	.getStepContext("s1")
	.stateInternals()
	.state(StateNamespaces.global(), intBag)
	.read(),
	contains(1));
	}

	@Test
	public void getExecutionContextDifferentKeysIndependentState() {
	DirectExecutionContext fooContext =
	context.getExecutionContext(createdProducer, StructuralKey.of("foo", StringUtf8Coder.of()));

	StateTag<BagState<Integer>> intBag = StateTags.bag("myBag", VarIntCoder.of());

	fooContext.getStepContext("s1").stateInternals().state(StateNamespaces.global(), intBag).add(1);

	DirectExecutionContext barContext =
	context.getExecutionContext(createdProducer, StructuralKey.of("bar", StringUtf8Coder.of()));
	assertThat(barContext, not(equalTo(fooContext)));
	assertThat(
	barContext
	.getStepContext("s1")
	.stateInternals()
	.state(StateNamespaces.global(), intBag)
	.read(),
	emptyIterable());
	}

	@Test
	public void getExecutionContextDifferentStepsIndependentState() {
	StructuralKey<?> myKey = StructuralKey.of("foo", StringUtf8Coder.of());
	DirectExecutionContext fooContext = context.getExecutionContext(createdProducer, myKey);

	StateTag<BagState<Integer>> intBag = StateTags.bag("myBag", VarIntCoder.of());

	fooContext.getStepContext("s1").stateInternals().state(StateNamespaces.global(), intBag).add(1);

	DirectExecutionContext barContext = context.getExecutionContext(downstreamProducer, myKey);
	assertThat(
	barContext
	.getStepContext("s1")
	.stateInternals()
	.state(StateNamespaces.global(), intBag)
	.read(),
	emptyIterable());
	}

	@Test
	public void handleResultStoresState() {
	StructuralKey<?> myKey = StructuralKey.of("foo".getBytes(UTF_8), ByteArrayCoder.of());
	DirectExecutionContext fooContext = context.getExecutionContext(downstreamProducer, myKey);

	StateTag<BagState<Integer>> intBag = StateTags.bag("myBag", VarIntCoder.of());

	CopyOnAccessInMemoryStateInternals<?> state = fooContext.getStepContext("s1").stateInternals();
	BagState<Integer> bag = state.state(StateNamespaces.global(), intBag);
	bag.add(1);
	bag.add(2);
	bag.add(4);

	TransformResult<?> stateResult =
	StepTransformResult.withoutHold(downstreamProducer).withState(state).build();

	context.handleResult(
	context.createKeyedBundle(myKey, created).commit(Instant.now()),
	ImmutableList.of(),
	stateResult);

	DirectExecutionContext afterResultContext =
	context.getExecutionContext(downstreamProducer, myKey);

	CopyOnAccessInMemoryStateInternals<?> afterResultState =
	afterResultContext.getStepContext("s1").stateInternals();
	assertThat(afterResultState.state(StateNamespaces.global(), intBag).read(), contains(1, 2, 4));
	}

	@Test
	public void callAfterOutputMustHaveBeenProducedAfterEndOfWatermarkCallsback() throws Exception {
	final CountDownLatch callLatch = new CountDownLatch(1);
	Runnable callback = callLatch::countDown;

	// Should call back after the end of the global window
	context.scheduleAfterOutputWouldBeProduced(
	downstream, GlobalWindow.INSTANCE, WindowingStrategy.globalDefault(), callback);

	TransformResult<?> result =
	StepTransformResult.withHold(createdProducer, new Instant(0)).build();

	context.handleResult(null, ImmutableList.of(), result);
	// Difficult to demonstrate that we took no action in a multithreaded world; poll for a bit
	// will likely be flaky if this logic is broken
	assertThat(callLatch.await(500L, TimeUnit.MILLISECONDS), is(false));

	TransformResult<?> finishedResult = StepTransformResult.withoutHold(createdProducer).build();
	context.handleResult(null, ImmutableList.of(), finishedResult);
	context.forceRefresh();
	// Obtain the value via blocking call
	assertThat(callLatch.await(1, TimeUnit.SECONDS), is(true));
	}

	@Test
	public void callAfterOutputMustHaveBeenProducedAlreadyAfterCallsImmediately() throws Exception {
	TransformResult<?> finishedResult = StepTransformResult.withoutHold(createdProducer).build();
	context.handleResult(null, ImmutableList.of(), finishedResult);

	final CountDownLatch callLatch = new CountDownLatch(1);
	context.extractFiredTimers();
	Runnable callback = callLatch::countDown;
	context.scheduleAfterOutputWouldBeProduced(
	downstream, GlobalWindow.INSTANCE, WindowingStrategy.globalDefault(), callback);
	assertThat(callLatch.await(1, TimeUnit.SECONDS), is(true));
	}

	@Test
	public void extractFiredTimersExtractsTimers() {
	TransformResult<?> holdResult =
	StepTransformResult.withHold(createdProducer, new Instant(0)).build();
	context.handleResult(null, ImmutableList.of(), holdResult);

	StructuralKey<?> key = StructuralKey.of("foo".length(), VarIntCoder.of());
	TimerData toFire =
	TimerData.of(StateNamespaces.global(), new Instant(100L), TimeDomain.EVENT_TIME);
	TransformResult<?> timerResult =
	StepTransformResult.withoutHold(downstreamProducer)
	.withState(CopyOnAccessInMemoryStateInternals.withUnderlying(key, null))
	.withTimerUpdate(TimerUpdate.builder(key).setTimer(toFire).build())
	.build();

	// haven't added any timers, must be empty
	assertThat(context.extractFiredTimers(), emptyIterable());
	context.handleResult(
	context.createKeyedBundle(key, created).commit(Instant.now()),
	ImmutableList.of(),
	timerResult);

	// timer hasn't fired
	assertThat(context.extractFiredTimers(), emptyIterable());

	TransformResult<?> advanceResult = StepTransformResult.withoutHold(createdProducer).build();
	// Should cause the downstream timer to fire
	context.handleResult(null, ImmutableList.of(), advanceResult);

	Collection<FiredTimers<AppliedPTransform<?, ?, ?>>> fired = context.extractFiredTimers();
	assertThat(Iterables.getOnlyElement(fired).getKey(), equalTo(key));

	FiredTimers<AppliedPTransform<?, ?, ?>> firedForKey = Iterables.getOnlyElement(fired);
	// Contains exclusively the fired timer
	assertThat(firedForKey.getTimers(), contains(toFire));

	// Don't reextract timers
	assertThat(context.extractFiredTimers(), emptyIterable());
	}

	@Test
	public void createKeyedBundleKeyed() {
	StructuralKey<String> key = StructuralKey.of("foo", StringUtf8Coder.of());
	CommittedBundle<KV<String, Integer>> keyedBundle =
	context.createKeyedBundle(key, downstream).commit(Instant.now());
	assertThat(keyedBundle.getKey(), equalTo(key));
	}

	@Test
	public void isDoneWithUnboundedPCollection() {
	assertThat(context.isDone(unboundedProducer), is(false));

	context.handleResult(
	null, ImmutableList.of(), StepTransformResult.withoutHold(unboundedProducer).build());
	context.extractFiredTimers();
	assertThat(context.isDone(unboundedProducer), is(true));
	}

	@Test
	public void isDoneWithPartiallyDone() {
	assertThat(context.isDone(), is(false));

	UncommittedBundle<Integer> rootBundle = context.createBundle(created);
	rootBundle.add(WindowedValue.valueInGlobalWindow(1));
	CommittedResult handleResult =
	context.handleResult(
	null,
	ImmutableList.of(),
	StepTransformResult.<Integer>withoutHold(createdProducer)
	.addOutput(rootBundle)
	.build());
	@SuppressWarnings("unchecked")
	CommittedBundle<Integer> committedBundle =
	(CommittedBundle<Integer>) Iterables.getOnlyElement(handleResult.getOutputs());
	context.handleResult(
	null, ImmutableList.of(), StepTransformResult.withoutHold(unboundedProducer).build());
	assertThat(context.isDone(), is(false));

	for (AppliedPTransform<?, ?, ?> consumers : graph.getPerElementConsumers(created)) {
	context.handleResult(
	committedBundle, ImmutableList.of(), StepTransformResult.withoutHold(consumers).build());
	}
	context.extractFiredTimers();
	assertThat(context.isDone(), is(true));
	}

	private static class TestBoundedWindow extends BoundedWindow {
	private final Instant ts;

	public TestBoundedWindow(Instant ts) {
	this.ts = ts;
	}

	@Override
	public Instant maxTimestamp() {
	return ts;
	}
	}
	}