Search in sources :

Example 31 with Server

use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server in project beam by apache.

the class StreamingDataflowWorkerTest method testHotKeyLogging.

@Test
public void testHotKeyLogging() throws Exception {
    // This is to test that the worker can correctly log the key from a hot key.
    List<ParallelInstruction> instructions = Arrays.asList(makeSourceInstruction(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())), makeSinkInstruction(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()), 0));
    FakeWindmillServer server = new FakeWindmillServer(errorCollector);
    server.setIsReady(false);
    StreamingConfigTask streamingConfig = new StreamingConfigTask();
    streamingConfig.setStreamingComputationConfigs(ImmutableList.of(makeDefaultStreamingComputationConfig(instructions)));
    streamingConfig.setWindmillServiceEndpoint("foo");
    WorkItem workItem = new WorkItem();
    workItem.setStreamingConfigTask(streamingConfig);
    when(mockWorkUnitClient.getGlobalStreamingConfigWorkItem()).thenReturn(Optional.of(workItem));
    StreamingDataflowWorkerOptions options = createTestingPipelineOptions(server, "--hotKeyLoggingEnabled=true");
    StreamingDataflowWorker worker = makeWorker(instructions, options, true);
    worker.start();
    final int numIters = 2000;
    for (int i = 0; i < numIters; ++i) {
        server.addWorkToOffer(makeInput(i, TimeUnit.MILLISECONDS.toMicros(i), "key", DEFAULT_SHARDING_KEY));
    }
    server.waitForAndGetCommits(numIters);
    worker.stop();
    verify(hotKeyLogger, atLeastOnce()).logHotKeyDetection(nullable(String.class), any(), eq("key"));
}
Also used : ParallelInstruction(com.google.api.services.dataflow.model.ParallelInstruction) StreamingConfigTask(com.google.api.services.dataflow.model.StreamingConfigTask) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) Structs.addString(org.apache.beam.runners.dataflow.util.Structs.addString) WorkItem(com.google.api.services.dataflow.model.WorkItem) StreamingDataflowWorkerOptions(org.apache.beam.runners.dataflow.worker.options.StreamingDataflowWorkerOptions) Test(org.junit.Test)

Example 32 with Server

use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server in project beam by apache.

the class StreamingDataflowWorkerTest method testHotKeyLoggingNotEnabled.

@Test
public void testHotKeyLoggingNotEnabled() throws Exception {
    // This is to test that the worker can correctly log the key from a hot key.
    List<ParallelInstruction> instructions = Arrays.asList(makeSourceInstruction(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())), makeSinkInstruction(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()), 0));
    FakeWindmillServer server = new FakeWindmillServer(errorCollector);
    server.setIsReady(false);
    StreamingConfigTask streamingConfig = new StreamingConfigTask();
    streamingConfig.setStreamingComputationConfigs(ImmutableList.of(makeDefaultStreamingComputationConfig(instructions)));
    streamingConfig.setWindmillServiceEndpoint("foo");
    WorkItem workItem = new WorkItem();
    workItem.setStreamingConfigTask(streamingConfig);
    when(mockWorkUnitClient.getGlobalStreamingConfigWorkItem()).thenReturn(Optional.of(workItem));
    StreamingDataflowWorkerOptions options = createTestingPipelineOptions(server);
    StreamingDataflowWorker worker = makeWorker(instructions, options, true);
    worker.start();
    final int numIters = 2000;
    for (int i = 0; i < numIters; ++i) {
        server.addWorkToOffer(makeInput(i, TimeUnit.MILLISECONDS.toMicros(i), "key", DEFAULT_SHARDING_KEY));
    }
    server.waitForAndGetCommits(numIters);
    worker.stop();
    verify(hotKeyLogger, atLeastOnce()).logHotKeyDetection(nullable(String.class), any());
}
Also used : ParallelInstruction(com.google.api.services.dataflow.model.ParallelInstruction) StreamingConfigTask(com.google.api.services.dataflow.model.StreamingConfigTask) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) Structs.addString(org.apache.beam.runners.dataflow.util.Structs.addString) WorkItem(com.google.api.services.dataflow.model.WorkItem) StreamingDataflowWorkerOptions(org.apache.beam.runners.dataflow.worker.options.StreamingDataflowWorkerOptions) Test(org.junit.Test)

Example 33 with Server

use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server in project beam by apache.

the class StreamingDataflowWorkerTest method testExceptionInvalidatesCache.

