/*
 * 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.assertTrue;

import com.google.common.collect.ImmutableList;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Function;
import org.apache.arrow.flight.perf.PerformanceTestServer;
import org.apache.arrow.flight.perf.TestPerf;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.types.Types.MinorType;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.Schema;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;

public class TestBackPressure {

  private static final int BATCH_SIZE = 4095;

  /** Make sure that failing to consume one stream doesn't block other streams. */
  @Disabled
  @Test
  public void ensureIndependentSteams() throws Exception {
    ensureIndependentSteams((b) -> (location -> new PerformanceTestServer(b, location)));
  }

  /** Make sure that failing to consume one stream doesn't block other streams. */
  @Disabled
  @Test
  public void ensureIndependentSteamsWithCallbacks() throws Exception {
    ensureIndependentSteams(
        (b) ->
            (location ->
                new PerformanceTestServer(
                    b, location, new BackpressureStrategy.CallbackBackpressureStrategy(), true)));
  }

  /** Test to make sure stream doesn't go faster than the consumer is consuming. */
  @Disabled
  @Test
  public void ensureWaitUntilProceed() throws Exception {
    ensureWaitUntilProceed(new PollingBackpressureStrategy(), false);
  }

  /**
   * Test to make sure stream doesn't go faster than the consumer is consuming using a
   * callback-based backpressure strategy.
   */
  @Disabled
  @Test
  public void ensureWaitUntilProceedWithCallbacks() throws Exception {
    ensureWaitUntilProceed(new RecordingCallbackBackpressureStrategy(), true);
  }

  /** Make sure that failing to consume one stream doesn't block other streams. */
  private static void ensureIndependentSteams(
      Function<BufferAllocator, Function<Location, PerformanceTestServer>> serverConstructor)
      throws Exception {
    try (final BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
        final PerformanceTestServer server =
            serverConstructor.apply(a).apply(forGrpcInsecure(LOCALHOST, 0)).start();
        final FlightClient client = FlightClient.builder(a, server.getLocation()).build()) {
      try (FlightStream fs1 =
          client.getStream(
              client
                  .getInfo(TestPerf.getPerfFlightDescriptor(110L * BATCH_SIZE, BATCH_SIZE, 1))
                  .getEndpoints()
                  .get(0)
                  .getTicket())) {
        consume(fs1, 10);

        // stop consuming fs1 but make sure we can consume a large amount of fs2.
        try (FlightStream fs2 =
            client.getStream(
                client
                    .getInfo(TestPerf.getPerfFlightDescriptor(200L * BATCH_SIZE, BATCH_SIZE, 1))
                    .getEndpoints()
                    .get(0)
                    .getTicket())) {
          consume(fs2, 100);

          consume(fs1, 100);
          consume(fs2, 100);

          consume(fs1);
          consume(fs2);
        }
      }
    }
  }

  /** Make sure that a stream doesn't go faster than the consumer is consuming. */
  private static void ensureWaitUntilProceed(
      SleepTimeRecordingBackpressureStrategy bpStrategy, boolean isNonBlocking) throws Exception {
    // request some values.
    final long wait = 3000;
    final long epsilon = 1000;

    try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) {

      final FlightProducer producer =
          new NoOpFlightProducer() {

            @Override
            public void getStream(
                CallContext context, Ticket ticket, ServerStreamListener listener) {
              bpStrategy.register(listener);
              final Runnable loadData =
                  () -> {
                    int batches = 0;
                    final Schema pojoSchema =
                        new Schema(
                            ImmutableList.of(Field.nullable("a", MinorType.BIGINT.getType())));
                    try (VectorSchemaRoot root = VectorSchemaRoot.create(pojoSchema, allocator)) {
                      listener.start(root);
                      while (true) {
                        bpStrategy.waitForListener(0);
                        if (batches > 100) {
                          root.clear();
                          listener.completed();
                          return;
                        }

                        root.allocateNew();
                        root.setRowCount(4095);
                        listener.putNext();
                        batches++;
                      }
                    }
                  };

              if (!isNonBlocking) {
                loadData.run();
              } else {
                final ExecutorService service = Executors.newSingleThreadExecutor();
                Future<?> unused = service.submit(loadData);
                service.shutdown();
              }
            }
          };

      try (BufferAllocator serverAllocator =
              allocator.newChildAllocator("server", 0, Long.MAX_VALUE);
          FlightServer server =
              FlightServer.builder(serverAllocator, forGrpcInsecure(LOCALHOST, 0), producer)
                  .build()
                  .start();
          BufferAllocator clientAllocator =
              allocator.newChildAllocator("client", 0, Long.MAX_VALUE);
          FlightClient client =
              FlightClient.builder(clientAllocator, server.getLocation()).build();
          FlightStream stream = client.getStream(new Ticket(new byte[1]))) {
        VectorSchemaRoot root = stream.getRoot();
        root.clear();
        Thread.sleep(wait);
        while (stream.next()) {
          root.clear();
        }
        long expected = wait - epsilon;
        assertTrue(
            bpStrategy.getSleepTime() > expected,
            String.format(
                "Expected a sleep of at least %dms but only slept for %d",
                expected, bpStrategy.getSleepTime()));
      }
    }
  }

  private static void consume(FlightStream stream) {
    VectorSchemaRoot root = stream.getRoot();
    while (stream.next()) {
      root.clear();
    }
  }

  private static void consume(FlightStream stream, int batches) {
    VectorSchemaRoot root = stream.getRoot();
    while (batches > 0 && stream.next()) {
      root.clear();
      batches--;
    }
  }

  private interface SleepTimeRecordingBackpressureStrategy extends BackpressureStrategy {
    /**
     * Returns the total time spent waiting on the listener to be ready.
     *
     * @return the total time spent waiting on the listener to be ready.
     */
    long getSleepTime();
  }

  /**
   * Implementation of a backpressure strategy that polls on isReady and records amount of time
   * spent in Thread.sleep().
   */
  private static class PollingBackpressureStrategy
      implements SleepTimeRecordingBackpressureStrategy {
    private final AtomicLong sleepTime = new AtomicLong(0);
    private FlightProducer.ServerStreamListener listener;

    @Override
    public long getSleepTime() {
      return sleepTime.get();
    }

    @Override
    public void register(FlightProducer.ServerStreamListener listener) {
      this.listener = listener;
    }

    @Override
    public WaitResult waitForListener(long timeout) {
      while (!listener.isReady()) {
        try {
          Thread.sleep(1);
          sleepTime.addAndGet(1L);
        } catch (InterruptedException expected) {
          // it is expected and no action needed
        }
      }
      return WaitResult.READY;
    }
  }

  /**
   * Implementation of a backpressure strategy that uses callbacks to detect changes in client
   * readiness state and records spent time waiting.
   */
  private static class RecordingCallbackBackpressureStrategy
      extends BackpressureStrategy.CallbackBackpressureStrategy
      implements SleepTimeRecordingBackpressureStrategy {
    private final AtomicLong sleepTime = new AtomicLong(0);

    @Override
    public long getSleepTime() {
      return sleepTime.get();
    }

    @Override
    public WaitResult waitForListener(long timeout) {
      final long startTime = System.currentTimeMillis();
      final WaitResult result = super.waitForListener(timeout);
      sleepTime.addAndGet(System.currentTimeMillis() - startTime);
      return result;
    }
  }
}
