use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server in project beam by apache.
the class BeamFnDataGrpcClientTest method testForInboundConsumer.
@Test
public void testForInboundConsumer() throws Exception {
CountDownLatch waitForClientToConnect = new CountDownLatch(1);
Collection<WindowedValue<String>> inboundValuesA = new ConcurrentLinkedQueue<>();
Collection<WindowedValue<String>> inboundValuesB = new ConcurrentLinkedQueue<>();
Collection<BeamFnApi.Elements> inboundServerValues = new ConcurrentLinkedQueue<>();
AtomicReference<StreamObserver<BeamFnApi.Elements>> outboundServerObserver = new AtomicReference<>();
CallStreamObserver<BeamFnApi.Elements> inboundServerObserver = TestStreams.withOnNext(inboundServerValues::add).build();
Endpoints.ApiServiceDescriptor apiServiceDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl(this.getClass().getName() + "-" + UUID.randomUUID()).build();
Server server = InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()).addService(new BeamFnDataGrpc.BeamFnDataImplBase() {
@Override
public StreamObserver<BeamFnApi.Elements> data(StreamObserver<BeamFnApi.Elements> outboundObserver) {
outboundServerObserver.set(outboundObserver);
waitForClientToConnect.countDown();
return inboundServerObserver;
}
}).build();
server.start();
try {
ManagedChannel channel = InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build();
BeamFnDataGrpcClient clientFactory = new BeamFnDataGrpcClient(PipelineOptionsFactory.create(), (Endpoints.ApiServiceDescriptor descriptor) -> channel, OutboundObserverFactory.trivial());
BeamFnDataInboundObserver2 observerA = BeamFnDataInboundObserver2.forConsumers(Arrays.asList(DataEndpoint.create(TRANSFORM_ID_A, CODER, inboundValuesA::add)), Collections.emptyList());
BeamFnDataInboundObserver2 observerB = BeamFnDataInboundObserver2.forConsumers(Arrays.asList(DataEndpoint.create(TRANSFORM_ID_B, CODER, inboundValuesB::add)), Collections.emptyList());
clientFactory.registerReceiver(INSTRUCTION_ID_A, Arrays.asList(apiServiceDescriptor), observerA);
waitForClientToConnect.await();
outboundServerObserver.get().onNext(ELEMENTS_A_1);
// Purposefully transmit some data before the consumer for B is bound showing that
// data is not lost
outboundServerObserver.get().onNext(ELEMENTS_B_1);
Thread.sleep(100);
clientFactory.registerReceiver(INSTRUCTION_ID_B, Arrays.asList(apiServiceDescriptor), observerB);
// Show that out of order stream completion can occur.
observerB.awaitCompletion();
assertThat(inboundValuesB, contains(valueInGlobalWindow("JKL"), valueInGlobalWindow("MNO")));
outboundServerObserver.get().onNext(ELEMENTS_A_2);
observerA.awaitCompletion();
assertThat(inboundValuesA, contains(valueInGlobalWindow("ABC"), valueInGlobalWindow("DEF"), valueInGlobalWindow("GHI")));
} finally {
server.shutdownNow();
}
}
use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server in project beam by apache.
the class BeamFnStateGrpcClientCacheTest method testCachingOfClient.
@Test
public void testCachingOfClient() throws Exception {
Endpoints.ApiServiceDescriptor otherApiServiceDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl(apiServiceDescriptor.getUrl() + "-other").build();
Server testServer2 = InProcessServerBuilder.forName(otherApiServiceDescriptor.getUrl()).addService(new BeamFnStateGrpc.BeamFnStateImplBase() {
@Override
public StreamObserver<StateRequest> state(StreamObserver<StateResponse> outboundObserver) {
throw new IllegalStateException("Unexpected in test.");
}
}).build();
testServer2.start();
try {
assertSame(clientCache.forApiServiceDescriptor(apiServiceDescriptor), clientCache.forApiServiceDescriptor(apiServiceDescriptor));
assertNotSame(clientCache.forApiServiceDescriptor(apiServiceDescriptor), clientCache.forApiServiceDescriptor(otherApiServiceDescriptor));
} finally {
testServer2.shutdownNow();
}
}
use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server in project beam by apache.
the class PubsubGrpcClientTest method pullOneMessage.
@Test
public void pullOneMessage() throws IOException {
initializeClient(null, null);
String expectedSubscription = SUBSCRIPTION.getPath();
final PullRequest expectedRequest = PullRequest.newBuilder().setSubscription(expectedSubscription).setReturnImmediately(true).setMaxMessages(10).build();
Timestamp timestamp = Timestamp.newBuilder().setSeconds(PUB_TIME_MS / 1000).setNanos((int) (PUB_TIME_MS % 1000) * 1000 * 1000).build();
PubsubMessage expectedPubsubMessage = PubsubMessage.newBuilder().setMessageId(MESSAGE_ID).setData(ByteString.copyFrom(DATA.getBytes(StandardCharsets.UTF_8))).setPublishTime(timestamp).build();
ReceivedMessage expectedReceivedMessage = ReceivedMessage.newBuilder().setMessage(expectedPubsubMessage).setAckId(ACK_ID).build();
final PullResponse response = PullResponse.newBuilder().addAllReceivedMessages(ImmutableList.of(expectedReceivedMessage)).build();
final List<PullRequest> requestsReceived = new ArrayList<>();
SubscriberImplBase subscriberImplBase = new SubscriberImplBase() {
@Override
public void pull(PullRequest request, StreamObserver<PullResponse> responseObserver) {
requestsReceived.add(request);
responseObserver.onNext(response);
responseObserver.onCompleted();
}
};
Server server = InProcessServerBuilder.forName(channelName).addService(subscriberImplBase).build().start();
try {
List<IncomingMessage> actualMessages = client.pull(REQ_TIME_MS, SUBSCRIPTION, 10, true);
assertEquals(1, actualMessages.size());
IncomingMessage actualMessage = actualMessages.get(0);
assertEquals(ACK_ID, actualMessage.ackId());
assertEquals(DATA, actualMessage.message().getData().toStringUtf8());
assertEquals(MESSAGE_ID, actualMessage.recordId());
assertEquals(REQ_TIME_MS, actualMessage.requestTimeMsSinceEpoch());
assertEquals(PUB_TIME_MS, actualMessage.timestampMsSinceEpoch());
assertEquals(expectedRequest, Iterables.getOnlyElement(requestsReceived));
} finally {
server.shutdownNow();
}
}
use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server in project beam by apache.
the class PubsubGrpcClientTest method pullOneMessageUsingAttributes.
@Test
public void pullOneMessageUsingAttributes() throws IOException {
initializeClient(TIMESTAMP_ATTRIBUTE, ID_ATTRIBUTE);
String expectedSubscription = SUBSCRIPTION.getPath();
final PullRequest expectedRequest = PullRequest.newBuilder().setSubscription(expectedSubscription).setReturnImmediately(true).setMaxMessages(10).build();
Timestamp timestamp = Timestamp.newBuilder().setSeconds(PUB_TIME_MS / 1000).setNanos((int) (PUB_TIME_MS % 1000) * 1000 * 1000).build();
PubsubMessage expectedPubsubMessage = PubsubMessage.newBuilder().setMessageId(MESSAGE_ID).setData(ByteString.copyFrom(DATA.getBytes(StandardCharsets.UTF_8))).setPublishTime(timestamp).putAllAttributes(ATTRIBUTES).putAllAttributes(ImmutableMap.of(TIMESTAMP_ATTRIBUTE, String.valueOf(MESSAGE_TIME_MS), ID_ATTRIBUTE, RECORD_ID)).build();
ReceivedMessage expectedReceivedMessage = ReceivedMessage.newBuilder().setMessage(expectedPubsubMessage).setAckId(ACK_ID).build();
final PullResponse response = PullResponse.newBuilder().addAllReceivedMessages(ImmutableList.of(expectedReceivedMessage)).build();
final List<PullRequest> requestsReceived = new ArrayList<>();
SubscriberImplBase subscriberImplBase = new SubscriberImplBase() {
@Override
public void pull(PullRequest request, StreamObserver<PullResponse> responseObserver) {
requestsReceived.add(request);
responseObserver.onNext(response);
responseObserver.onCompleted();
}
};
Server server = InProcessServerBuilder.forName(channelName).addService(subscriberImplBase).build().start();
try {
List<IncomingMessage> actualMessages = client.pull(REQ_TIME_MS, SUBSCRIPTION, 10, true);
assertEquals(1, actualMessages.size());
IncomingMessage actualMessage = actualMessages.get(0);
assertEquals(ACK_ID, actualMessage.ackId());
assertEquals(DATA, actualMessage.message().getData().toStringUtf8());
assertEquals(RECORD_ID, actualMessage.recordId());
assertEquals(REQ_TIME_MS, actualMessage.requestTimeMsSinceEpoch());
assertEquals(MESSAGE_TIME_MS, actualMessage.timestampMsSinceEpoch());
assertEquals(expectedRequest, Iterables.getOnlyElement(requestsReceived));
} finally {
server.shutdownNow();
}
}
use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server in project beam by apache.
the class GrpcFnServer method allocatePortAndCreateFor.
/**
* Create {@link GrpcFnServer}s for the provided {@link FnService}s running on an arbitrary port.
*/
public static List<GrpcFnServer<? extends FnService>> allocatePortAndCreateFor(List<? extends FnService> services, ServerFactory factory) throws IOException {
ApiServiceDescriptor.Builder apiServiceDescriptor = ApiServiceDescriptor.newBuilder();
Server server = factory.allocateAddressAndCreate(Collections.unmodifiableList(services), apiServiceDescriptor);
AtomicInteger countdown = new AtomicInteger(services.size());
return Lists.transform(services, service -> new SharedGrpcFnServer<>(server, service, apiServiceDescriptor.build(), countdown));
}
Aggregations