TestServerOptions.java
/*
* 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.arrow.flight;
import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST;
import static org.apache.arrow.flight.Location.forGrpcInsecure;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
import io.grpc.Channel;
import io.grpc.MethodDescriptor;
import io.grpc.ServerServiceDefinition;
import io.grpc.health.v1.HealthCheckRequest;
import io.grpc.health.v1.HealthCheckResponse;
import io.grpc.health.v1.HealthGrpc;
import io.grpc.netty.NettyChannelBuilder;
import io.grpc.netty.NettyServerBuilder;
import io.grpc.protobuf.services.HealthStatusManager;
import java.io.File;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;
import org.apache.arrow.flight.TestBasicOperation.Producer;
import org.apache.arrow.flight.auth.ServerAuthHandler;
import org.apache.arrow.flight.impl.FlightServiceGrpc;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.junit.jupiter.api.Assumptions;
import org.junit.jupiter.api.Test;
public class TestServerOptions {
@Test
public void builderConsumer() throws Exception {
final AtomicBoolean consumerCalled = new AtomicBoolean();
final Consumer<NettyServerBuilder> consumer = (builder) -> consumerCalled.set(true);
try (BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
Producer producer = new Producer(a);
FlightServer s =
FlightServer.builder(a, forGrpcInsecure(LOCALHOST, 0), producer)
.transportHint("grpc.builderConsumer", consumer)
.build()
.start()) {
assertTrue(consumerCalled.get());
}
}
/**
* Make sure that if Flight supplies a default executor to gRPC, then it is closed along with the
* server.
*/
@Test
public void defaultExecutorClosed() throws Exception {
final ExecutorService executor;
try (BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
FlightServer server =
FlightServer.builder(a, forGrpcInsecure(LOCALHOST, 0), new NoOpFlightProducer())
.build()
.start()) {
assertNotNull(server.grpcExecutor);
executor = server.grpcExecutor;
}
assertTrue(executor.isShutdown());
}
/** Make sure that if the user provides an executor to gRPC, then Flight does not close it. */
@Test
public void suppliedExecutorNotClosed() throws Exception {
final ExecutorService executor = Executors.newSingleThreadExecutor();
try {
try (BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
FlightServer server =
FlightServer.builder(a, forGrpcInsecure(LOCALHOST, 0), new NoOpFlightProducer())
.executor(executor)
.build()
.start()) {
assertNull(server.grpcExecutor);
}
assertFalse(executor.isShutdown());
} finally {
executor.shutdown();
}
}
@Test
public void domainSocket() throws Exception {
Assumptions.assumeTrue(
FlightTestUtil.isNativeTransportAvailable(), "We have a native transport available");
final File domainSocket = File.createTempFile("flight-unit-test-", ".sock");
assertTrue(domainSocket.delete());
// Domain socket paths have a platform-dependent limit. Set a conservative limit and skip the
// test if the temporary
// file name is too long. (We do not assume a particular platform-dependent temporary directory
// path.)
Assumptions.assumeTrue(
domainSocket.getAbsolutePath().length() < 100, "The domain socket path is not too long");
final Location location = Location.forGrpcDomainSocket(domainSocket.getAbsolutePath());
try (BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
Producer producer = new Producer(a);
FlightServer s = FlightServer.builder(a, location, producer).build().start(); ) {
try (FlightClient c = FlightClient.builder(a, location).build()) {
try (FlightStream stream = c.getStream(new Ticket(new byte[0]))) {
VectorSchemaRoot root = stream.getRoot();
IntVector iv = (IntVector) root.getVector("c1");
int value = 0;
while (stream.next()) {
for (int i = 0; i < root.getRowCount(); i++) {
assertEquals(value, iv.get(i));
value++;
}
}
}
}
}
}
@Test
public void checkReflectionMetadata() {
// This metadata is needed for gRPC reflection to work.
final ExecutorService executorService = Executors.newSingleThreadExecutor();
try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE)) {
final FlightBindingService service =
new FlightBindingService(
allocator, new NoOpFlightProducer(), ServerAuthHandler.NO_OP, executorService);
final ServerServiceDefinition definition = service.bindService();
assertEquals(
FlightServiceGrpc.getServiceDescriptor().getSchemaDescriptor(),
definition.getServiceDescriptor().getSchemaDescriptor());
final Map<String, MethodDescriptor<?, ?>> definedMethods = new HashMap<>();
final Map<String, MethodDescriptor<?, ?>> serviceMethods = new HashMap<>();
// Make sure that the reflection metadata object is identical across all the places where it's
// accessible
definition
.getMethods()
.forEach(
method ->
definedMethods.put(
method.getMethodDescriptor().getFullMethodName(),
method.getMethodDescriptor()));
definition
.getServiceDescriptor()
.getMethods()
.forEach(method -> serviceMethods.put(method.getFullMethodName(), method));
for (final MethodDescriptor<?, ?> descriptor :
FlightServiceGrpc.getServiceDescriptor().getMethods()) {
final String methodName = descriptor.getFullMethodName();
assertTrue(
definedMethods.containsKey(methodName),
"Method is missing from ServerServiceDefinition: " + methodName);
assertTrue(
definedMethods.containsKey(methodName),
"Method is missing from ServiceDescriptor: " + methodName);
assertEquals(
descriptor.getSchemaDescriptor(), definedMethods.get(methodName).getSchemaDescriptor());
assertEquals(
descriptor.getSchemaDescriptor(), serviceMethods.get(methodName).getSchemaDescriptor());
}
} finally {
executorService.shutdown();
}
}
/*
* This is an extension of builderConsumer test.
* Test that Flight interceptors don't break other registered services
*/
@Test
public void addHealthCheckService() throws Exception {
final HealthStatusManager statusManager = new HealthStatusManager();
final Consumer<NettyServerBuilder> consumer =
(builder) -> {
builder.addService(statusManager.getHealthService());
};
final Location location = forGrpcInsecure(LOCALHOST, 5555);
try (BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
Producer producer = new Producer(a);
FlightServer s =
FlightServer.builder(a, location, producer)
.transportHint("grpc.builderConsumer", consumer)
.build()
.start(); ) {
Channel channel =
NettyChannelBuilder.forAddress(location.toSocketAddress()).usePlaintext().build();
HealthCheckResponse response =
HealthGrpc.newBlockingStub(channel).check(HealthCheckRequest.getDefaultInstance());
assertEquals(response.getStatus(), HealthCheckResponse.ServingStatus.SERVING);
}
}
}