/*
 * 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.dataflow.worker.streaming.harness;

import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import javax.annotation.concurrent.ThreadSafe;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GetWorkRequest;
import org.apache.beam.runners.dataflow.worker.windmill.WindmillConnection;
import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.CommitWorkStream;
import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.GetDataStream;
import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.GetWorkStream;
import org.apache.beam.runners.dataflow.worker.windmill.client.commits.WorkCommitter;
import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.GetDataClient;
import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.GrpcWindmillStreamFactory;
import org.apache.beam.runners.dataflow.worker.windmill.work.WorkItemScheduler;
import org.apache.beam.runners.dataflow.worker.windmill.work.budget.GetWorkBudget;
import org.apache.beam.runners.dataflow.worker.windmill.work.budget.GetWorkBudgetSpender;
import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.FixedStreamHeartbeatSender;
import org.apache.beam.sdk.annotations.Internal;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder;

/**
 * Owns and maintains a set of streams used to communicate with a specific Windmill worker.
 *
 * <p>Once started, the underlying streams are "alive" until they are manually closed via {@link
 * #close()}.
 *
 * <p>If closed, it means that the backend endpoint is no longer in the worker set. Once closed,
 * these instances are not reused.
 *
 * @implNote Does not manage streams for fetching {@link
 *     org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalData} for side inputs.
 */
@Internal
@ThreadSafe
final class WindmillStreamSender implements GetWorkBudgetSpender, StreamSender {
  private static final String STREAM_STARTER_THREAD_NAME = "StartWindmillStreamThread-%d";
  private static final int TERMINATION_TIMEOUT_SECONDS = 5;
  private final AtomicBoolean started;
  private final AtomicReference<GetWorkBudget> getWorkBudget;
  private final GetWorkStream getWorkStream;
  private final GetDataStream getDataStream;
  private final CommitWorkStream commitWorkStream;
  private final WorkCommitter workCommitter;
  private final ExecutorService streamStarter;

  private WindmillStreamSender(
      WindmillConnection connection,
      GetWorkRequest getWorkRequest,
      AtomicReference<GetWorkBudget> getWorkBudget,
      GrpcWindmillStreamFactory streamingEngineStreamFactory,
      WorkItemScheduler workItemScheduler,
      Function<GetDataStream, GetDataClient> getDataClientFactory,
      Function<CommitWorkStream, WorkCommitter> workCommitterFactory) {
    this.started = new AtomicBoolean(false);
    this.getWorkBudget = getWorkBudget;

    // Stream instances connect/reconnect internally, so we can reuse the same instance through the
    // entire lifecycle of WindmillStreamSender.
    this.getDataStream = streamingEngineStreamFactory.createDirectGetDataStream(connection);
    this.commitWorkStream = streamingEngineStreamFactory.createDirectCommitWorkStream(connection);
    this.workCommitter = workCommitterFactory.apply(commitWorkStream);
    this.getWorkStream =
        streamingEngineStreamFactory.createDirectGetWorkStream(
            connection,
            withRequestBudget(getWorkRequest, getWorkBudget.get()),
            FixedStreamHeartbeatSender.create(getDataStream),
            getDataClientFactory.apply(getDataStream),
            workCommitter,
            workItemScheduler);
    // 3 threads, 1 for each stream type (GetWork, GetData, CommitWork).
    this.streamStarter =
        Executors.newFixedThreadPool(
            3, new ThreadFactoryBuilder().setNameFormat(STREAM_STARTER_THREAD_NAME).build());
  }

  static WindmillStreamSender create(
      WindmillConnection connection,
      GetWorkRequest getWorkRequest,
      GetWorkBudget getWorkBudget,
      GrpcWindmillStreamFactory streamingEngineStreamFactory,
      WorkItemScheduler workItemScheduler,
      Function<GetDataStream, GetDataClient> getDataClientFactory,
      Function<CommitWorkStream, WorkCommitter> workCommitterFactory) {
    return new WindmillStreamSender(
        connection,
        getWorkRequest,
        new AtomicReference<>(getWorkBudget),
        streamingEngineStreamFactory,
        workItemScheduler,
        getDataClientFactory,
        workCommitterFactory);
  }

  private static GetWorkRequest withRequestBudget(GetWorkRequest request, GetWorkBudget budget) {
    return request.toBuilder().setMaxItems(budget.items()).setMaxBytes(budget.bytes()).build();
  }

  synchronized void start() {
    if (!started.get()) {
      checkState(!streamStarter.isShutdown(), "WindmillStreamSender has already been shutdown.");

      // Start these 3 streams in parallel since they each may perform blocking IO.
      CompletableFuture.allOf(
              CompletableFuture.runAsync(getWorkStream::start, streamStarter),
              CompletableFuture.runAsync(getDataStream::start, streamStarter),
              CompletableFuture.runAsync(commitWorkStream::start, streamStarter))
          .join();
      workCommitter.start();
      started.set(true);
    }
  }

  @Override
  public synchronized void close() {
    streamStarter.shutdown();
    getWorkStream.shutdown();
    getDataStream.shutdown();
    workCommitter.stop();
    commitWorkStream.shutdown();
    try {
      if (!Preconditions.checkNotNull(streamStarter)
          .awaitTermination(TERMINATION_TIMEOUT_SECONDS, TimeUnit.SECONDS)) {
        streamStarter.shutdownNow();
      }
      Preconditions.checkNotNull(getWorkStream)
          .awaitTermination(TERMINATION_TIMEOUT_SECONDS, TimeUnit.SECONDS);
      Preconditions.checkNotNull(getDataStream)
          .awaitTermination(TERMINATION_TIMEOUT_SECONDS, TimeUnit.SECONDS);
      Preconditions.checkNotNull(commitWorkStream)
          .awaitTermination(TERMINATION_TIMEOUT_SECONDS, TimeUnit.SECONDS);
    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
      throw new RuntimeException("Interrupted while waiting for streams to terminate", e);
    }
  }

  @Override
  public void setBudget(long items, long bytes) {
    GetWorkBudget budget = GetWorkBudget.builder().setItems(items).setBytes(bytes).build();
    getWorkBudget.set(budget);
    if (started.get()) {
      getWorkStream.setBudget(budget);
    }
  }

  long getCurrentActiveCommitBytes() {
    return workCommitter.currentActiveCommitBytes();
  }
}
