<NME> Watchdog.java
<BEF> /*
 * Copyright 2022 Google Inc. All Rights Reserved.
 *
 *  Licensed 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 com.google.cloud.hadoop.gcsio;

import static com.google.common.base.Preconditions.checkNotNull;
import static java.util.concurrent.Executors.newSingleThreadScheduledExecutor;
import static java.util.concurrent.TimeUnit.MILLISECONDS;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.flogger.GoogleLogger;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import io.grpc.ClientCall;
import io.grpc.Context.CancellableContext;
import io.grpc.stub.StreamObserver;
import java.time.Clock;
import java.time.Duration;
import java.util.Iterator;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentHashMap.KeySetView;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeoutException;
import javax.annotation.Nonnull;
import javax.annotation.concurrent.GuardedBy;

/**
 * Prevents the streams from hanging indefinitely. This middleware garbage collects idle streams in
 * case the user forgot to close a ServerStream or if a connection is reset and GRPC does not get
 * notified.
 *
 * <p>For every {@code checkInterval}, this class checks thresholds:
 *
 * <ul>
 *   <li>idleTimeout: the amount of time to wait before assuming that the caller forgot to close the
 *       stream and forcefully closing the stream. This is measured from the last time the caller
 *       had no outstanding demand. Duration.ZERO disables the timeout.
 * </ul>
 */
final class Watchdog implements Runnable {
  private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();

  private final KeySetView<WatchdogStream, Boolean> openStreams = ConcurrentHashMap.newKeySet();

  private final Duration scheduleInterval;

  private final Clock clock = Clock.systemUTC();

  private final ScheduledExecutorService executor =
      newSingleThreadScheduledExecutor(
          new ThreadFactoryBuilder()
              .setNameFormat("gcs-background-watchdog-pool-%d")
              .setDaemon(true)
              .build());

  /** returns a Watchdog which is scheduled at the provided interval. */
  public static Watchdog create(Duration scheduleInterval) {
    Watchdog watchdog = new Watchdog(scheduleInterval);
    watchdog.start();
    return watchdog;
  }

  private Watchdog(Duration scheduleInterval) {
    this.scheduleInterval = scheduleInterval;
  }

  @SuppressWarnings("FutureReturnValueIgnored")
  private void start() {
    executor.scheduleAtFixedRate(
        this, /* initialDelay= */ 0, scheduleInterval.toMillis(), MILLISECONDS);
  }

  /**
   * Wraps the target {@code StreamObserver} with timing constraints. Used for ServerStreamingRPC
   * calls
   *
   * @param requestContext request to be cancelled when the timer is overrun
   * @param responseIterator responseIterator to monitor the state of the RPC
   * @param idleTimeout Idle time duration that triggers client call cancellation
   * @param <T> Type reference for the RPC response
   * @return {@link Iterator}
   */
  public <T> Iterator<T> watch(
      CancellableContext requestContext,
      Iterator<T> responseIterator,
      @Nonnull Duration idleTimeout) {
    checkNotNull(responseIterator, "responseIterator can't be null");
    checkNotNull(idleTimeout, "idleTimeout can't be null");

    if (idleTimeout.isZero()) {
      return responseIterator;
    }

    ServerStreamingRPCWatchdogStream<T> stream =
        new ServerStreamingRPCWatchdogStream<>(requestContext, responseIterator, idleTimeout);
    openStreams.add(stream);
    return stream;
  }

  /**
   * Wraps the target {@code StreamObserver} with timing constraints. Used for ClientStreamingRPC
   * calls
   *
   * @param clientCall ClientCall to be cancelled when the timer is overrun
   * @param streamObserver Observer instance to monitor the state of the RPC
   * @param idleTimeout Idle time duration that triggers client call cancellation
   * @param <T> Type reference for the RPC response
   * @return {@link StreamObserver}
   */
  public <R, T> StreamObserver<R> watch(
      ClientCall<R, T> clientCall,
      StreamObserver<R> streamObserver,
      @Nonnull Duration idleTimeout) {
    checkNotNull(streamObserver, "streamObserver can't be null");
    checkNotNull(idleTimeout, "idleTimeout can't be null");

    if (idleTimeout.isZero()) {
      return streamObserver;
    }

    ClientStreamingRPCWatchdogStream<R, T> stream =
        new ClientStreamingRPCWatchdogStream<>(clientCall, streamObserver, idleTimeout);
    openStreams.add(stream);
    return stream;
  }

  @Override
  public void run() {
    try {
      runUnsafe();
    } catch (RuntimeException e) {
      logger.atWarning().withCause(e).log(
          "Caught RuntimeException in periodic Watchdog run, continuing.");
    }
  }

