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.beam.runners.dataflow.worker.windmill.client;
  
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.concurrent.ThreadLocalRandom;
  import java.util.function.Supplier;
  import javax.annotation.concurrent.GuardedBy;
  import javax.annotation.concurrent.ThreadSafe;
  import org.apache.beam.sdk.annotations.Internal;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
  import org.checkerframework.checker.nullness.qual.Nullable;
  import org.joda.time.Duration;
  import org.joda.time.Instant;
  
  /**
   * Pool of homogeneous streams to Windmill.
   *
   * <p>The pool holds a fixed total number of streams, and keeps each stream open for a specified
   * time to allow for better load-balancing.
   */
  @Internal
  @ThreadSafe
  public class WindmillStreamPool<StreamT extends WindmillStream> {
  
    private final Duration streamTimeout;
    private final Supplier<StreamT> streamSupplier;
  
    /** @implNote Size of streams never changes once initialized. */
    private final List<@Nullable StreamData<StreamT>> streams;
  
    @GuardedBy("this")
    private final Map<StreamT, StreamData<StreamT>> holds;
  
    private WindmillStreamPool(
        Duration streamTimeout,
        Supplier<StreamT> streamSupplier,
        List<@Nullable StreamData<StreamT>> streams,
        Map<StreamT, StreamData<StreamT>> holds) {
      this.streams = streams;
      this.streamTimeout = streamTimeout;
      this.streamSupplier = streamSupplier;
      this.holds = holds;
    }
  
    public static <StreamT extends WindmillStream> WindmillStreamPool<StreamT> create(
        int numStreams, Duration streamTimeout, Supplier<StreamT> streamSupplier) {
      return new WindmillStreamPool<>(
          streamTimeout, streamSupplier, newStreamList(numStreams), new HashMap<>());
    }
  
    @VisibleForTesting
    static <StreamT extends WindmillStream> WindmillStreamPool<StreamT> forTesting(
        Duration streamTimeout,
        Supplier<StreamT> streamSupplier,
        List<@Nullable StreamData<StreamT>> streamPool,
        Map<StreamT, StreamData<StreamT>> holds) {
      return new WindmillStreamPool<>(streamTimeout, streamSupplier, streamPool, holds);
    }
  
    /**
     * Creates a new list of streams of the given capacity with all values initialized to null. This
     * is because we randomly load balance across all the streams in the pool.
     */
    @VisibleForTesting
    static <StreamT extends WindmillStream> List<@Nullable StreamData<StreamT>> newStreamList(
        int numStreams) {
      List<@Nullable StreamData<StreamT>> streamPool = new ArrayList<>(numStreams);
      for (int i = 0; i < numStreams; i++) {
        streamPool.add(null);
      }
      return streamPool;
    }
  
    /**
     * Returns a stream for use that may be cached from a previous call. Each call of getStream must
     * be matched with a call of {@link WindmillStreamPool#releaseStream(WindmillStream)}. If the
     * stream has been cached but has timed out and drained (no longer has any holds), the stream will
     * be closed.
     */
    public StreamT getStream() {
      int index = streams.size() == 1 ? 0 : ThreadLocalRandom.current().nextInt(streams.size());
      // We will return this stream
      StreamT resultStream;
      StreamT closeThisStream = null;
      try {
        synchronized (this) {
          WindmillStreamPool.StreamData<StreamT> existingStreamData = streams.get(index);
          // There are 3 possible states that can result from fetching the stream from the cache.
          if (existingStreamData == null) {
            // 1. Stream doesn't exist create and cache a new one.
            resultStream = createAndCacheStream(index).stream;
          } else if (existingStreamData.hasTimedOut(streamTimeout)) {
            // 2. The stream exists, but has timed out. The timed out stream is not returned (a new
            // one is created and returned here) and evicted from the cache if the stream has
            // completely drained. Every call to getStream(), is matched with a call to
            // releaseStream(), so the stream will eventually drain and be closed.
            if (--existingStreamData.holds == 0) {
              holds.remove(existingStreamData.stream);
              closeThisStream = existingStreamData.stream;
            }
            // Create and cache a new stream at the timed out stream's index.
            resultStream = createAndCacheStream(index).stream;
          } else {
            // 3. The stream exists and is in a valid state.
            existingStreamData.holds++;
            resultStream = existingStreamData.stream;
          }
        }
        return resultStream;
      } finally {
        if (closeThisStream != null) {
          closeThisStream.halfClose();
        }
      }
    }
  
    public CloseableStream<StreamT> getCloseableStream() {
      StreamT stream = getStream();
      return CloseableStream.create(stream, () -> releaseStream(stream));
    }
  
    private synchronized WindmillStreamPool.StreamData<StreamT> createAndCacheStream(int cacheKey) {
      WindmillStreamPool.StreamData<StreamT> newStreamData =
          new WindmillStreamPool.StreamData<>(streamSupplier.get());
      newStreamData.holds++;
      streams.set(cacheKey, newStreamData);
      holds.put(newStreamData.stream, newStreamData);
      return newStreamData;
    }
  
    /** Releases a stream that was obtained with {@link WindmillStreamPool#getStream()}. */
    public void releaseStream(StreamT stream) {
      boolean closeStream = false;
      synchronized (this) {
        StreamData<StreamT> streamData = holds.get(stream);
        // All streams that are created by an instance of a pool will be present.
        if (streamData == null) {
          throw new IllegalStateException(
              "Attempted to release stream that does not exist in this pool. This stream "
                  + "may not have been created by this pool or may have been released more "
                  + "times than acquired.");
        }
        if (--streamData.holds == 0) {
          closeStream = true;
          holds.remove(stream);
        }
      }
  
      if (closeStream) {
        stream.halfClose();
      }
    }
  
    @VisibleForTesting
    static final class StreamData<StreamT extends WindmillStream> {
      final StreamT stream;
      int holds;
  
      @VisibleForTesting
      StreamData(StreamT stream) {
        this.stream = stream;
        holds = 1;
      }
  
      private boolean hasTimedOut(Duration timeout) {
        return stream.startTime().isBefore(Instant.now().minus(timeout));
      }
    }
  }
  