@Test
public void testExceptionInvalidatesCache() throws Exception {
    // We'll need to force the system to limit bundles to one message at a time.
    // Sequence is as follows:
    // 01. GetWork[0] (token 0)
    // 02. Create counter reader
    // 03. Counter yields 0
    // 04. GetData[0] (state as null)
    // 05. Read state as null
    // 06. Set state as 42
    // 07. THROW on taking counter reader checkpoint
    // 08. Create counter reader
    // 09. Counter yields 0
    // 10. GetData[1] (state as null)
    // 11. Read state as null (*** not 42 ***)
    // 12. Take counter reader checkpoint as 0
    // 13. CommitWork[0] (message 0:0, state 42, checkpoint 0)
    // 14. GetWork[1] (token 1, checkpoint as 0)
    // 15. Counter yields 1
    // 16. Read (cached) state as 42
    // 17. Take counter reader checkpoint 1
    // 18. CommitWork[1] (message 0:1, checkpoint 1)
    // 19. GetWork[2] (token 2, checkpoint as 1)
    // 20. Counter yields 2
    // 21. THROW on processElement
    // 22. Recreate reader from checkpoint 1
    // 23. Counter yields 2 (*** not eof ***)
    // 24. GetData[2] (state as 42)
    // 25. Read state as 42
    // 26. Take counter reader checkpoint 2
    // 27. CommitWork[2] (message 0:2, checkpoint 2)
    FakeWindmillServer server = new FakeWindmillServer(errorCollector);
    server.setExpectedExceptionCount(2);
    DataflowPipelineOptions options = createTestingPipelineOptions(server);
    options.setNumWorkers(1);
    DataflowPipelineDebugOptions debugOptions = options.as(DataflowPipelineDebugOptions.class);
    debugOptions.setUnboundedReaderMaxElements(1);
    CloudObject codec = CloudObjects.asCloudObject(WindowedValue.getFullCoder(ValueWithRecordId.ValueWithRecordIdCoder.of(KvCoder.of(VarIntCoder.of(), VarIntCoder.of())), GlobalWindow.Coder.INSTANCE), /*sdkComponents=*/
    null);
    TestCountingSource counter = new TestCountingSource(3).withThrowOnFirstSnapshot(true);
    List<ParallelInstruction> instructions = Arrays.asList(new ParallelInstruction().setOriginalName("OriginalReadName").setSystemName("Read").setName(DEFAULT_PARDO_USER_NAME).setRead(new ReadInstruction().setSource(CustomSources.serializeToCloudSource(counter, options).setCodec(codec))).setOutputs(Arrays.asList(new InstructionOutput().setName("read_output").setOriginalName(DEFAULT_OUTPUT_ORIGINAL_NAME).setSystemName(DEFAULT_OUTPUT_SYSTEM_NAME).setCodec(codec))), makeDoFnInstruction(new TestExceptionInvalidatesCacheFn(), 0, StringUtf8Coder.of(), WindowingStrategy.globalDefault()), makeSinkInstruction(StringUtf8Coder.of(), 1, GlobalWindow.Coder.INSTANCE));
    StreamingDataflowWorker worker = makeWorker(instructions, options.as(StreamingDataflowWorkerOptions.class), true);
    worker.setRetryLocallyDelayMs(100);
    worker.start();
    // Three GetData requests
    for (int i = 0; i < 3; i++) {
        ByteString state;
        if (i == 0 || i == 1) {
            state = ByteString.EMPTY;
        } else {
            state = ByteString.copyFrom(new byte[] { 42 });
        }
        Windmill.GetDataResponse.Builder dataResponse = Windmill.GetDataResponse.newBuilder();
        dataResponse.addDataBuilder().setComputationId(DEFAULT_COMPUTATION_ID).addDataBuilder().setKey(ByteString.copyFromUtf8("0000000000000001")).setShardingKey(1).addValuesBuilder().setTag(ByteString.copyFromUtf8("//+uint")).setStateFamily(DEFAULT_PARDO_STATE_FAMILY).getValueBuilder().setTimestamp(0).setData(state);
        server.addDataToOffer(dataResponse.build());
    }
    // Three GetWork requests and commits
    for (int i = 0; i < 3; i++) {
        StringBuilder sb = new StringBuilder();
        sb.append("work {\n");
        sb.append("  computation_id: \"computation\"\n");
        sb.append("  input_data_watermark: 0\n");
        sb.append("  work {\n");
        sb.append("    key: \"0000000000000001\"\n");
        sb.append("    sharding_key: 1\n");
        sb.append("    work_token: ");
        sb.append(i);
        sb.append("    cache_token: 1");
        sb.append("\n");
        if (i > 0) {
            int previousCheckpoint = i - 1;
            sb.append("    source_state {\n");
            sb.append("      state: \"");
            sb.append((char) previousCheckpoint);
            sb.append("\"\n");
            // We'll elide the finalize ids since it's not necessary to trigger the finalizer
            // for this test.
            sb.append("    }\n");
        }
        sb.append("  }\n");
        sb.append("}\n");
        server.addWorkToOffer(buildInput(sb.toString(), null));
        Map<Long, Windmill.WorkItemCommitRequest> result = server.waitForAndGetCommits(1);
        Windmill.WorkItemCommitRequest commit = result.get((long) i);
        UnsignedLong finalizeId = UnsignedLong.fromLongBits(commit.getSourceStateUpdates().getFinalizeIds(0));
        sb = new StringBuilder();
        sb.append("key: \"0000000000000001\"\n");
        sb.append("sharding_key: 1\n");
        sb.append("work_token: ");
        sb.append(i);
        sb.append("\n");
        sb.append("cache_token: 1\n");
        sb.append("output_messages {\n");
        sb.append("  destination_stream_id: \"out\"\n");
        sb.append("  bundles {\n");
        sb.append("    key: \"0000000000000001\"\n");
        int messageNum = i;
        sb.append("    messages {\n");
        sb.append("      timestamp: ");
        sb.append(messageNum * 1000);
        sb.append("\n");
        sb.append("      data: \"0:");
        sb.append(messageNum);
        sb.append("\"\n");
        sb.append("    }\n");
        sb.append("    messages_ids: \"\"\n");
        sb.append("  }\n");
        sb.append("}\n");
        if (i == 0) {
            sb.append("value_updates {\n");
            sb.append("  tag: \"//+uint\"\n");
            sb.append("  value {\n");
            sb.append("    timestamp: 0\n");
            sb.append("    data: \"");
            sb.append((char) 42);
            sb.append("\"\n");
            sb.append("  }\n");
            sb.append("  state_family: \"parDoStateFamily\"\n");
            sb.append("}\n");
        }
        int sourceState = i;
        sb.append("source_state_updates {\n");
        sb.append("  state: \"");
        sb.append((char) sourceState);
        sb.append("\"\n");
        sb.append("  finalize_ids: ");
        sb.append(finalizeId);
        sb.append("}\n");
        sb.append("source_watermark: ");
        sb.append((sourceState + 1) * 1000);
        sb.append("\n");
        sb.append("source_backlog_bytes: 7\n");
        assertThat(// for the current test.
        setValuesTimestamps(commit.toBuilder().clearOutputTimers()).build(), equalTo(setMessagesMetadata(PaneInfo.NO_FIRING, CoderUtils.encodeToByteArray(CollectionCoder.of(GlobalWindow.Coder.INSTANCE), ImmutableList.of(GlobalWindow.INSTANCE)), parseCommitRequest(sb.toString())).build()));
    }
}
Also used : DataflowPipelineOptions(org.apache.beam.runners.dataflow.options.DataflowPipelineOptions) UnsignedLong(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.primitives.UnsignedLong) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) InstructionOutput(com.google.api.services.dataflow.model.InstructionOutput) ReadInstruction(com.google.api.services.dataflow.model.ReadInstruction) WorkItemCommitRequest(org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItemCommitRequest) StreamingDataflowWorkerOptions(org.apache.beam.runners.dataflow.worker.options.StreamingDataflowWorkerOptions) ParallelInstruction(com.google.api.services.dataflow.model.ParallelInstruction) CloudObject(org.apache.beam.runners.dataflow.util.CloudObject) WorkItemCommitRequest(org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItemCommitRequest) TestCountingSource(org.apache.beam.runners.dataflow.worker.testing.TestCountingSource) AtomicLong(java.util.concurrent.atomic.AtomicLong) DataflowCounterUpdateExtractor.splitIntToLong(org.apache.beam.runners.dataflow.worker.counters.DataflowCounterUpdateExtractor.splitIntToLong) UnsignedLong(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.primitives.UnsignedLong) ComputationGetDataResponse(org.apache.beam.runners.dataflow.worker.windmill.Windmill.ComputationGetDataResponse) KeyedGetDataResponse(org.apache.beam.runners.dataflow.worker.windmill.Windmill.KeyedGetDataResponse) GetDataResponse(org.apache.beam.runners.dataflow.worker.windmill.Windmill.GetDataResponse) Windmill(org.apache.beam.runners.dataflow.worker.windmill.Windmill) DataflowPipelineDebugOptions(org.apache.beam.runners.dataflow.options.DataflowPipelineDebugOptions) Test(org.junit.Test)

