beam-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From aljos...@apache.org
Subject [1/2] beam git commit: [BEAM-1456] Make UnboundedSourceWrapper snapshot to rescalable operator state in Flink Runner
Date Fri, 24 Feb 2017 11:13:21 GMT
Repository: beam
Updated Branches:
  refs/heads/master 2c71354d0 -> 3d0fe8539


[BEAM-1456] Make UnboundedSourceWrapper snapshot to rescalable operator state in Flink Runner


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

Branch: refs/heads/master
Commit: 5e3fb8c353d40b17cd4ed7332a9ecb7425b64c86
Parents: 2c71354
Author: JingsongLi <lzljs3620320@aliyun.com>
Authored: Sat Feb 18 13:04:48 2017 +0800
Committer: Aljoscha Krettek <aljoscha.krettek@gmail.com>
Committed: Fri Feb 24 12:04:47 2017 +0100

----------------------------------------------------------------------
 .../streaming/io/UnboundedSourceWrapper.java    | 136 ++++++++++---------
 .../streaming/UnboundedSourceWrapperTest.java   | 116 +++++++++++++++-
 2 files changed, 185 insertions(+), 67 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/beam/blob/5e3fb8c3/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapper.java
----------------------------------------------------------------------
diff --git a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapper.java
b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapper.java
index 237e5a3..2849464 100644
--- a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapper.java
+++ b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapper.java
@@ -18,17 +18,14 @@
 package org.apache.beam.runners.flink.translation.wrappers.streaming.io;
 
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.base.Function;
-import com.google.common.collect.Lists;
-import java.io.ByteArrayInputStream;
 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.List;
+import org.apache.beam.runners.flink.translation.types.CoderTypeInformation;
 import org.apache.beam.runners.flink.translation.utils.SerializedPipelineOptions;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.KvCoder;
-import org.apache.beam.sdk.coders.ListCoder;
 import org.apache.beam.sdk.coders.SerializableCoder;
 import org.apache.beam.sdk.io.UnboundedSource;
 import org.apache.beam.sdk.options.PipelineOptions;
@@ -37,11 +34,17 @@ import org.apache.beam.sdk.transforms.windowing.PaneInfo;
 import org.apache.beam.sdk.util.WindowedValue;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.TypeDescriptor;
-import org.apache.commons.io.output.ByteArrayOutputStream;
+import org.apache.flink.api.common.ExecutionConfig;
 import org.apache.flink.api.common.functions.StoppableFunction;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.state.OperatorStateStore;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.runtime.state.CheckpointListener;
-import org.apache.flink.streaming.api.checkpoint.Checkpointed;
+import org.apache.flink.runtime.state.DefaultOperatorStateBackend;
+import org.apache.flink.runtime.state.FunctionInitializationContext;
+import org.apache.flink.runtime.state.FunctionSnapshotContext;
+import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
 import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;
 import org.apache.flink.streaming.api.operators.StreamingRuntimeContext;
 import org.apache.flink.streaming.api.watermark.Watermark;
@@ -56,7 +59,8 @@ import org.slf4j.LoggerFactory;
 public class UnboundedSourceWrapper<
     OutputT, CheckpointMarkT extends UnboundedSource.CheckpointMark>
     extends RichParallelSourceFunction<WindowedValue<OutputT>>
