content: /*
   * 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);
      }
    }
  }
  