Example 34 with Server

use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server in project beam by apache.

the class StateFetcherTest method testFetchGlobalDataBasic.

@Test
public void testFetchGlobalDataBasic() throws Exception {
    StateFetcher fetcher = new StateFetcher(server);
    ByteString.Output stream = ByteString.newOutput();
    ListCoder.of(StringUtf8Coder.of()).encode(Arrays.asList("data"), stream, Coder.Context.OUTER);
    ByteString encodedIterable = stream.toByteString();
    PCollectionView<String> view = TestPipeline.create().apply(Create.empty(StringUtf8Coder.of())).apply(View.asSingleton());
    String tag = view.getTagInternal().getId();
    // Test three calls in a row. First, data is not ready, then data is ready,
    // then the data is already cached.
    when(server.getSideInputData(any(Windmill.GlobalDataRequest.class))).thenReturn(buildGlobalDataResponse(tag, ByteString.EMPTY, false, null), buildGlobalDataResponse(tag, ByteString.EMPTY, true, encodedIterable));
    assertEquals(null, fetcher.fetchSideInput(view, GlobalWindow.INSTANCE, STATE_FAMILY, SideInputState.UNKNOWN, readStateSupplier));
    assertEquals(null, fetcher.fetchSideInput(view, GlobalWindow.INSTANCE, STATE_FAMILY, SideInputState.UNKNOWN, readStateSupplier));
    assertEquals("data", fetcher.fetchSideInput(view, GlobalWindow.INSTANCE, STATE_FAMILY, SideInputState.KNOWN_READY, readStateSupplier).orNull());
    assertEquals("data", fetcher.fetchSideInput(view, GlobalWindow.INSTANCE, STATE_FAMILY, SideInputState.KNOWN_READY, readStateSupplier).orNull());
    verify(server, times(2)).getSideInputData(buildGlobalDataRequest(tag, ByteString.EMPTY));
    verifyNoMoreInteractions(server);
}
Also used : ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) Test(org.junit.Test)

