beam-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From k...@apache.org
Subject [07/50] incubator-beam git commit: Flatten FiredTimers and ExtractFiredTimers
Date Tue, 08 Nov 2016 03:41:06 GMT
Flatten FiredTimers and ExtractFiredTimers

Pass a single collection of fired timers, and have those objects contain
the associated transform and key that they fired for. Timers already
contain the domain they are in.


Project: http://git-wip-us.apache.org/repos/asf/incubator-beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-beam/commit/5dca2674
Tree: http://git-wip-us.apache.org/repos/asf/incubator-beam/tree/5dca2674
Diff: http://git-wip-us.apache.org/repos/asf/incubator-beam/diff/5dca2674

Branch: refs/heads/gearpump-runner
Commit: 5dca2674a8d145c6e619005c2282c6064cd7aab7
Parents: 6e1e57b
Author: Thomas Groh <tgroh@google.com>
Authored: Thu Nov 3 14:10:37 2016 -0700
Committer: Thomas Groh <tgroh@google.com>
Committed: Fri Nov 4 13:05:21 2016 -0700

----------------------------------------------------------------------
 .../beam/runners/direct/EvaluationContext.java  |   6 +-
 .../direct/ExecutorServiceParallelExecutor.java |  41 ++++----
 .../beam/runners/direct/WatermarkManager.java   |  79 ++++++++------
 .../runners/direct/EvaluationContextTest.java   |  23 ++---
 .../runners/direct/WatermarkManagerTest.java    | 102 ++++++-------------
 5 files changed, 109 insertions(+), 142 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/5dca2674/runners/direct-java/src/main/java/org/apache/beam/runners/direct/EvaluationContext.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/EvaluationContext.java