-    implements ProcessingTimeCallback, StoppableFunction, Checkpointed<byte[]>, CheckpointListener
{
+    implements ProcessingTimeCallback, StoppableFunction,
+    CheckpointListener, CheckpointedFunction {
 
   private static final Logger LOG = LoggerFactory.getLogger(UnboundedSourceWrapper.class);
 
@@ -68,8 +72,8 @@ public class UnboundedSourceWrapper<
   /**
    * For snapshot and restore.
    */
-  private final ListCoder<
-      KV<? extends UnboundedSource<OutputT, CheckpointMarkT>, CheckpointMarkT>>
checkpointCoder;
+  private final KvCoder<
+      ? extends UnboundedSource<OutputT, CheckpointMarkT>, CheckpointMarkT> checkpointCoder;
 
   /**
    * The split sources. We split them in the constructor to ensure that all parallel
@@ -117,12 +121,13 @@ public class UnboundedSourceWrapper<
    */
   private static final int MAX_NUMBER_PENDING_CHECKPOINTS = 32;
 
+  private transient ListState<KV<? extends
+      UnboundedSource<OutputT, CheckpointMarkT>, CheckpointMarkT>> stateForCheckpoint;
+
   /**
-   * When restoring from a snapshot we put the restored sources/checkpoint marks here
-   * and open in {@link #open(Configuration)}.
+   * false if checkpointCoder is null or no restore state by starting first.
    */
-  private transient List<
-      KV<? extends UnboundedSource<OutputT, CheckpointMarkT>, CheckpointMarkT>>
restoredState;
+  private transient boolean isRestored = false;
 
   @SuppressWarnings("unchecked")
   public UnboundedSourceWrapper(
@@ -145,7 +150,7 @@ public class UnboundedSourceWrapper<
           (Coder) SerializableCoder.of(new TypeDescriptor<UnboundedSource>() {
           });
 
-      checkpointCoder = (ListCoder) ListCoder.of(KvCoder.of(sourceCoder, checkpointMarkCoder));
+      checkpointCoder = KvCoder.of(sourceCoder, checkpointMarkCoder);
     }
 
     // get the splits early. we assume that the generated splits are stable,
@@ -171,30 +176,14 @@ public class UnboundedSourceWrapper<
 
     pendingCheckpoints = new LinkedHashMap<>();
 
-    if (restoredState != null) {
-
+    if (isRestored) {
       // restore the splitSources from the checkpoint to ensure consistent ordering
-      // do it using a transform because otherwise we would have to do
-      // unchecked casts
-      localSplitSources = Lists.transform(
-          restoredState,
-          new Function<
-              KV<? extends UnboundedSource<OutputT, CheckpointMarkT>, CheckpointMarkT>,
-              UnboundedSource<OutputT, CheckpointMarkT>>() {
-            @Override
-            public UnboundedSource<OutputT, CheckpointMarkT> apply(
-                KV<? extends UnboundedSource<OutputT, CheckpointMarkT>, CheckpointMarkT>
input) {
-              return input.getKey();
-            }
-          });
-
       for (KV<? extends UnboundedSource<OutputT, CheckpointMarkT>, CheckpointMarkT>
restored:
-          restoredState) {
-        localReaders.add(
-            restored.getKey().createReader(
-                serializedOptions.getPipelineOptions(), restored.getValue()));
+          stateForCheckpoint.get()) {
+        localSplitSources.add(restored.getKey());
+        localReaders.add(restored.getKey().createReader(
+            serializedOptions.getPipelineOptions(), restored.getValue()));
       }
-      restoredState = null;
     } else {
       // initialize localReaders and localSources from scratch
       for (int i = 0; i < splitSources.size(); i++) {
@@ -342,37 +331,42 @@ public class UnboundedSourceWrapper<
     isRunning = false;
   }
 
+  // ------------------------------------------------------------------------
+  //  Checkpoint and restore
+  // ------------------------------------------------------------------------
+
   @Override
-  public byte[] snapshotState(long checkpointId, long checkpointTimestamp) throws Exception
{
+  public void snapshotState(FunctionSnapshotContext functionSnapshotContext) throws Exception
{
+    if (!isRunning) {
+      LOG.debug("snapshotState() called on closed source");
+    } else {
 
-    if (checkpointCoder == null) {
-      // no checkpoint coder available in this source
-      return null;
-    }
+      if (checkpointCoder == null) {
+        // no checkpoint coder available in this source
+        return;
+      }
 
-    // we checkpoint the sources along with the CheckpointMarkT to ensure
-    // than we have a correct mapping of checkpoints to sources when
-    // restoring
-    List<KV<? extends UnboundedSource<OutputT, CheckpointMarkT>, CheckpointMarkT>>
checkpoints =
-        new ArrayList<>(localSplitSources.size());
-    List<CheckpointMarkT> checkpointMarks = new ArrayList<>(localSplitSources.size());
-
-    for (int i = 0; i < localSplitSources.size(); i++) {
-      UnboundedSource<OutputT, CheckpointMarkT> source = localSplitSources.get(i);
-      UnboundedSource.UnboundedReader<OutputT> reader = localReaders.get(i);
-
-      @SuppressWarnings("unchecked")
-      CheckpointMarkT mark = (CheckpointMarkT) reader.getCheckpointMark();
-      checkpointMarks.add(mark);
-      KV<UnboundedSource<OutputT, CheckpointMarkT>, CheckpointMarkT> kv =
-          KV.of(source, mark);
-      checkpoints.add(kv);
-    }
+      stateForCheckpoint.clear();
+
+      long checkpointId = functionSnapshotContext.getCheckpointId();
+
+      // we checkpoint the sources along with the CheckpointMarkT to ensure
+      // than we have a correct mapping of checkpoints to sources when
+      // restoring
+      List<CheckpointMarkT> checkpointMarks = new ArrayList<>(localSplitSources.size());
+
+      for (int i = 0; i < localSplitSources.size(); i++) {
+        UnboundedSource<OutputT, CheckpointMarkT> source = localSplitSources.get(i);
+        UnboundedSource.UnboundedReader<OutputT> reader = localReaders.get(i);
+
+        @SuppressWarnings("unchecked")
+        CheckpointMarkT mark = (CheckpointMarkT) reader.getCheckpointMark();
+        checkpointMarks.add(mark);
+        KV<UnboundedSource<OutputT, CheckpointMarkT>, CheckpointMarkT> kv =
+            KV.of(source, mark);
+        stateForCheckpoint.add(kv);
+      }
 
-    try (ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
-      checkpointCoder.encode(checkpoints, baos, Coder.Context.OUTER);
-      return baos.toByteArray();
-    } finally {
       // cleanup old pending checkpoints and add new checkpoint
       int diff = pendingCheckpoints.size() - MAX_NUMBER_PENDING_CHECKPOINTS;
       if (diff >= 0) {
@@ -384,18 +378,30 @@ public class UnboundedSourceWrapper<
         }
       }
       pendingCheckpoints.put(checkpointId, checkpointMarks);
+
     }
   }
 
   @Override
-  public void restoreState(byte[] bytes) throws Exception {
+  public void initializeState(FunctionInitializationContext context) throws Exception {
     if (checkpointCoder == null) {
       // no checkpoint coder available in this source
       return;
     }
 
-    try (ByteArrayInputStream bais = new ByteArrayInputStream(bytes)) {
-      restoredState = checkpointCoder.decode(bais, Coder.Context.OUTER);
+    OperatorStateStore stateStore = context.getOperatorStateStore();
+    CoderTypeInformation<
+        KV<? extends UnboundedSource<OutputT, CheckpointMarkT>, CheckpointMarkT>>
+        typeInformation = (CoderTypeInformation) new CoderTypeInformation<>(checkpointCoder);
+    stateForCheckpoint = stateStore.getOperatorState(
+        new ListStateDescriptor<>(DefaultOperatorStateBackend.DEFAULT_OPERATOR_STATE_NAME,
+            typeInformation.createSerializer(new ExecutionConfig())));
+
+    if (context.isRestored()) {
+      isRestored = true;
+      LOG.info("Having restore state in the UnbounedSourceWrapper.");
+    } else {
+      LOG.info("No restore state for UnbounedSourceWrapper.");
     }
   }
 

http://git-wip-us.apache.org/repos/asf/beam/blob/5e3fb8c3/runners/flink/runner/src/test/java/org/apache/beam/runners/flink/streaming/UnboundedSourceWrapperTest.java
----------------------------------------------------------------------
diff --git a/runners/flink/runner/src/test/java/org/apache/beam/runners/flink/streaming/UnboundedSourceWrapperTest.java
b/runners/flink/runner/src/test/java/org/apache/beam/runners/flink/streaming/UnboundedSourceWrapperTest.java
index 5b3d088..90f95d6 100644
--- a/runners/flink/runner/src/test/java/org/apache/beam/runners/flink/streaming/UnboundedSourceWrapperTest.java
+++ b/runners/flink/runner/src/test/java/org/apache/beam/runners/flink/streaming/UnboundedSourceWrapperTest.java
@@ -28,17 +28,25 @@ import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Set;
 import org.apache.beam.runners.flink.translation.wrappers.streaming.io.UnboundedSourceWrapper;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.io.UnboundedSource;
 import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.util.WindowedValue;
 import org.apache.beam.sdk.values.KV;
 import org.apache.flink.api.common.ExecutionConfig;
 import org.apache.flink.api.common.accumulators.Accumulator;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.state.OperatorStateStore;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.runtime.execution.Environment;
 import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContextSynchronousImpl;
 import org.apache.flink.streaming.api.TimeCharacteristic;
 import org.apache.flink.streaming.api.graph.StreamConfig;
 import org.apache.flink.streaming.api.operators.Output;
@@ -53,6 +61,7 @@ import org.junit.Test;
 import org.junit.experimental.runners.Enclosed;
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
+import org.mockito.Matchers;
 
 /**
  * Tests for {@link UnboundedSourceWrapper}.
@@ -180,6 +189,22 @@ public class UnboundedSourceWrapperTest {
               KV<Integer, Integer>,
               TestCountingSource.CounterMark>> sourceOperator = new StreamSource<>(flinkWrapper);
 
+
+      OperatorStateStore backend = mock(OperatorStateStore.class);
+
+      TestingListState<KV<UnboundedSource, TestCountingSource.CounterMark>>
+          listState = new TestingListState<>();
+
+      when(backend.getOperatorState(Matchers.any(ListStateDescriptor.class)))
+          .thenReturn(listState);
+
+      StateInitializationContext initializationContext = mock(StateInitializationContext.class);
+
+      when(initializationContext.getOperatorStateStore()).thenReturn(backend);
+      when(initializationContext.isRestored()).thenReturn(false, true);
+
+      flinkWrapper.initializeState(initializationContext);
+
       setupSourceOperator(sourceOperator, numTasks);
 
       final Set<KV<Integer, Integer>> emittedElements = new HashSet<>();
@@ -224,7 +249,16 @@ public class UnboundedSourceWrapperTest {
       assertTrue("Did not successfully read first batch of elements.", readFirstBatchOfElements);
 
       // draw a snapshot
-      byte[] snapshot = flinkWrapper.snapshotState(0, 0);
+      flinkWrapper.snapshotState(new StateSnapshotContextSynchronousImpl(0, 0));
+
+      // test snapshot offsets
+      assertEquals(flinkWrapper.getLocalSplitSources().size(),
+          listState.getList().size());
+      int totalEmit = 0;
+      for (KV<UnboundedSource, TestCountingSource.CounterMark> kv : listState.get())
{
+        totalEmit += kv.getValue().current + 1;
+      }
+      assertEquals(numElements / 2, totalEmit);
 
       // test that finalizeCheckpoint on CheckpointMark is called
       final ArrayList<Integer> finalizeList = new ArrayList<>();
@@ -250,7 +284,7 @@ public class UnboundedSourceWrapperTest {
       setupSourceOperator(restoredSourceOperator, numTasks);
 
       // restore snapshot
-      restoredFlinkWrapper.restoreState(snapshot);
+      restoredFlinkWrapper.initializeState(initializationContext);
 
       boolean readSecondBatchOfElements = false;
 
@@ -297,6 +331,58 @@ public class UnboundedSourceWrapperTest {
       assertTrue(emittedElements.size() == numElements);
     }
 
+    @Test
+    public void testNullCheckpoint() throws Exception {
+      final int numElements = 20;
+      PipelineOptions options = PipelineOptionsFactory.create();
+
+      TestCountingSource source = new TestCountingSource(numElements) {
+        @Override
+        public Coder<CounterMark> getCheckpointMarkCoder() {
+          return null;
+        }
+      };
+      UnboundedSourceWrapper<KV<Integer, Integer>, TestCountingSource.CounterMark>
flinkWrapper =
+          new UnboundedSourceWrapper<>(options, source, numSplits);
+
+      OperatorStateStore backend = mock(OperatorStateStore.class);
+
+      TestingListState<KV<UnboundedSource, TestCountingSource.CounterMark>>
+          listState = new TestingListState<>();
+
+      when(backend.getOperatorState(Matchers.any(ListStateDescriptor.class)))
+          .thenReturn(listState);
+
+      StateInitializationContext initializationContext = mock(StateInitializationContext.class);
+
+      when(initializationContext.getOperatorStateStore()).thenReturn(backend);
+      when(initializationContext.isRestored()).thenReturn(false, true);
+
+      flinkWrapper.initializeState(initializationContext);
+
+      StreamSource sourceOperator = new StreamSource<>(flinkWrapper);
+      setupSourceOperator(sourceOperator, numTasks);
+      sourceOperator.open();
+
+      flinkWrapper.snapshotState(new StateSnapshotContextSynchronousImpl(0, 0));
+
+      assertEquals(0, listState.getList().size());
+
+      UnboundedSourceWrapper<
+          KV<Integer, Integer>, TestCountingSource.CounterMark> restoredFlinkWrapper
=
+          new UnboundedSourceWrapper<>(options, new TestCountingSource(numElements),
+              numSplits);
+
+      StreamSource restoredSourceOperator = new StreamSource<>(flinkWrapper);
+      setupSourceOperator(restoredSourceOperator, numTasks);
+      sourceOperator.open();
+
+      restoredFlinkWrapper.initializeState(initializationContext);
+
+      assertEquals(Math.max(1, numSplits / numTasks), flinkWrapper.getLocalSplitSources().size());
+
+    }
+
     @SuppressWarnings("unchecked")
     private static <T> void setupSourceOperator(StreamSource<T, ?> operator,
int numSubTasks) {
       ExecutionConfig executionConfig = new ExecutionConfig();
@@ -349,4 +435,30 @@ public class UnboundedSourceWrapperTest {
     }
 
   }
+
+  private static final class TestingListState<T> implements ListState<T> {
+
+    private final List<T> list = new ArrayList<>();
+
+    @Override
+    public void clear() {
+      list.clear();
+    }
+
+    @Override
+    public Iterable<T> get() throws Exception {
+      return list;
+    }
+
+    @Override
+    public void add(T value) throws Exception {
+      list.add(value);
+    }
+
+    public List<T> getList() {
+      return list;
+    }
+
+  }
+
 }


Mime
View raw message