  @VisibleForTesting
  KeySetView<WatchdogStream, Boolean> getOpenStreams() {
    return openStreams;
  }

  private void runUnsafe() {
    openStreams.removeIf(WatchdogStream::cancelIfStale);
  }

  /** clean up resources used by the class */
  public void shutdown() {
    executor.shutdown();
  }

  enum State {
    /** Stream has been started, but doesn't have any outstanding requests. */
    IDLE,
    /** Stream is awaiting a response from upstream. */
    WAITING,
    /**
     * Stream received a response from upstream, and is awaiting outerResponseObserver processing.
     */
    DELIVERING
  }

  interface WatchdogStream {
    /**
     * Checks if this stream has overrun any of timeouts and cancels it if it does.
     *
     * @return True if the stream was canceled.
     */
    boolean cancelIfStale();
  }

  class ServerStreamingRPCWatchdogStream<T> implements Iterator<T>, WatchdogStream {
    private final Object lock = new Object();

    private final Duration waitTimeout;

    private final Iterator<T> innerIterator;

    private final CancellableContext requestContext;

    @GuardedBy("lock")
    private long lastActivityAt = clock.millis();

    @GuardedBy("lock")
    private State state = State.IDLE;

    ServerStreamingRPCWatchdogStream(
        CancellableContext requestContext, Iterator<T> iterator, Duration waitTimeout) {
      this.waitTimeout = waitTimeout;
      this.requestContext = requestContext;
      this.innerIterator = iterator;
    }
        return false;
      }

      Throwable throwable = null;
      synchronized (lock) {
        long waitTime = clock.millis() - lastActivityAt;
      if (innerIterator == null || requestContext == null) {
        return false;
      }

      // No need to monitor the stream, If the request is cancelled outside of watchdog
      if (requestContext.isCancelled()) {
        openStreams.remove(this);
        return false;
      }

      Throwable throwable = null;
      synchronized (lock) {
        long waitTime = clock.millis() - lastActivityAt;
        if (state == State.WAITING && !waitTimeout.isZero() && waitTime >= waitTimeout.toMillis()) {
          throwable = new TimeoutException("Canceled due to timeout waiting for next response");
        }
      }
      if (throwable != null) {
        requestContext.cancel(throwable);
        return true;
      }
      return false;
    }

    @Override
    public boolean hasNext() {
      boolean hasNext = false;
      try {
        synchronized (lock) {
          lastActivityAt = clock.millis();
          state = State.WAITING;
        }
        hasNext = innerIterator.hasNext();
      } finally {
        // stream is complete successfully with no more items or has thrown an exception
        if (!hasNext) {
          openStreams.remove(this);
        }
      }
      return hasNext;
    }
  }

  class ClientStreamingRPCWatchdogStream<R, T> implements StreamObserver<R>, WatchdogStream {

    private final Object lock = new Object();

    private final Duration waitTimeout;

    private final StreamObserver<R> innerStreamObserver;

    private final ClientCall<R, T> clientCall;

    @GuardedBy("lock")
    private long lastActivityAt = clock.millis();

    @GuardedBy("lock")
    private State state = State.IDLE;

    public ClientStreamingRPCWatchdogStream(
        ClientCall<R, T> clientCall, StreamObserver<R> innerStreamObserver, Duration waitTimeout) {
      this.clientCall = clientCall;
      this.innerStreamObserver = innerStreamObserver;
      this.waitTimeout = waitTimeout;
    }

    @Override
    public boolean cancelIfStale() {
      if (innerStreamObserver == null || clientCall == null) {
        return false;
      }

      Throwable throwable = null;
      synchronized (lock) {
        long waitTime = clock.millis() - lastActivityAt;
        if (state == State.WAITING && !waitTimeout.isZero() && waitTime >= waitTimeout.toMillis()) {
          throwable = new TimeoutException("Canceled due to timeout waiting for next response");
        }
      }
      if (throwable != null) {
        clientCall.cancel(throwable.getMessage(), throwable);
        return true;
      }
      return false;
    }

    @Override
    public void onNext(R value) {
      synchronized (lock) {
        lastActivityAt = clock.millis();
        state = State.WAITING;
      }
      innerStreamObserver.onNext(value);
      synchronized (lock) {
        state = State.IDLE;
      }
    }

    @Override
    public void onError(Throwable t) {
      openStreams.remove(this);
      innerStreamObserver.onError(t);
    }

    @Override
    public void onCompleted() {
      openStreams.remove(this);
      innerStreamObserver.onCompleted();
    }
  }
}