b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/EvaluationContext.java
index 965e77d..b814def 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/EvaluationContext.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/EvaluationContext.java
@@ -391,11 +391,9 @@ class EvaluationContext {
    * <p>This is a destructive operation. Timers will only appear in the result of this
method once
    * for each time they are set.
    */
-  public Map<AppliedPTransform<?, ?, ?>, Map<StructuralKey<?>, FiredTimers>>
extractFiredTimers() {
+  public Collection<FiredTimers> extractFiredTimers() {
     forceRefresh();
-    Map<AppliedPTransform<?, ?, ?>, Map<StructuralKey<?>, FiredTimers>>
fired =
-        watermarkManager.extractFiredTimers();
-    return fired;
+    return watermarkManager.extractFiredTimers();
   }
 
   /**

http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/5dca2674/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ExecutorServiceParallelExecutor.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ExecutorServiceParallelExecutor.java
b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ExecutorServiceParallelExecutor.java
index e32f671..d1ffea1 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ExecutorServiceParallelExecutor.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ExecutorServiceParallelExecutor.java
@@ -47,7 +47,6 @@ import org.apache.beam.sdk.transforms.AppliedPTransform;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.util.KeyedWorkItem;
 import org.apache.beam.sdk.util.KeyedWorkItems;
-import org.apache.beam.sdk.util.TimeDomain;
 import org.apache.beam.sdk.util.TimerInternals.TimerData;
 import org.apache.beam.sdk.util.UserCodeException;
 import org.apache.beam.sdk.util.WindowedValue;
@@ -440,29 +439,23 @@ final class ExecutorServiceParallelExecutor implements PipelineExecutor
{
      */
     private void fireTimers() throws Exception {
       try {
-        for (Map.Entry<
-               AppliedPTransform<?, ?, ?>, Map<StructuralKey<?>, FiredTimers>>
transformTimers :
-            evaluationContext.extractFiredTimers().entrySet()) {
-          AppliedPTransform<?, ?, ?> transform = transformTimers.getKey();
-          for (Map.Entry<StructuralKey<?>, FiredTimers> keyTimers :
-              transformTimers.getValue().entrySet()) {
-            for (TimeDomain domain : TimeDomain.values()) {
-              Collection<TimerData> delivery = keyTimers.getValue().getTimers(domain);
-              if (delivery.isEmpty()) {
-                continue;
-              }
-              KeyedWorkItem<?, Object> work =
-                  KeyedWorkItems.timersWorkItem(keyTimers.getKey().getKey(), delivery);
-              @SuppressWarnings({"unchecked", "rawtypes"})
-              CommittedBundle<?> bundle =
-                  evaluationContext
-                      .createKeyedBundle(keyTimers.getKey(), (PCollection) transform.getInput())
-                      .add(WindowedValue.valueInGlobalWindow(work))
-                      .commit(evaluationContext.now());
-              scheduleConsumption(transform, bundle, new TimerIterableCompletionCallback(delivery));
-              state.set(ExecutorState.ACTIVE);
-            }
-          }
+        for (FiredTimers transformTimers : evaluationContext.extractFiredTimers()) {
+          Collection<TimerData> delivery = transformTimers.getTimers();
+          KeyedWorkItem<?, Object> work =
+              KeyedWorkItems.timersWorkItem(transformTimers.getKey().getKey(), delivery);
+          @SuppressWarnings({"unchecked", "rawtypes"})
+          CommittedBundle<?> bundle =
+              evaluationContext
+                  .createKeyedBundle(
+                      transformTimers.getKey(),
+                      (PCollection) transformTimers.getTransform().getInput())
+                  .add(WindowedValue.valueInGlobalWindow(work))
+                  .commit(evaluationContext.now());
+          scheduleConsumption(
+              transformTimers.getTransform(),
+              bundle,
+              new TimerIterableCompletionCallback(delivery));
+          state.set(ExecutorState.ACTIVE);
         }
       } catch (Exception e) {
         LOG.error("Internal Error while delivering timers", e);

http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/5dca2674/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkManager.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkManager.java
b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkManager.java
index 31b8091..f01c13c 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkManager.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkManager.java
@@ -740,14 +740,17 @@ public class WatermarkManager {
 
       wms =
           new TransformWatermarks(
-              inputWatermark, outputWatermark, inputProcessingWatermark, outputProcessingWatermark);
+              transform,
+              inputWatermark,
+              outputWatermark,
+              inputProcessingWatermark,
+              outputProcessingWatermark);
       transformToWatermarks.put(transform, wms);
     }
     return wms;
   }
 
-  private Collection<Watermark> getInputProcessingWatermarks(
-      AppliedPTransform<?, ?, ?> transform) {
+  private Collection<Watermark> getInputProcessingWatermarks(AppliedPTransform<?,
?, ?> transform) {
     ImmutableList.Builder<Watermark> inputWmsBuilder = ImmutableList.builder();
     Collection<? extends PValue> inputs = transform.getInput().expand();
     if (inputs.isEmpty()) {
@@ -924,15 +927,12 @@ public class WatermarkManager {
    * Returns a map of each {@link PTransform} that has pending timers to those timers. All
of the
    * pending timers will be removed from this {@link WatermarkManager}.
    */
-  public Map<AppliedPTransform<?, ?, ?>, Map<StructuralKey<?>, FiredTimers>>
extractFiredTimers() {
-    Map<AppliedPTransform<?, ?, ?>, Map<StructuralKey<?>, FiredTimers>>
allTimers = new HashMap<>();
+  public Collection<FiredTimers> extractFiredTimers() {
+    Collection<FiredTimers> allTimers = new ArrayList<>();
     for (Map.Entry<AppliedPTransform<?, ?, ?>, TransformWatermarks> watermarksEntry
:
         transformToWatermarks.entrySet()) {
-      Map<StructuralKey<?>, FiredTimers> keyFiredTimers =
-          watermarksEntry.getValue().extractFiredTimers();
-      if (!keyFiredTimers.isEmpty()) {
-        allTimers.put(watermarksEntry.getKey(), keyFiredTimers);
-      }
+      Collection<FiredTimers> firedTimers = watermarksEntry.getValue().extractFiredTimers();
+      allTimers.addAll(firedTimers);
     }
     return allTimers;
   }
@@ -1043,6 +1043,8 @@ public class WatermarkManager {
    * A reference to the input and output watermarks of an {@link AppliedPTransform}.
    */
   public class TransformWatermarks {
+    private final AppliedPTransform<?, ?, ?> transform;
+
     private final AppliedPTransformInputWatermark inputWatermark;
     private final AppliedPTransformOutputWatermark outputWatermark;
 
@@ -1053,10 +1055,12 @@ public class WatermarkManager {
     private Instant latestSynchronizedOutputWm;
 
     private TransformWatermarks(
+        AppliedPTransform<?, ?, ?> transform,
         AppliedPTransformInputWatermark inputWatermark,
         AppliedPTransformOutputWatermark outputWatermark,
         SynchronizedProcessingTimeInputWatermark inputSynchProcessingWatermark,
         SynchronizedProcessingTimeOutputWatermark outputSynchProcessingWatermark) {
+      this.transform = transform;
       this.inputWatermark = inputWatermark;
       this.outputWatermark = outputWatermark;
 
@@ -1128,7 +1132,7 @@ public class WatermarkManager {
       synchronizedProcessingInputWatermark.addPending(bundle);
     }
 
-    private Map<StructuralKey<?>, FiredTimers> extractFiredTimers() {
+    private Collection<FiredTimers> extractFiredTimers() {
       Map<StructuralKey<?>, List<TimerData>> eventTimeTimers =
           inputWatermark.extractFiredEventTimeTimers();
       Map<StructuralKey<?>, List<TimerData>> processingTimers;
@@ -1137,31 +1141,33 @@ public class WatermarkManager {
           TimeDomain.PROCESSING_TIME, clock.now());
       synchronizedTimers = synchronizedProcessingInputWatermark.extractFiredDomainTimers(
           TimeDomain.SYNCHRONIZED_PROCESSING_TIME, getSynchronizedProcessingInputTime());
-      Map<StructuralKey<?>, Map<TimeDomain, List<TimerData>>> groupedTimers
= new HashMap<>();
-      groupFiredTimers(groupedTimers, eventTimeTimers, processingTimers, synchronizedTimers);
 
-      Map<StructuralKey<?>, FiredTimers> keyFiredTimers = new HashMap<>();
-      for (Map.Entry<StructuralKey<?>, Map<TimeDomain, List<TimerData>>>
firedTimers :
-          groupedTimers.entrySet()) {
-        keyFiredTimers.put(firedTimers.getKey(), new FiredTimers(firedTimers.getValue()));
+      Map<StructuralKey<?>, List<TimerData>> timersPerKey =
+          groupFiredTimers(eventTimeTimers, processingTimers, synchronizedTimers);
+      Collection<FiredTimers> keyFiredTimers = new ArrayList<>(timersPerKey.size());
+      for (Map.Entry<StructuralKey<?>, List<TimerData>> firedTimers :
+          timersPerKey.entrySet()) {
+        keyFiredTimers.add(
+            new FiredTimers(transform, firedTimers.getKey(), firedTimers.getValue()));
       }
       return keyFiredTimers;
     }
 
     @SafeVarargs
-    private final void groupFiredTimers(
-        Map<StructuralKey<?>, Map<TimeDomain, List<TimerData>>> groupedToMutate,
+    private final Map<StructuralKey<?>, List<TimerData>> groupFiredTimers(
         Map<StructuralKey<?>, List<TimerData>>... timersToGroup) {
+      Map<StructuralKey<?>, List<TimerData>> groupedTimers = new HashMap<>();
       for (Map<StructuralKey<?>, List<TimerData>> subGroup : timersToGroup)
{
         for (Map.Entry<StructuralKey<?>, List<TimerData>> newTimers : subGroup.entrySet())
{
-          Map<TimeDomain, List<TimerData>> grouped = groupedToMutate.get(newTimers.getKey());
+          List<TimerData> grouped = groupedTimers.get(newTimers.getKey());
           if (grouped == null) {
-            grouped = new HashMap<>();
-            groupedToMutate.put(newTimers.getKey(), grouped);
+            grouped = new ArrayList<>();
+            groupedTimers.put(newTimers.getKey(), grouped);
           }
-          grouped.put(newTimers.getValue().get(0).getDomain(), newTimers.getValue());
+          grouped.addAll(newTimers.getValue());
         }
       }
+      return groupedTimers;
     }
 
     private void updateTimers(TimerUpdate update) {
@@ -1334,24 +1340,35 @@ public class WatermarkManager {
    * {@link WatermarkManager}.
    */
   public static class FiredTimers {
-    private final Map<TimeDomain, ? extends Collection<TimerData>> timers;
+    /** The transform the timers were set at and will be delivered to. */
+    private final AppliedPTransform<?, ?, ?> transform;
+    /** The key the timers were set for and will be delivered to. */
+    private final StructuralKey<?> key;
+    private final Collection<TimerData> timers;
 
-    private FiredTimers(Map<TimeDomain, ? extends Collection<TimerData>> timers)
{
+    private FiredTimers(
+        AppliedPTransform<?, ?, ?> transform, StructuralKey<?> key, Collection<TimerData>
timers) {
+      this.transform = transform;
+      this.key = key;
       this.timers = timers;
     }
 
+    public AppliedPTransform<?, ?, ?> getTransform() {
+      return transform;
+    }
+
+    public StructuralKey<?> getKey() {
+      return key;
+    }
+
     /**
      * Gets all of the timers that have fired within the provided {@link TimeDomain}. If
no timers
      * fired within the provided domain, return an empty collection.
      *
      * <p>Timers within a {@link TimeDomain} are guaranteed to be in order of increasing
timestamp.
      */
-    public Collection<TimerData> getTimers(TimeDomain domain) {
-      Collection<TimerData> domainTimers = timers.get(domain);
-      if (domainTimers == null) {
-        return Collections.emptyList();
-      }
-      return domainTimers;
+    public Collection<TimerData> getTimers() {
+      return timers;
     }
 
     @Override

http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/5dca2674/runners/direct-java/src/test/java/org/apache/beam/runners/direct/EvaluationContextTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/EvaluationContextTest.java
b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/EvaluationContextTest.java
index bc53570..e1277ac 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/EvaluationContextTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/EvaluationContextTest.java
@@ -373,36 +373,31 @@ public class EvaluationContextTest {
             .build();
 
     // haven't added any timers, must be empty
-    assertThat(context.extractFiredTimers().entrySet(), emptyIterable());
+    assertThat(context.extractFiredTimers(), emptyIterable());
     context.handleResult(
         context.createKeyedBundle(key, created).commit(Instant.now()),
         ImmutableList.<TimerData>of(),
         timerResult);
 
     // timer hasn't fired
-    assertThat(context.extractFiredTimers().entrySet(), emptyIterable());
+    assertThat(context.extractFiredTimers(), emptyIterable());
 
     TransformResult advanceResult =
         StepTransformResult.withoutHold(created.getProducingTransformInternal()).build();
     // Should cause the downstream timer to fire
     context.handleResult(null, ImmutableList.<TimerData>of(), advanceResult);
 
-    Map<AppliedPTransform<?, ?, ?>, Map<StructuralKey<?>, FiredTimers>>
fired =
-        context.extractFiredTimers();
+    Collection<FiredTimers> fired = context.extractFiredTimers();
     assertThat(
-        fired,
-        Matchers.<AppliedPTransform<?, ?, ?>>hasKey(downstream.getProducingTransformInternal()));
-    Map<StructuralKey<?>, FiredTimers> downstreamFired =
-        fired.get(downstream.getProducingTransformInternal());
-    assertThat(downstreamFired, Matchers.<Object>hasKey(key));
+        Iterables.getOnlyElement(fired).getKey(),
+        Matchers.<StructuralKey<?>>equalTo(key));
 
-    FiredTimers firedForKey = downstreamFired.get(key);
-    assertThat(firedForKey.getTimers(TimeDomain.PROCESSING_TIME), emptyIterable());
-    assertThat(firedForKey.getTimers(TimeDomain.SYNCHRONIZED_PROCESSING_TIME), emptyIterable());
-    assertThat(firedForKey.getTimers(TimeDomain.EVENT_TIME), contains(toFire));
+    FiredTimers firedForKey = Iterables.getOnlyElement(fired);
+    // Contains exclusively the fired timer
+    assertThat(firedForKey.getTimers(), contains(toFire));
 
     // Don't reextract timers
-    assertThat(context.extractFiredTimers().entrySet(), emptyIterable());
+    assertThat(context.extractFiredTimers(), emptyIterable());
   }
 
   @Test

http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/5dca2674/runners/direct-java/src/test/java/org/apache/beam/runners/direct/WatermarkManagerTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/WatermarkManagerTest.java
b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/WatermarkManagerTest.java
index 1954005..6bde462 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/WatermarkManagerTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/WatermarkManagerTest.java
@@ -22,7 +22,6 @@ import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.emptyIterable;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.not;
-import static org.hamcrest.Matchers.nullValue;
 import static org.junit.Assert.assertThat;
 
 import com.google.common.collect.ImmutableList;
@@ -68,6 +67,7 @@ import org.apache.beam.sdk.values.TimestampedValue;
 import org.hamcrest.BaseMatcher;
 import org.hamcrest.Description;
 import org.hamcrest.Matcher;
+import org.hamcrest.Matchers;
 import org.joda.time.Instant;
 import org.joda.time.ReadableInstant;
 import org.junit.Before;
@@ -915,12 +915,9 @@ public class WatermarkManagerTest implements Serializable {
         filteredDoubledWms.getSynchronizedProcessingOutputTime(),
         not(earlierThan(initialFilteredDoubledWm)));
 
-    Map<AppliedPTransform<?, ?, ?>, Map<StructuralKey<?>, FiredTimers>>
firedTimers =
-        manager.extractFiredTimers();
+    Collection<FiredTimers> firedTimers = manager.extractFiredTimers();
     assertThat(
-        firedTimers.get(filtered.getProducingTransformInternal())
-            .get(key)
-            .getTimers(TimeDomain.PROCESSING_TIME),
+        Iterables.getOnlyElement(firedTimers).getTimers(),
         contains(pastTimer));
     // Our timer has fired, but has not been completed, so it holds our synchronized processing
WM
     assertThat(filteredWms.getSynchronizedProcessingOutputTime(), not(laterThan(startTime)));
@@ -1099,10 +1096,9 @@ public class WatermarkManagerTest implements Serializable {
 
   @Test
   public void extractFiredTimersReturnsFiredEventTimeTimers() {
-    Map<AppliedPTransform<?, ?, ?>, Map<StructuralKey<?>, FiredTimers>>
initialTimers =
-        manager.extractFiredTimers();
+    Collection<FiredTimers> initialTimers = manager.extractFiredTimers();
     // Watermarks haven't advanced
-    assertThat(initialTimers.entrySet(), emptyIterable());
+    assertThat(initialTimers, emptyIterable());
 
     // Advance WM of keyed past the first timer, but ahead of the second and third
     CommittedBundle<Integer> createdBundle = multiWindowedBundle(filtered);
@@ -1136,15 +1132,11 @@ public class WatermarkManagerTest implements Serializable {
         new Instant(1000L));
     manager.refreshAll();
 
-    Map<AppliedPTransform<?, ?, ?>, Map<StructuralKey<?>, FiredTimers>>
firstTransformFiredTimers =
+    Collection<FiredTimers> firstFiredTimers =
         manager.extractFiredTimers();
-    assertThat(
-        firstTransformFiredTimers.get(filtered.getProducingTransformInternal()), not(nullValue()));
-    Map<StructuralKey<?>, FiredTimers> firstFilteredTimers =
-        firstTransformFiredTimers.get(filtered.getProducingTransformInternal());
-    assertThat(firstFilteredTimers.get(key), not(nullValue()));
-    FiredTimers firstFired = firstFilteredTimers.get(key);
-    assertThat(firstFired.getTimers(TimeDomain.EVENT_TIME), contains(earliestTimer));
+    assertThat(firstFiredTimers, not(Matchers.<FiredTimers>emptyIterable()));
+    FiredTimers firstFired = Iterables.getOnlyElement(firstFiredTimers);
+    assertThat(firstFired.getTimers(), contains(earliestTimer));
 
     manager.updateWatermarks(null,
         TimerUpdate.empty(),
@@ -1153,24 +1145,18 @@ public class WatermarkManagerTest implements Serializable {
             Collections.<CommittedBundle<?>>emptyList()),
         new Instant(50_000L));
     manager.refreshAll();
-    Map<AppliedPTransform<?, ?, ?>, Map<StructuralKey<?>, FiredTimers>>
secondTransformFiredTimers =
-        manager.extractFiredTimers();
-    assertThat(
-        secondTransformFiredTimers.get(filtered.getProducingTransformInternal()), not(nullValue()));
-    Map<StructuralKey<?>, FiredTimers> secondFilteredTimers =
-        secondTransformFiredTimers.get(filtered.getProducingTransformInternal());
-    assertThat(secondFilteredTimers.get(key), not(nullValue()));
-    FiredTimers secondFired = secondFilteredTimers.get(key);
+    Collection<FiredTimers> secondFiredTimers = manager.extractFiredTimers();
+    assertThat(secondFiredTimers, not(Matchers.<FiredTimers>emptyIterable()));
+    FiredTimers secondFired = Iterables.getOnlyElement(secondFiredTimers);
     // Contains, in order, middleTimer and then lastTimer
-    assertThat(secondFired.getTimers(TimeDomain.EVENT_TIME), contains(middleTimer, lastTimer));
+    assertThat(secondFired.getTimers(), contains(middleTimer, lastTimer));
   }
 
   @Test
   public void extractFiredTimersReturnsFiredProcessingTimeTimers() {
-    Map<AppliedPTransform<?, ?, ?>, Map<StructuralKey<?>, FiredTimers>>
initialTimers =
-        manager.extractFiredTimers();
+    Collection<FiredTimers> initialTimers = manager.extractFiredTimers();
     // Watermarks haven't advanced
-    assertThat(initialTimers.entrySet(), emptyIterable());
+    assertThat(initialTimers, emptyIterable());
 
     // Advance WM of keyed past the first timer, but ahead of the second and third
     CommittedBundle<Integer> createdBundle = multiWindowedBundle(filtered);
@@ -1204,15 +1190,10 @@ public class WatermarkManagerTest implements Serializable {
         new Instant(1000L));
     manager.refreshAll();
 
-    Map<AppliedPTransform<?, ?, ?>, Map<StructuralKey<?>, FiredTimers>>
firstTransformFiredTimers =
-        manager.extractFiredTimers();
-    assertThat(
-        firstTransformFiredTimers.get(filtered.getProducingTransformInternal()), not(nullValue()));
-    Map<StructuralKey<?>, FiredTimers> firstFilteredTimers =
-        firstTransformFiredTimers.get(filtered.getProducingTransformInternal());
-    assertThat(firstFilteredTimers.get(key), not(nullValue()));
-    FiredTimers firstFired = firstFilteredTimers.get(key);
-    assertThat(firstFired.getTimers(TimeDomain.PROCESSING_TIME), contains(earliestTimer));
+    Collection<FiredTimers> firstFiredTimers = manager.extractFiredTimers();
+    assertThat(firstFiredTimers, not(Matchers.<FiredTimers>emptyIterable()));
+    FiredTimers firstFired = Iterables.getOnlyElement(firstFiredTimers);
+    assertThat(firstFired.getTimers(), contains(earliestTimer));
 
     clock.set(new Instant(50_000L));
     manager.updateWatermarks(null,
@@ -1222,24 +1203,19 @@ public class WatermarkManagerTest implements Serializable {
             Collections.<CommittedBundle<?>>emptyList()),
         new Instant(50_000L));
     manager.refreshAll();
-    Map<AppliedPTransform<?, ?, ?>, Map<StructuralKey<?>, FiredTimers>>
secondTransformFiredTimers =
+    Collection<FiredTimers> secondFiredTimers =
         manager.extractFiredTimers();
-    assertThat(
-        secondTransformFiredTimers.get(filtered.getProducingTransformInternal()), not(nullValue()));
-    Map<StructuralKey<?>, FiredTimers> secondFilteredTimers =
-        secondTransformFiredTimers.get(filtered.getProducingTransformInternal());
-    assertThat(secondFilteredTimers.get(key), not(nullValue()));
-    FiredTimers secondFired = secondFilteredTimers.get(key);
+    assertThat(secondFiredTimers, not(Matchers.<FiredTimers>emptyIterable()));
+    FiredTimers secondFired = Iterables.getOnlyElement(secondFiredTimers);
     // Contains, in order, middleTimer and then lastTimer
-    assertThat(secondFired.getTimers(TimeDomain.PROCESSING_TIME), contains(middleTimer, lastTimer));
+    assertThat(secondFired.getTimers(), contains(middleTimer, lastTimer));
   }
 
   @Test
   public void extractFiredTimersReturnsFiredSynchronizedProcessingTimeTimers() {
-    Map<AppliedPTransform<?, ?, ?>, Map<StructuralKey<?>, FiredTimers>>
initialTimers =
-        manager.extractFiredTimers();
+    Collection<FiredTimers> initialTimers = manager.extractFiredTimers();
     // Watermarks haven't advanced
-    assertThat(initialTimers.entrySet(), emptyIterable());
+    assertThat(initialTimers, emptyIterable());
 
     // Advance WM of keyed past the first timer, but ahead of the second and third
     CommittedBundle<Integer> createdBundle = multiWindowedBundle(filtered);
@@ -1273,16 +1249,11 @@ public class WatermarkManagerTest implements Serializable {
         new Instant(1000L));
     manager.refreshAll();
 
-    Map<AppliedPTransform<?, ?, ?>, Map<StructuralKey<?>, FiredTimers>>
firstTransformFiredTimers =
+    Collection<FiredTimers> firstFiredTimers =
         manager.extractFiredTimers();
-    assertThat(
-        firstTransformFiredTimers.get(filtered.getProducingTransformInternal()), not(nullValue()));
-    Map<StructuralKey<?>, FiredTimers> firstFilteredTimers =
-        firstTransformFiredTimers.get(filtered.getProducingTransformInternal());
-    assertThat(firstFilteredTimers.get(key), not(nullValue()));
-    FiredTimers firstFired = firstFilteredTimers.get(key);
-    assertThat(
-        firstFired.getTimers(TimeDomain.SYNCHRONIZED_PROCESSING_TIME), contains(earliestTimer));
+    assertThat(firstFiredTimers, not(Matchers.<FiredTimers>emptyIterable()));
+    FiredTimers firstFired = Iterables.getOnlyElement(firstFiredTimers);
+    assertThat(firstFired.getTimers(), contains(earliestTimer));
 
     clock.set(new Instant(50_000L));
     manager.updateWatermarks(null,
@@ -1292,18 +1263,11 @@ public class WatermarkManagerTest implements Serializable {
             Collections.<CommittedBundle<?>>emptyList()),
         new Instant(50_000L));
     manager.refreshAll();
-    Map<AppliedPTransform<?, ?, ?>, Map<StructuralKey<?>, FiredTimers>>
secondTransformFiredTimers =
-        manager.extractFiredTimers();
-    assertThat(
-        secondTransformFiredTimers.get(filtered.getProducingTransformInternal()), not(nullValue()));
-    Map<StructuralKey<?>, FiredTimers> secondFilteredTimers =
-        secondTransformFiredTimers.get(filtered.getProducingTransformInternal());
-    assertThat(secondFilteredTimers.get(key), not(nullValue()));
-    FiredTimers secondFired = secondFilteredTimers.get(key);
+    Collection<FiredTimers> secondFiredTimers = manager.extractFiredTimers();
+    assertThat(secondFiredTimers, not(Matchers.<FiredTimers>emptyIterable()));
+    FiredTimers secondFired = Iterables.getOnlyElement(secondFiredTimers);
     // Contains, in order, middleTimer and then lastTimer
-    assertThat(
-        secondFired.getTimers(TimeDomain.SYNCHRONIZED_PROCESSING_TIME),
-        contains(middleTimer, lastTimer));
+    assertThat(secondFired.getTimers(), contains(middleTimer, lastTimer));
   }
 
   @Test


Mime
View raw message