use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.CallStreamObserver in project beam by apache.
the class ServerFactoryTest method runTestUsing.
private Endpoints.ApiServiceDescriptor runTestUsing(ServerFactory serverFactory, ManagedChannelFactory channelFactory) throws Exception {
Endpoints.ApiServiceDescriptor.Builder apiServiceDescriptorBuilder = Endpoints.ApiServiceDescriptor.newBuilder();
final Collection<Elements> serverElements = new ArrayList<>();
final CountDownLatch clientHangedUp = new CountDownLatch(1);
CallStreamObserver<Elements> serverInboundObserver = TestStreams.withOnNext(serverElements::add).withOnCompleted(clientHangedUp::countDown).build();
TestDataService service = new TestDataService(serverInboundObserver);
Server server = serverFactory.allocateAddressAndCreate(ImmutableList.of(service), apiServiceDescriptorBuilder);
assertFalse(server.isShutdown());
ManagedChannel channel = channelFactory.forDescriptor(apiServiceDescriptorBuilder.build());
BeamFnDataGrpc.BeamFnDataStub stub = BeamFnDataGrpc.newStub(channel);
final Collection<BeamFnApi.Elements> clientElements = new ArrayList<>();
final CountDownLatch serverHangedUp = new CountDownLatch(1);
CallStreamObserver<BeamFnApi.Elements> clientInboundObserver = TestStreams.withOnNext(clientElements::add).withOnCompleted(serverHangedUp::countDown).build();
StreamObserver<Elements> clientOutboundObserver = stub.data(clientInboundObserver);
StreamObserver<BeamFnApi.Elements> serverOutboundObserver = service.outboundObservers.take();
clientOutboundObserver.onNext(CLIENT_DATA);
serverOutboundObserver.onNext(SERVER_DATA);
clientOutboundObserver.onCompleted();
clientHangedUp.await();
serverOutboundObserver.onCompleted();
serverHangedUp.await();
assertThat(clientElements, contains(SERVER_DATA));
assertThat(serverElements, contains(CLIENT_DATA));
return apiServiceDescriptorBuilder.build();
}
use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.CallStreamObserver in project beam by apache.
the class BeamFnLoggingClientTest method testLogging.
@Test
public void testLogging() throws Exception {
BeamFnLoggingMDC.setInstructionId("instruction-1");
AtomicBoolean clientClosedStream = new AtomicBoolean();
Collection<BeamFnApi.LogEntry> values = new ConcurrentLinkedQueue<>();
AtomicReference<StreamObserver<BeamFnApi.LogControl>> outboundServerObserver = new AtomicReference<>();
CallStreamObserver<BeamFnApi.LogEntry.List> inboundServerObserver = TestStreams.withOnNext((BeamFnApi.LogEntry.List logEntries) -> values.addAll(logEntries.getLogEntriesList())).withOnCompleted(() -> {
// Remember that the client told us that this stream completed
clientClosedStream.set(true);
outboundServerObserver.get().onCompleted();
}).build();
Endpoints.ApiServiceDescriptor apiServiceDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl(this.getClass().getName() + "-" + UUID.randomUUID().toString()).build();
Server server = InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()).addService(new BeamFnLoggingGrpc.BeamFnLoggingImplBase() {
@Override
public StreamObserver<BeamFnApi.LogEntry.List> logging(StreamObserver<BeamFnApi.LogControl> outboundObserver) {
outboundServerObserver.set(outboundObserver);
return inboundServerObserver;
}
}).build();
server.start();
ManagedChannel channel = InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build();
try {
BeamFnLoggingClient client = new BeamFnLoggingClient(PipelineOptionsFactory.fromArgs(new String[] { "--defaultSdkHarnessLogLevel=OFF", "--sdkHarnessLogLevelOverrides={\"ConfiguredLogger\": \"DEBUG\"}" }).create(), apiServiceDescriptor, (Endpoints.ApiServiceDescriptor descriptor) -> channel);
// Keep a strong reference to the loggers in this block. Otherwise the call to client.close()
// removes the only reference and the logger may get GC'd before the assertions (BEAM-4136).
Logger rootLogger = LogManager.getLogManager().getLogger("");
Logger configuredLogger = LogManager.getLogManager().getLogger("ConfiguredLogger");
// Ensure that log levels were correctly set.
assertEquals(Level.OFF, rootLogger.getLevel());
assertEquals(Level.FINE, configuredLogger.getLevel());
// Should be filtered because the default log level override is OFF
rootLogger.log(FILTERED_RECORD);
// Should not be filtered because the default log level override for ConfiguredLogger is DEBUG
configuredLogger.log(TEST_RECORD);
configuredLogger.log(TEST_RECORD_WITH_EXCEPTION);
client.close();
// Verify that after close, log levels are reset.
assertEquals(Level.INFO, rootLogger.getLevel());
assertNull(configuredLogger.getLevel());
assertTrue(clientClosedStream.get());
assertTrue(channel.isShutdown());
assertThat(values, contains(TEST_ENTRY, TEST_ENTRY_WITH_EXCEPTION));
} finally {
server.shutdownNow();
}
}
use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.CallStreamObserver in project beam by apache.
the class BeamFnLoggingClientTest method testWhenServerHangsUpEarlyThatClientIsAbleCleanup.
@Test
public void testWhenServerHangsUpEarlyThatClientIsAbleCleanup() throws Exception {
BeamFnLoggingMDC.setInstructionId("instruction-1");
Collection<BeamFnApi.LogEntry> values = new ConcurrentLinkedQueue<>();
AtomicReference<StreamObserver<BeamFnApi.LogControl>> outboundServerObserver = new AtomicReference<>();
CallStreamObserver<BeamFnApi.LogEntry.List> inboundServerObserver = TestStreams.withOnNext((BeamFnApi.LogEntry.List logEntries) -> values.addAll(logEntries.getLogEntriesList())).build();
Endpoints.ApiServiceDescriptor apiServiceDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl(this.getClass().getName() + "-" + UUID.randomUUID().toString()).build();
Server server = InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()).addService(new BeamFnLoggingGrpc.BeamFnLoggingImplBase() {
@Override
public StreamObserver<BeamFnApi.LogEntry.List> logging(StreamObserver<BeamFnApi.LogControl> outboundObserver) {
outboundServerObserver.set(outboundObserver);
outboundObserver.onCompleted();
return inboundServerObserver;
}
}).build();
server.start();
ManagedChannel channel = InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build();
try {
BeamFnLoggingClient client = new BeamFnLoggingClient(PipelineOptionsFactory.fromArgs(new String[] { "--defaultSdkHarnessLogLevel=OFF", "--sdkHarnessLogLevelOverrides={\"ConfiguredLogger\": \"DEBUG\"}" }).create(), apiServiceDescriptor, (Endpoints.ApiServiceDescriptor descriptor) -> channel);
// Keep a strong reference to the loggers in this block. Otherwise the call to client.close()
// removes the only reference and the logger may get GC'd before the assertions (BEAM-4136).
Logger rootLogger = LogManager.getLogManager().getLogger("");
Logger configuredLogger = LogManager.getLogManager().getLogger("ConfiguredLogger");
client.close();
// Verify that after close, log levels are reset.
assertEquals(Level.INFO, rootLogger.getLevel());
assertNull(configuredLogger.getLevel());
} finally {
assertTrue(channel.isShutdown());
server.shutdownNow();
}
}
use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.CallStreamObserver in project beam by apache.
the class BeamFnControlClientTest method testDelegation.
@Test
public void testDelegation() throws Exception {
AtomicBoolean clientClosedStream = new AtomicBoolean();
BlockingQueue<BeamFnApi.InstructionResponse> values = new LinkedBlockingQueue<>();
BlockingQueue<StreamObserver<BeamFnApi.InstructionRequest>> outboundServerObservers = new LinkedBlockingQueue<>();
CallStreamObserver<BeamFnApi.InstructionResponse> inboundServerObserver = TestStreams.withOnNext(values::add).withOnCompleted(() -> clientClosedStream.set(true)).build();
Endpoints.ApiServiceDescriptor apiServiceDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl(this.getClass().getName() + "-" + UUID.randomUUID().toString()).build();
Server server = InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()).addService(new BeamFnControlGrpc.BeamFnControlImplBase() {
@Override
public StreamObserver<BeamFnApi.InstructionResponse> control(StreamObserver<BeamFnApi.InstructionRequest> outboundObserver) {
Uninterruptibles.putUninterruptibly(outboundServerObservers, outboundObserver);
return inboundServerObserver;
}
}).build();
server.start();
try {
EnumMap<BeamFnApi.InstructionRequest.RequestCase, ThrowingFunction<BeamFnApi.InstructionRequest, BeamFnApi.InstructionResponse.Builder>> handlers = new EnumMap<>(BeamFnApi.InstructionRequest.RequestCase.class);
handlers.put(BeamFnApi.InstructionRequest.RequestCase.PROCESS_BUNDLE, value -> {
assertEquals(value.getInstructionId(), BeamFnLoggingMDC.getInstructionId());
return BeamFnApi.InstructionResponse.newBuilder().setProcessBundle(BeamFnApi.ProcessBundleResponse.getDefaultInstance());
});
handlers.put(BeamFnApi.InstructionRequest.RequestCase.REGISTER, value -> {
assertEquals(value.getInstructionId(), BeamFnLoggingMDC.getInstructionId());
throw FAILURE;
});
ExecutorService executor = Executors.newCachedThreadPool();
BeamFnControlClient client = new BeamFnControlClient(apiServiceDescriptor, ManagedChannelFactory.createInProcess(), OutboundObserverFactory.trivial(), executor, handlers);
// Get the connected client and attempt to send and receive an instruction
StreamObserver<BeamFnApi.InstructionRequest> outboundServerObserver = outboundServerObservers.take();
outboundServerObserver.onNext(SUCCESSFUL_REQUEST);
assertEquals(SUCCESSFUL_RESPONSE, values.take());
// Ensure that conversion of an unknown request type is properly converted to a
// failure response.
outboundServerObserver.onNext(UNKNOWN_HANDLER_REQUEST);
assertEquals(UNKNOWN_HANDLER_RESPONSE, values.take());
// Ensure that all exceptions are caught and translated to failures
outboundServerObserver.onNext(FAILURE_REQUEST);
assertEquals(FAILURE_RESPONSE, values.take());
// Ensure that the server completing the stream translates to the completable future
// being completed allowing for a successful shutdown of the client.
outboundServerObserver.onCompleted();
client.waitForTermination();
} finally {
server.shutdownNow();
}
}
use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.CallStreamObserver in project beam by apache.
the class BeamFnControlClientTest method testJavaErrorResponse.
@Test
public void testJavaErrorResponse() throws Exception {
BlockingQueue<StreamObserver<BeamFnApi.InstructionRequest>> outboundServerObservers = new LinkedBlockingQueue<>();
BlockingQueue<Throwable> error = new LinkedBlockingQueue<>();
CallStreamObserver<BeamFnApi.InstructionResponse> inboundServerObserver = TestStreams.<BeamFnApi.InstructionResponse>withOnNext(response -> fail(String.format("Unexpected Response %s", response))).withOnError(error::add).build();
Endpoints.ApiServiceDescriptor apiServiceDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl(this.getClass().getName() + "-" + UUID.randomUUID().toString()).build();
Server server = InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()).addService(new BeamFnControlGrpc.BeamFnControlImplBase() {
@Override
public StreamObserver<BeamFnApi.InstructionResponse> control(StreamObserver<BeamFnApi.InstructionRequest> outboundObserver) {
Uninterruptibles.putUninterruptibly(outboundServerObservers, outboundObserver);
return inboundServerObserver;
}
}).build();
server.start();
try {
EnumMap<BeamFnApi.InstructionRequest.RequestCase, ThrowingFunction<BeamFnApi.InstructionRequest, BeamFnApi.InstructionResponse.Builder>> handlers = new EnumMap<>(BeamFnApi.InstructionRequest.RequestCase.class);
handlers.put(BeamFnApi.InstructionRequest.RequestCase.REGISTER, value -> {
assertEquals(value.getInstructionId(), BeamFnLoggingMDC.getInstructionId());
throw new Error("Test Error");
});
ExecutorService executor = Executors.newCachedThreadPool();
BeamFnControlClient client = new BeamFnControlClient(apiServiceDescriptor, ManagedChannelFactory.createInProcess(), OutboundObserverFactory.trivial(), executor, handlers);
// Get the connected client and attempt to send and receive an instruction
StreamObserver<BeamFnApi.InstructionRequest> outboundServerObserver = outboundServerObservers.take();
// Ensure that all exceptions are caught and translated to failures
outboundServerObserver.onNext(InstructionRequest.newBuilder().setInstructionId("0").setRegister(RegisterRequest.getDefaultInstance()).build());
// There should be an error reported to the StreamObserver.
assertThat(error.take(), not(nullValue()));
// Ensure that the client shuts down when an Error is thrown from the harness
try {
client.waitForTermination();
throw new IllegalStateException("The future should have terminated with an error");
} catch (ExecutionException errorWrapper) {
assertThat(errorWrapper.getCause().getMessage(), containsString("Test Error"));
}
} finally {
server.shutdownNow();
}
}
Aggregations