Add DrainMode to PaneInfo

Author: kennknowles

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/PaneInfo.java

@@ -30,6 +30,7 @@
 import org.apache.beam.sdk.transforms.DoFn.WindowedContext;
 import org.apache.beam.sdk.transforms.GroupByKey;
 import org.apache.beam.sdk.util.VarInt;
+import org.apache.beam.sdk.values.DrainMode;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
@@ -140,16 +141,32 @@ private static byte encodedByte(boolean isFirst, boolean isLast, Timing timing)
     return result;
   }
 
+  private static byte encodeExtendedMetadataByte(DrainMode drainMode) {
+    return (byte) (drainMode.ordinal() << 6);
+  }
+
+  private static DrainMode drainModeFromExtendedMetadata(byte extendedMetadataByte) {
+    return DrainMode.values()[(extendedMetadataByte >>> 6) & 0b11];
+  }
+
   private static final ImmutableMap<Byte, PaneInfo> BYTE_TO_PANE_INFO;
 
   static {
     ImmutableMap.Builder<Byte, PaneInfo> decodingBuilder = ImmutableMap.builder();
     for (Timing timing : Timing.values()) {
       long onTimeIndex = timing == Timing.EARLY ? -1 : 0;
-      register(decodingBuilder, new PaneInfo(true, true, timing, 0, onTimeIndex));
-      register(decodingBuilder, new PaneInfo(true, false, timing, 0, onTimeIndex));
-      register(decodingBuilder, new PaneInfo(false, true, timing, -1, onTimeIndex));
-      register(decodingBuilder, new PaneInfo(false, false, timing, -1, onTimeIndex));
+      register(
+          decodingBuilder,
+          new PaneInfo(true, true, timing, DrainMode.NOT_DRAINING, 0, onTimeIndex));
+      register(
+          decodingBuilder,
+          new PaneInfo(true, false, timing, DrainMode.NOT_DRAINING, 0, onTimeIndex));
+      register(
+          decodingBuilder,
+          new PaneInfo(false, true, timing, DrainMode.NOT_DRAINING, -1, onTimeIndex));
+      register(
+          decodingBuilder,
+          new PaneInfo(false, false, timing, DrainMode.NOT_DRAINING, -1, onTimeIndex));
     }
     BYTE_TO_PANE_INFO = decodingBuilder.build();
   }
@@ -158,11 +175,17 @@ private static void register(ImmutableMap.Builder<Byte, PaneInfo> builder, PaneI
     builder.put(info.encodedByte, info);
   }
 
+  // Byte containing the encoding tag, timing, isFirst, and isLast.
+  // These bytes are cached for all known possibilities to avoid repeat processing
   private final byte encodedByte;
 
+  // Extended metadata byte, containing drain mode
+  private final byte extendedMetadataByte;
+
   private final boolean isFirst;
   private final boolean isLast;
   private final Timing timing;
+  private final DrainMode drainMode;
   private final long index;
   private final long nonSpeculativeIndex;
 
