/*
 * 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.beam.runners.fnexecution.logging;

import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.containsInAnyOrder;

import java.util.ArrayList;
import java.util.Collection;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.function.Consumer;
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.LogControl;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.LogEntry;
import org.apache.beam.model.fnexecution.v1.BeamFnLoggingGrpc;
import org.apache.beam.sdk.fn.server.GrpcFnServer;
import org.apache.beam.sdk.fn.server.InProcessServerFactory;
import org.apache.beam.sdk.fn.test.TestStreams;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ManagedChannel;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.inprocess.InProcessChannelBuilder;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.StreamObserver;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.Timeout;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;

/** Tests for {@link GrpcLoggingService}. */
@RunWith(JUnit4.class)
public class GrpcLoggingServiceTest {
  @Rule public transient Timeout globalTimeout = Timeout.seconds(600);
  private Consumer<LogControl> messageDiscarder =
      item -> {
        // Ignore
      };

  @Test
  public void testMultipleClientsSuccessfullyProcessed() throws Exception {
    ConcurrentLinkedQueue<BeamFnApi.LogEntry> logs = new ConcurrentLinkedQueue<>();
    GrpcLoggingService service =
        GrpcLoggingService.forWriter(new CollectionAppendingLogWriter(logs));
    try (GrpcFnServer<GrpcLoggingService> server =
        GrpcFnServer.allocatePortAndCreateFor(service, InProcessServerFactory.create())) {
      ExecutorService executorService = Executors.newCachedThreadPool();
      Collection<Future<?>> futures = new ArrayList<>();
      CountDownLatch waitForServerHangup = new CountDownLatch(3);
      for (int i = 1; i <= 3; ++i) {
        final int instructionId = i;
        futures.add(
            executorService.submit(
                () -> {
                  String url = server.getApiServiceDescriptor().getUrl();
                  ManagedChannel channel = InProcessChannelBuilder.forName(url).build();
                  StreamObserver<LogEntry.List> outboundObserver =
                      BeamFnLoggingGrpc.newStub(channel)
                          .logging(
                              TestStreams.withOnNext(messageDiscarder)
                                  .withOnCompleted(new CountDown(waitForServerHangup))
                                  .build());
                  outboundObserver.onNext(createLogsWithIds(instructionId, -instructionId));
                  outboundObserver.onCompleted();
                }));
      }
      // Make sure all streams were created and issued client operations.
      for (Future<?> f : futures) {
        f.get();
      }
      // Ensure all the streams were completed as expected before closing the server.
      waitForServerHangup.await();
    }
    assertThat(
        logs,
        containsInAnyOrder(
            createLogWithId(1L),
            createLogWithId(2L),
            createLogWithId(3L),
            createLogWithId(-1L),
            createLogWithId(-2L),
            createLogWithId(-3L)));
  }

  @Test
  public void testMultipleClientsFailingIsHandledGracefullyByServer() throws Exception {
    ConcurrentLinkedQueue<BeamFnApi.LogEntry> logs = new ConcurrentLinkedQueue<>();
    GrpcLoggingService service =
        GrpcLoggingService.forWriter(new CollectionAppendingLogWriter(logs));
    try (GrpcFnServer<GrpcLoggingService> server =
        GrpcFnServer.allocatePortAndCreateFor(service, InProcessServerFactory.create())) {

      CountDownLatch waitForTermination = new CountDownLatch(3);
      final Collection<StreamObserver<LogEntry.List>> outboundObservers = new ArrayList<>();
      // Create all the streams
      for (int instructionId = 1; instructionId <= 3; ++instructionId) {
        ManagedChannel channel =
            InProcessChannelBuilder.forName(server.getApiServiceDescriptor().getUrl()).build();
        StreamObserver<LogEntry.List> outboundObserver =
            BeamFnLoggingGrpc.newStub(channel)
                .logging(
                    TestStreams.withOnNext(messageDiscarder)
                        .withOnCompleted(new CountDown(waitForTermination))
                        .withOnError(new CountDown(waitForTermination))
                        .build());
        outboundObserver.onNext(createLogsWithIds(instructionId, -instructionId));
        outboundObservers.add(outboundObserver);
      }

      for (StreamObserver<LogEntry.List> outboundObserver : outboundObservers) {
        outboundObserver.onError(new RuntimeException("Client"));
      }
      waitForTermination.await();
    }
  }

  @Test
  public void testServerCloseHangsUpClients() throws Exception {
    LinkedBlockingQueue<LogEntry> logs = new LinkedBlockingQueue<>();
    ExecutorService executorService = Executors.newCachedThreadPool();
    final GrpcLoggingService service =
        GrpcLoggingService.forWriter(new CollectionAppendingLogWriter(logs));
    CountDownLatch waitForServerHangup = new CountDownLatch(3);
    try (GrpcFnServer<GrpcLoggingService> server =
        GrpcFnServer.allocatePortAndCreateFor(service, InProcessServerFactory.create())) {

      Collection<Future<?>> futures = new ArrayList<>();
      for (int i = 1; i <= 3; ++i) {
        final long instructionId = i;
        futures.add(
            executorService.submit(
                () -> {
                  {
                    ManagedChannel channel =
                        InProcessChannelBuilder.forName(server.getApiServiceDescriptor().getUrl())
                            .build();
                    StreamObserver<LogEntry.List> outboundObserver =
                        BeamFnLoggingGrpc.newStub(channel)
                            .logging(
                                TestStreams.withOnNext(messageDiscarder)
                                    .withOnCompleted(new CountDown(waitForServerHangup))
                                    .build());
                    outboundObserver.onNext(createLogsWithIds(instructionId));
                    return null;
                  }
                }));
      }
      // Ensure all the streams have started and sent their instruction.
      for (Future<?> f : futures) {
        f.get();
      }
      // Wait till each client has sent their message showing that they have connected.
      for (int i = 1; i <= 3; ++i) {
        logs.take();
      }
      // Close the server without closing the streams and ensure they observe the hangup.
    }
    waitForServerHangup.await();
  }

  private BeamFnApi.LogEntry.List createLogsWithIds(long... ids) {
    BeamFnApi.LogEntry.List.Builder builder = BeamFnApi.LogEntry.List.newBuilder();
    for (long id : ids) {
      builder.addLogEntries(createLogWithId(id));
    }
    return builder.build();
  }

  private BeamFnApi.LogEntry createLogWithId(long id) {
    return BeamFnApi.LogEntry.newBuilder().setInstructionId(Long.toString(id)).build();
  }

  private static class CollectionAppendingLogWriter implements LogWriter {
    private final Collection<BeamFnApi.LogEntry> entries;

    private CollectionAppendingLogWriter(Collection<LogEntry> entries) {
      this.entries = entries;
    }

    @Override
    public void log(LogEntry entry) {
      entries.add(entry);
    }
  }

  /**
   * A {@link Runnable} that calls {@link CountDownLatch#countDown()} on a {@link CountDownLatch}.
   */
  private static class CountDown implements Runnable {
    private final CountDownLatch latch;

    CountDown(CountDownLatch latch) {
      this.latch = latch;
    }

    @Override
    public void run() {
      latch.countDown();
    }
  }
}