Example 35 with Server

use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server in project beam by apache.

the class StateFetcherTest method testFetchGlobalDataNull.

@Test
public void testFetchGlobalDataNull() throws Exception {
    StateFetcher fetcher = new StateFetcher(server);
    ByteString.Output stream = ByteString.newOutput();
    ListCoder.of(VoidCoder.of()).encode(Arrays.asList((Void) null), stream, Coder.Context.OUTER);
    ByteString encodedIterable = stream.toByteString();
    PCollectionView<Void> view = TestPipeline.create().apply(Create.empty(VoidCoder.of())).apply(View.asSingleton());
    String tag = view.getTagInternal().getId();
    // Test three calls in a row. First, data is not ready, then data is ready,
    // then the data is already cached.
    when(server.getSideInputData(any(Windmill.GlobalDataRequest.class))).thenReturn(buildGlobalDataResponse(tag, ByteString.EMPTY, false, null), buildGlobalDataResponse(tag, ByteString.EMPTY, true, encodedIterable));
    assertEquals(null, fetcher.fetchSideInput(view, GlobalWindow.INSTANCE, STATE_FAMILY, SideInputState.UNKNOWN, readStateSupplier));
    assertEquals(null, fetcher.fetchSideInput(view, GlobalWindow.INSTANCE, STATE_FAMILY, SideInputState.UNKNOWN, readStateSupplier));
    assertEquals(null, fetcher.fetchSideInput(view, GlobalWindow.INSTANCE, STATE_FAMILY, SideInputState.KNOWN_READY, readStateSupplier).orNull());
    assertEquals(null, fetcher.fetchSideInput(view, GlobalWindow.INSTANCE, STATE_FAMILY, SideInputState.KNOWN_READY, readStateSupplier).orNull());
    verify(server, times(2)).getSideInputData(buildGlobalDataRequest(tag, ByteString.EMPTY));
    verifyNoMoreInteractions(server);
}
Also used : ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) Test(org.junit.Test)

Aggregations

Test (org.junit.Test)76 Server (io.grpc.Server)68 IOException (java.io.IOException)27 ByteString (org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString)24 ArrayList (java.util.ArrayList)21 StreamObserver (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver)21 ExecutionException (java.util.concurrent.ExecutionException)20 TimeoutException (java.util.concurrent.TimeoutException)20 CountDownLatch (java.util.concurrent.CountDownLatch)19 ManagedChannel (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ManagedChannel)19 Server (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server)18 StatusException (io.grpc.StatusException)17 Metadata (io.grpc.Metadata)12 List (java.util.List)12 BeamFnApi (org.apache.beam.model.fnexecution.v1.BeamFnApi)12 FilterChain (io.grpc.xds.EnvoyServerProtoData.FilterChain)11 ExecutorService (java.util.concurrent.ExecutorService)11 ParallelInstruction (com.google.api.services.dataflow.model.ParallelInstruction)10 ServerCall (io.grpc.ServerCall)10 StreamObserver (io.grpc.stub.StreamObserver)10