@@ -177,18 +200,28 @@ private static void register(ImmutableMap.Builder<Byte, PaneInfo> builder, PaneI
   public static final PaneInfo ON_TIME_AND_ONLY_FIRING =
       PaneInfo.createPane(true, true, Timing.ON_TIME, 0, 0);
 
-  private PaneInfo(boolean isFirst, boolean isLast, Timing timing, long index, long onTimeIndex) {
+  private PaneInfo(
+      boolean isFirst,
+      boolean isLast,
+      Timing timing,
+      DrainMode drainMode,
+      long index,
+      long onTimeIndex) {
     this.encodedByte = encodedByte(isFirst, isLast, timing);
+    this.extendedMetadataByte =
+        drainMode == DrainMode.DRAINING ? encodeExtendedMetadataByte(drainMode) : 0x00;
     this.isFirst = isFirst;
     this.isLast = isLast;
     this.timing = timing;
+    this.drainMode = drainMode;
     this.index = index;
     this.nonSpeculativeIndex = onTimeIndex;
   }
 
   public static PaneInfo createPane(boolean isFirst, boolean isLast, Timing timing) {
     checkArgument(isFirst, "Indices must be provided for non-first pane info.");
-    return createPane(isFirst, isLast, timing, 0, timing == Timing.EARLY ? -1 : 0);
+    return createPane(
+        isFirst, isLast, timing, DrainMode.NOT_DRAINING, 0, timing == Timing.EARLY ? -1 : 0);
   }
 
   /** Factory method to create a {@link PaneInfo} with the specified parameters. */
@@ -197,7 +230,22 @@ public static PaneInfo createPane(
     if (isFirst || timing == Timing.UNKNOWN) {
       return checkNotNull(BYTE_TO_PANE_INFO.get(encodedByte(isFirst, isLast, timing)));
     } else {
-      return new PaneInfo(isFirst, isLast, timing, index, onTimeIndex);
+      return new PaneInfo(isFirst, isLast, timing, DrainMode.NOT_DRAINING, index, onTimeIndex);
+    }
+  }
+
+  /** Factory method to create a {@link PaneInfo} with the specified parameters. */
+  public static PaneInfo createPane(
+      boolean isFirst,
+      boolean isLast,
+      Timing timing,
+      DrainMode drainMode,
+      long index,
+      long onTimeIndex) {
+    if (drainMode != DrainMode.DRAINING && (isFirst || timing == Timing.UNKNOWN)) {
+      return checkNotNull(BYTE_TO_PANE_INFO.get(encodedByte(isFirst, isLast, timing)));
+    } else {
+      return new PaneInfo(isFirst, isLast, timing, drainMode, index, onTimeIndex);
     }
   }
 
@@ -241,6 +289,14 @@ public long getIndex() {
     return index;
   }
 
+  /**
+   * Indicates whether this element resulted from an aggregation that fired during a drain
+   * operation.
+   */
+  public DrainMode getDrainMode() {
+    return drainMode;
+  }
+
   /**
    * The zero-based index of this trigger firing among non-speculative panes.
    *
@@ -253,13 +309,17 @@ public long getNonSpeculativeIndex() {
     return nonSpeculativeIndex;
   }
 
-  int getEncodedByte() {
+  byte getEncodedByte() {
     return encodedByte;
   }
 
+  byte getExtendedMetadataByte() {
+    return extendedMetadataByte;
+  }
+
   @Override
   public int hashCode() {
-    return Objects.hash(encodedByte, index, nonSpeculativeIndex);
+    return Objects.hash(encodedByte, extendedMetadataByte, index, nonSpeculativeIndex);
   }
 
   @Override
@@ -270,6 +330,7 @@ public boolean equals(@Nullable Object obj) {
     } else if (obj instanceof PaneInfo) {
       PaneInfo that = (PaneInfo) obj;
       return this.encodedByte == that.encodedByte
+          && this.extendedMetadataByte == that.extendedMetadataByte
           && this.index == that.index
           && this.nonSpeculativeIndex == that.nonSpeculativeIndex;
     } else {
@@ -290,6 +351,7 @@ public String toString() {
         .add("timing", timing)
         .add("index", index)
         .add("onTimeIndex", nonSpeculativeIndex != -1 ? nonSpeculativeIndex : null)
+        .add("drainMode", drainMode == DrainMode.DRAINING ? "DRAINING" : null)
         .toString();
   }
 
@@ -298,7 +360,8 @@ public static class PaneInfoCoder extends AtomicCoder<PaneInfo> {
     private enum Encoding {
       FIRST,
       ONE_INDEX,
-      TWO_INDICES;
+      TWO_INDICES,
+      EXTENDED_METADATA;
 
       // NOTE: Do not reorder fields. The ordinal is used as part of
       // the encoding.
@@ -311,12 +374,15 @@ private enum Encoding {
       }
 
       public static Encoding fromTag(byte b) {
-        return Encoding.values()[b >> 4];
+        return Encoding.values()[(b >>> 4) & 0xF];
       }
     }
 
     private Encoding chooseEncoding(PaneInfo value) {
-      if ((value.index == 0 && value.nonSpeculativeIndex == 0) || value.timing == Timing.UNKNOWN) {
+      if (value.drainMode == DrainMode.DRAINING) {
+        return Encoding.EXTENDED_METADATA;
+      } else if ((value.index == 0 && value.nonSpeculativeIndex == 0)
+          || value.timing == Timing.UNKNOWN) {
         return Encoding.FIRST;
       } else if (value.index == value.nonSpeculativeIndex || value.timing == Timing.EARLY) {
         return Encoding.ONE_INDEX;
@@ -350,6 +416,12 @@ public void encode(PaneInfo value, final OutputStream outStream)
           VarInt.encode(value.index, outStream);
           VarInt.encode(value.nonSpeculativeIndex, outStream);
           break;
+        case EXTENDED_METADATA:
+          outStream.write(value.encodedByte | encoding.tag);
+          outStream.write(value.extendedMetadataByte);
+          VarInt.encode(value.index, outStream);
+          VarInt.encode(value.nonSpeculativeIndex, outStream);
+          break;
         default:
           throw new CoderException("Unknown encoding " + encoding);
       }
@@ -360,6 +432,7 @@ public PaneInfo decode(final InputStream inStream) throws CoderException, IOExce
       byte keyAndTag = (byte) inStream.read();
       PaneInfo base = Preconditions.checkNotNull(BYTE_TO_PANE_INFO.get((byte) (keyAndTag & 0x0F)));
       long index, onTimeIndex;
+      DrainMode drainMode = DrainMode.NOT_DRAINING;
       switch (Encoding.fromTag(keyAndTag)) {
         case FIRST:
           return base;
@@ -371,10 +444,16 @@ public PaneInfo decode(final InputStream inStream) throws CoderException, IOExce
           index = VarInt.decodeLong(inStream);
           onTimeIndex = VarInt.decodeLong(inStream);
           break;
+        case EXTENDED_METADATA:
+          byte extendedMetadata = (byte) inStream.read();
+          drainMode = drainModeFromExtendedMetadata(extendedMetadata);
+          index = VarInt.decodeLong(inStream);
+          onTimeIndex = VarInt.decodeLong(inStream);
+          break;
         default:
           throw new CoderException("Unknown encoding " + (keyAndTag & 0xF0));
       }
-      return new PaneInfo(base.isFirst, base.isLast, base.timing, index, onTimeIndex);
+      return new PaneInfo(base.isFirst, base.isLast, base.timing, drainMode, index, onTimeIndex);
     }
 
     @Override

sdks/java/core/src/main/java/org/apache/beam/sdk/values/DrainMode.java

@@ -0,0 +1,24 @@
+/*
+ * 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.sdk.values;
+
+/** {@link DrainMode} indicates whether a pipeline is known to be draining or not. */
+public enum DrainMode {
+  NOT_DRAINING,
+  DRAINING;
+}

sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/windowing/PaneInfoTest.java

@@ -23,6 +23,7 @@
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.testing.CoderProperties;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo.Timing;
+import org.apache.beam.sdk.values.DrainMode;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -41,19 +42,23 @@ public void testInterned() throws Exception {
   @Test
   public void testEncodingRoundTrip() throws Exception {
     Coder<PaneInfo> coder = PaneInfo.PaneInfoCoder.INSTANCE;
-    for (Timing timing : Timing.values()) {
-      long onTimeIndex = timing == Timing.EARLY ? -1 : 37;
-      CoderProperties.coderDecodeEncodeEqual(
-          coder, PaneInfo.createPane(false, false, timing, 389, onTimeIndex));
-      CoderProperties.coderDecodeEncodeEqual(
-          coder, PaneInfo.createPane(false, true, timing, 5077, onTimeIndex));
-      CoderProperties.coderDecodeEncodeEqual(coder, PaneInfo.createPane(true, false, timing, 0, 0));
-      CoderProperties.coderDecodeEncodeEqual(coder, PaneInfo.createPane(true, true, timing, 0, 0));
+    for (DrainMode drainMode : DrainMode.values()) {
+      for (Timing timing : Timing.values()) {
+        long onTimeIndex = timing == Timing.EARLY ? -1 : 37;
+        CoderProperties.coderDecodeEncodeEqual(
+            coder, PaneInfo.createPane(false, false, timing, drainMode, 389, onTimeIndex));
+        CoderProperties.coderDecodeEncodeEqual(
+            coder, PaneInfo.createPane(false, true, timing, drainMode, 5077, onTimeIndex));
+        CoderProperties.coderDecodeEncodeEqual(
+            coder, PaneInfo.createPane(true, false, timing, drainMode, 0, 0));
+        CoderProperties.coderDecodeEncodeEqual(
+            coder, PaneInfo.createPane(true, true, timing, drainMode, 0, 0));
+      }
     }
   }
 
   @Test
-  public void testEncodings() {
+  public void testHeaderByteEncoding() {
     assertEquals(
         "PaneInfo encoding assumes that there are only 4 Timing values.",
         4,
@@ -83,4 +88,17 @@ public void testEncodings() {
         0xF,
         PaneInfo.createPane(true, true, Timing.UNKNOWN).getEncodedByte());
   }
+
+  /**
+   * NOTE: this does not reproduce the logic in the main code by doing exhaustive testing
+   * programmatically. It is deliberately a smoke test of constant values.
+   */
+  @Test
+  public void testExtendedMetadataByteEncoding() {
+    assertEquals(
+        "PaneInfo with drainMode should have known extended metadata byte",
+        (byte) 0b01000000,
+        PaneInfo.createPane(true, true, Timing.UNKNOWN, DrainMode.DRAINING, 1, 1)
+            .getExtendedMetadataByte());
+  }
 }