<MSG> Remove tracking if request is cancelled (#680)


<DFF> @@ -201,6 +201,12 @@ final class Watchdog implements Runnable {
         return false;
       }
 
+      // No need to monitor the stream, If the request is cancelled outside of watchdog
+      if (requestContext.isCancelled()) {
+        openStreams.remove(this);
+        return false;
+      }
+
       Throwable throwable = null;
       synchronized (lock) {
         long waitTime = clock.millis() - lastActivityAt;

{'commit': 'diff --git a/gcsio/src/main/java/com/google/cloud/hadoop/gcsio/Watchdog.java b/gcsio/src/main/java/com/google/cloud/hadoop/gcsio/Watchdog.java\nindex 4d528d58..19aa2927 100644\n--- a/gcsio/src/main/java/com/google/cloud/hadoop/gcsio/Watchdog.java\n+++ b/gcsio/src/main/java/com/google/cloud/hadoop/gcsio/Watchdog.java\n@@ -201,6 +201,12 @@ final class Watchdog implements Runnable {\n         return false;\n       }\n \n+      // No need to monitor the stream, If the request is cancelled outside of watchdog\n+      if (requestContext.isCancelled()) {\n+        openStreams.remove(this);\n+        return false;\n+      }\n+\n       Throwable throwable = null;\n       synchronized (lock) {\n         long waitTime = clock.millis() - lastActivityAt;\n', 'commit_subject': 'Remove tracking if request is cancelled (#680)', 'repo_name': 'GoogleCloudDataproc/hadoop-connectors', 'committer': {'name': 'GitHub', 'email': '1505422b2465e9a84f6fdfaa161078890c593f06@github.com', 'time_sec': '1641896691', 'tz_offset': '330', 'date': {'seconds': '1641896691'}}, 'commit_tree': '37d88c43f57bc7d0dd279e9e155b5e3223e34185', 'commit_message': 'Remove tracking if request is cancelled (#680)\n\n', 'language_name': 'Java', 'license': 'apache-2.0', 'hunks': '@@ -201,6 +201,12 @@ final class Watchdog implements Runnable {\n         return false;\n       }\n \n+      // No need to monitor the stream, If the request is cancelled outside of watchdog\n+      if (requestContext.isCancelled()) {\n+        openStreams.remove(this);\n+        return false;\n+      }\n+\n       Throwable throwable = null;\n       synchronized (lock) {\n         long waitTime = clock.millis() - lastActivityAt;\n', 'addition_count': 6, 'deletion_count': 0, 'src_file': 'a/gcsio/src/main/java/com/google/cloud/hadoop/gcsio/Watchdog.java', 'tgt_file': 'b/gcsio/src/main/java/com/google/cloud/hadoop/gcsio/Watchdog.java', 'file_extension': '.java', 'before_file': '/*\n * Copyright 2022 Google Inc. All Rights Reserved.\n *\n *  Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except\n * in compliance with the License. You may obtain a copy of the License at\n *\n *  http://www.apache.org/licenses/LICENSE-2.0\n *\n *  Unless required by applicable law or agreed to in writing, software distributed under the\n * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either\n * express or implied. See the License for the specific language governing permissions and\n * limitations under the License.\n */\npackage com.google.cloud.hadoop.gcsio;\n\nimport static com.google.common.base.Preconditions.checkNotNull;\nimport static java.util.concurrent.Executors.newSingleThreadScheduledExecutor;\nimport static java.util.concurrent.TimeUnit.MILLISECONDS;\n\nimport com.google.common.annotations.VisibleForTesting;\nimport com.google.common.flogger.GoogleLogger;\nimport com.google.common.util.concurrent.ThreadFactoryBuilder;\nimport io.grpc.ClientCall;\nimport io.grpc.Context.CancellableContext;\nimport io.grpc.stub.StreamObserver;\nimport java.time.Clock;\nimport java.time.Duration;\nimport java.util.Iterator;\nimport java.util.concurrent.ConcurrentHashMap;\nimport java.util.concurrent.ConcurrentHashMap.KeySetView;\nimport java.util.concurrent.ScheduledExecutorService;\nimport java.util.concurrent.TimeoutException;\nimport javax.annotation.Nonnull;\nimport javax.annotation.concurrent.GuardedBy;\n\n/**\n * Prevents the streams from hanging indefinitely. This middleware garbage collects idle streams in\n * case the user forgot to close a ServerStream or if a connection is reset and GRPC does not get\n * notified.\n *\n * <p>For every {@code checkInterval}, this class checks thresholds:\n *\n * <ul>\n *   <li>idleTimeout: the amount of time to wait before assuming that the caller forgot to close the\n *       stream and forcefully closing the stream. This is measured from the last time the caller\n *       had no outstanding demand. Duration.ZERO disables the timeout.\n * </ul>\n */\nfinal class Watchdog implements Runnable {\n  private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();\n\n  private final KeySetView<WatchdogStream, Boolean> openStreams = ConcurrentHashMap.newKeySet();\n\n  private final Duration scheduleInterval;\n\n  private final Clock clock = Clock.systemUTC();\n\n  private final ScheduledExecutorService executor =\n      newSingleThreadScheduledExecutor(\n          new ThreadFactoryBuilder()\n              .setNameFormat("gcs-background-watchdog-pool-%d")\n              .setDaemon(true)\n              .build());\n\n  /** returns a Watchdog which is scheduled at the provided interval. */\n  public static Watchdog create(Duration scheduleInterval) {\n    Watchdog watchdog = new Watchdog(scheduleInterval);\n    watchdog.start();\n    return watchdog;\n  }\n\n  private Watchdog(Duration scheduleInterval) {\n    this.scheduleInterval = scheduleInterval;\n  }\n\n  @SuppressWarnings("FutureReturnValueIgnored")\n  private void start() {\n    executor.scheduleAtFixedRate(\n        this, /* initialDelay= */ 0, scheduleInterval.toMillis(), MILLISECONDS);\n  }\n\n  /**\n   * Wraps the target {@code StreamObserver} with timing constraints. Used for ServerStreamingRPC\n   * calls\n   *\n   * @param requestContext request to be cancelled when the timer is overrun\n   * @param responseIterator responseIterator to monitor the state of the RPC\n   * @param idleTimeout Idle time duration that triggers client call cancellation\n   * @param <T> Type reference for the RPC response\n   * @return {@link Iterator}\n   */\n  public <T> Iterator<T> watch(\n      CancellableContext requestContext,\n      Iterator<T> responseIterator,\n      @Nonnull Duration idleTimeout) {\n    checkNotNull(responseIterator, "responseIterator can\'t be null");\n    checkNotNull(idleTimeout, "idleTimeout can\'t be null");\n\n    if (idleTimeout.isZero()) {\n      return responseIterator;\n    }\n\n    ServerStreamingRPCWatchdogStream<T> stream =\n        new ServerStreamingRPCWatchdogStream<>(requestContext, responseIterator, idleTimeout);\n    openStreams.add(stream);\n    return stream;\n  }\n\n  /**\n   * Wraps the target {@code StreamObserver} with timing constraints. Used for ClientStreamingRPC\n   * calls\n   *\n   * @param clientCall ClientCall to be cancelled when the timer is overrun\n   * @param streamObserver Observer instance to monitor the state of the RPC\n   * @param idleTimeout Idle time duration that triggers client call cancellation\n   * @param <T> Type reference for the RPC response\n   * @return {@link StreamObserver}\n   */\n  public <R, T> StreamObserver<R> watch(\n      ClientCall<R, T> clientCall,\n      StreamObserver<R> streamObserver,\n      @Nonnull Duration idleTimeout) {\n    checkNotNull(streamObserver, "streamObserver can\'t be null");\n    checkNotNull(idleTimeout, "idleTimeout can\'t be null");\n\n    if (idleTimeout.isZero()) {\n      return streamObserver;\n    }\n\n    ClientStreamingRPCWatchdogStream<R, T> stream =\n        new ClientStreamingRPCWatchdogStream<>(clientCall, streamObserver, idleTimeout);\n    openStreams.add(stream);\n    return stream;\n  }\n\n  @Override\n  public void run() {\n    try {\n      runUnsafe();\n    } catch (RuntimeException e) {\n      logger.atWarning().withCause(e).log(\n          "Caught RuntimeException in periodic Watchdog run, continuing.");\n    }\n  }\n\n  @VisibleForTesting\n  KeySetView<WatchdogStream, Boolean> getOpenStreams() {\n    return openStreams;\n  }\n\n  private void runUnsafe() {\n    openStreams.removeIf(WatchdogStream::cancelIfStale);\n  }\n\n  /** clean up resources used by the class */\n  public void shutdown() {\n    executor.shutdown();\n  }\n\n  enum State {\n    /** Stream has been started, but doesn\'t have any outstanding requests. */\n    IDLE,\n    /** Stream is awaiting a response from upstream. */\n    WAITING,\n    /**\n     * Stream received a response from upstream, and is awaiting outerResponseObserver processing.\n     */\n    DELIVERING\n  }\n\n  interface WatchdogStream {\n    /**\n     * Checks if this stream has overrun any of timeouts and cancels it if it does.\n     *\n     * @return True if the stream was canceled.\n     */\n    boolean cancelIfStale();\n  }\n\n  class ServerStreamingRPCWatchdogStream<T> implements Iterator<T>, WatchdogStream {\n    private final Object lock = new Object();\n\n    private final Duration waitTimeout;\n\n    private final Iterator<T> innerIterator;\n\n    private final CancellableContext requestContext;\n\n    @GuardedBy("lock")\n    private long lastActivityAt = clock.millis();\n\n    @GuardedBy("lock")\n    private State state = State.IDLE;\n\n    ServerStreamingRPCWatchdogStream(\n        CancellableContext requestContext, Iterator<T> iterator, Duration waitTimeout) {\n      this.waitTimeout = waitTimeout;\n      this.requestContext = requestContext;\n      this.innerIterator = iterator;\n    }\n        return false;\n      }\n\n      Throwable throwable = null;\n      synchronized (lock) {\n        long waitTime = clock.millis() - lastActivityAt;\n      if (innerIterator == null || requestContext == null) {\n        return false;\n      }\n\n      // No need to monitor the stream, If the request is cancelled outside of watchdog\n      if (requestContext.isCancelled()) {\n        openStreams.remove(this);\n        return false;\n      }\n\n      Throwable throwable = null;\n      synchronized (lock) {\n        long waitTime = clock.millis() - lastActivityAt;\n        if (state == State.WAITING && !waitTimeout.isZero() && waitTime >= waitTimeout.toMillis()) {\n          throwable = new TimeoutException("Canceled due to timeout waiting for next response");\n        }\n      }\n      if (throwable != null) {\n        requestContext.cancel(throwable);\n        return true;\n      }\n      return false;\n    }\n\n    @Override\n    public boolean hasNext() {\n      boolean hasNext = false;\n      try {\n        synchronized (lock) {\n          lastActivityAt = clock.millis();\n          state = State.WAITING;\n        }\n        hasNext = innerIterator.hasNext();\n      } finally {\n        // stream is complete successfully with no more items or has thrown an exception\n        if (!hasNext) {\n          openStreams.remove(this);\n        }\n      }\n      return hasNext;\n    }\n  }\n\n  class ClientStreamingRPCWatchdogStream<R, T> implements StreamObserver<R>, WatchdogStream {\n\n    private final Object lock = new Object();\n\n    private final Duration waitTimeout;\n\n    private final StreamObserver<R> innerStreamObserver;\n\n    private final ClientCall<R, T> clientCall;\n\n    @GuardedBy("lock")\n    private long lastActivityAt = clock.millis();\n\n    @GuardedBy("lock")\n    private State state = State.IDLE;\n\n    public ClientStreamingRPCWatchdogStream(\n        ClientCall<R, T> clientCall, StreamObserver<R> innerStreamObserver, Duration waitTimeout) {\n      this.clientCall = clientCall;\n      this.innerStreamObserver = innerStreamObserver;\n      this.waitTimeout = waitTimeout;\n    }\n\n    @Override\n    public boolean cancelIfStale() {\n      if (innerStreamObserver == null || clientCall == null) {\n        return false;\n      }\n\n      Throwable throwable = null;\n      synchronized (lock) {\n        long waitTime = clock.millis() - lastActivityAt;\n        if (state == State.WAITING && !waitTimeout.isZero() && waitTime >= waitTimeout.toMillis()) {\n          throwable = new TimeoutException("Canceled due to timeout waiting for next response");\n        }\n      }\n      if (throwable != null) {\n        clientCall.cancel(throwable.getMessage(), throwable);\n        return true;\n      }\n      return false;\n    }\n\n    @Override\n    public void onNext(R value) {\n      synchronized (lock) {\n        lastActivityAt = clock.millis();\n        state = State.WAITING;\n      }\n      innerStreamObserver.onNext(value);\n      synchronized (lock) {\n        state = State.IDLE;\n      }\n    }\n\n    @Override\n    public void onError(Throwable t) {\n      openStreams.remove(this);\n      innerStreamObserver.onError(t);\n    }\n\n    @Override\n    public void onCompleted() {\n      openStreams.remove(this);\n      innerStreamObserver.onCompleted();\n    }\n  }\n}\n', 'message': 'Remove tracking if request is cancelled (#680)\n\n', 'author_name': 'GitHub', 'lang': 'java'}