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.bookkeeper.metadata.etcd;
  
  import static io.etcd.jetcd.common.exception.EtcdExceptionFactory.newClosedWatchClientException;
  import static io.etcd.jetcd.common.exception.EtcdExceptionFactory.newEtcdException;
  import static io.etcd.jetcd.common.exception.EtcdExceptionFactory.toEtcdException;
  
  import com.google.common.base.Strings;
  import com.google.protobuf.ByteString;
  import com.google.protobuf.UnsafeByteOperations;
  import io.etcd.jetcd.ByteSequence;
  import io.etcd.jetcd.Client;
  import io.etcd.jetcd.api.WatchCancelRequest;
  import io.etcd.jetcd.api.WatchCreateRequest;
  import io.etcd.jetcd.api.WatchGrpc;
  import io.etcd.jetcd.api.WatchRequest;
  import io.etcd.jetcd.api.WatchResponse;
  import io.etcd.jetcd.common.exception.ErrorCode;
  import io.etcd.jetcd.common.exception.EtcdException;
  import io.etcd.jetcd.common.exception.EtcdExceptionFactory;
  import io.etcd.jetcd.impl.EtcdConnectionManager;
  import io.etcd.jetcd.options.WatchOption;
  import io.etcd.jetcd.watch.WatchResponseWithError;
  import io.grpc.Status;
  import io.grpc.Status.Code;
  import io.grpc.stub.StreamObserver;
  import java.util.LinkedList;
  import java.util.Optional;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.ScheduledExecutorService;
  import java.util.concurrent.TimeUnit;
  import java.util.function.BiConsumer;
  import lombok.extern.slf4j.Slf4j;
  import org.apache.bookkeeper.common.concurrent.FutureUtils;
  import org.apache.bookkeeper.common.util.OrderedScheduler;
  import org.apache.bookkeeper.util.collections.ConcurrentLongHashMap;
  import org.apache.bookkeeper.util.collections.ConcurrentLongHashSet;
  
  /**
   * An async watch implementation.
   */
  @Slf4j
  public class EtcdWatchClient implements AutoCloseable {
  
      private final EtcdConnectionManager connMgr;
      private final WatchGrpc.WatchStub stub;
      private volatile StreamObserver<WatchRequest> grpcWatchStreamObserver;
      // watchers stores a mapping between watchID -> EtcdWatcher.
      private final ConcurrentLongHashMap<EtcdWatcher> watchers =
              ConcurrentLongHashMap.<EtcdWatcher>newBuilder().build();
      private final LinkedList<EtcdWatcher> pendingWatchers = new LinkedList<>();
      private final ConcurrentLongHashSet cancelSet = ConcurrentLongHashSet.newBuilder().build();
  
      // scheduler
      private final OrderedScheduler scheduler;
      private final ScheduledExecutorService watchExecutor;
  
      // close state
      private CompletableFuture<Void> closeFuture = null;
  
      public EtcdWatchClient(Client client) {
          this.connMgr = new EtcdConnectionManager(client);
          this.stub = connMgr.newWatchStub();
          this.scheduler = OrderedScheduler.newSchedulerBuilder()
              .name("etcd-watcher-scheduler")
              .numThreads(Runtime.getRuntime().availableProcessors())
              .build();
          this.watchExecutor = this.scheduler.chooseThread();
      }
  
      public synchronized boolean isClosed() {
          return closeFuture != null;
      }
  
      public CompletableFuture<EtcdWatcher> watch(ByteSequence key,
                                                  BiConsumer<io.etcd.jetcd.watch.WatchResponse, Throwable> consumer) {
          return watch(key, WatchOption.DEFAULT, consumer);
      }
  
      public CompletableFuture<EtcdWatcher> watch(ByteSequence key,
                                                  WatchOption watchOption,
                                                  BiConsumer<io.etcd.jetcd.watch.WatchResponse, Throwable> consumer) {
          return CompletableFuture.supplyAsync(() -> {
              if (isClosed()) {
                  throw EtcdExceptionFactory.newClosedWatchClientException();
              }
  
              EtcdWatcher watcher = new EtcdWatcher(key, watchOption, scheduler.chooseThread(), this);
              watcher.addConsumer(consumer);
              pendingWatchers.add(watcher);
              if (pendingWatchers.size() == 1) {
                  WatchRequest request = toWatchCreateRequest(watcher);
                  getGrpcWatchStreamObserver().onNext(request);
              }
              return watcher;
          }, watchExecutor);
      }
  
      // notifies all watchers about a exception. it doesn't close watchers.
      // it is the responsibility of user to close watchers.
      private void notifyWatchers(EtcdException e) {
          WatchResponseWithError wre = new WatchResponseWithError(e);
          this.pendingWatchers.forEach(watcher -> watcher.notifyWatchResponse(wre));
          this.pendingWatchers.clear();
          this.watchers.values().forEach(watcher -> watcher.notifyWatchResponse(wre));
          this.watchers.clear();
      }
  
      public CompletableFuture<Void> unwatch(EtcdWatcher watcher) {
          return CompletableFuture.runAsync(() -> cancelWatcher(watcher.getWatchID()), watchExecutor);
      }
  
      private void cancelWatcher(long watchID) {
          if (isClosed()) {
              return;
          }
  
          if (cancelSet.contains(watchID)) {
              return;
          }
  
          watchers.remove(watchID);
          cancelSet.add(watchID);
  
          WatchCancelRequest watchCancelRequest = WatchCancelRequest.newBuilder()
              .setWatchId(watchID)
              .build();
          WatchRequest cancelRequest = WatchRequest.newBuilder()
              .setCancelRequest(watchCancelRequest)
              .build();
          getGrpcWatchStreamObserver().onNext(cancelRequest);
      }
  
      public CompletableFuture<Void> closeAsync() {
          CompletableFuture<Void> future;
          synchronized (this) {
              if (null == closeFuture) {
                  log.info("Closing watch client");
                  closeFuture = CompletableFuture.runAsync(() -> {
                      notifyWatchers(newClosedWatchClientException());
                      closeGrpcWatchStreamObserver();
                  }, watchExecutor);
              }
              future = closeFuture;
          }
          return future.whenComplete((ignored, cause) -> {
              this.scheduler.shutdown();
          });
      }
  
      @Override
      public void close() {
          try {
              FutureUtils.result(closeAsync());
          } catch (Exception e) {
              log.warn("Encountered exceptions on closing watch client", e);
          }
          this.scheduler.forceShutdown(10, TimeUnit.SECONDS);
      }
  
      private StreamObserver<WatchResponse> createWatchStreamObserver() {
          return new StreamObserver<WatchResponse>() {
              @Override
              public void onNext(WatchResponse watchResponse) {
                  if (isClosed()) {
                      return;
                  }
                  watchExecutor.submit(() -> processWatchResponse(watchResponse));
              }
  
              @Override
              public void onError(Throwable t) {
                  if (isClosed()) {
                      return;
                  }
                  watchExecutor.submit(() -> processError(t));
              }
  
              @Override
              public void onCompleted() {
              }
          };
      }
  
      private void processWatchResponse(WatchResponse watchResponse) {
          // prevents grpc on sending watchResponse to a closed watch client.
          if (isClosed()) {
              return;
          }
  
          if (watchResponse.getCreated()) {
              processCreate(watchResponse);
          } else if (watchResponse.getCanceled()) {
              processCanceled(watchResponse);
          } else {
              processEvents(watchResponse);
          }
      }
  
      private void processError(Throwable t) {
          // prevents grpc on sending error to a closed watch client.
          if (this.isClosed()) {
              return;
          }
  
          Status status = Status.fromThrowable(t);
          if (this.isHaltError(status) || this.isNoLeaderError(status)) {
              this.notifyWatchers(toEtcdException(status));
              this.closeGrpcWatchStreamObserver();
              this.cancelSet.clear();
              return;
          }
          // resume with a delay; avoiding immediate retry on a long connection downtime.
          scheduler.schedule(this::resume, 500, TimeUnit.MILLISECONDS);
      }
  
      private void resume() {
          this.closeGrpcWatchStreamObserver();
          this.cancelSet.clear();
          this.resumeWatchers();
      }
  
      private synchronized StreamObserver<WatchRequest> getGrpcWatchStreamObserver() {
          if (this.grpcWatchStreamObserver == null) {
              this.grpcWatchStreamObserver = this.stub.watch(this.createWatchStreamObserver());
          }
          return this.grpcWatchStreamObserver;
      }
  
      // closeGrpcWatchStreamObserver closes the underlying grpc watch stream.
      private void closeGrpcWatchStreamObserver() {
          if (this.grpcWatchStreamObserver == null) {
              return;
          }
          this.grpcWatchStreamObserver.onCompleted();
          this.grpcWatchStreamObserver = null;
      }
  
      private void processCreate(WatchResponse response) {
          EtcdWatcher watcher = this.pendingWatchers.poll();
  
          this.sendNextWatchCreateRequest();
  
          if (watcher == null) {
              // shouldn't happen
              // may happen due to duplicate watch create responses.
              log.warn("Watch client receives watch create response but find no corresponding watcher");
              return;
          }
  
          if (watcher.isClosed()) {
              return;
          }
  
          if (response.getWatchId() == -1) {
              watcher.notifyWatchResponse(new WatchResponseWithError(
                  newEtcdException(ErrorCode.INTERNAL, "etcd server failed to create watch id")));
              return;
          }
  
          if (watcher.getRevision() == 0) {
              watcher.setRevision(response.getHeader().getRevision());
          }
  
          watcher.setWatchID(response.getWatchId());
          this.watchers.put(watcher.getWatchID(), watcher);
      }
  
      /**
       * chooses the next resuming watcher to register with the grpc stream.
       */
      private Optional<WatchRequest> nextResume() {
          EtcdWatcher pendingWatcher = this.pendingWatchers.peek();
          if (pendingWatcher != null) {
              return Optional.of(this.toWatchCreateRequest(pendingWatcher));
          }
          return Optional.empty();
      }
  
      private void sendNextWatchCreateRequest() {
          this.nextResume().ifPresent(
              (nextWatchRequest -> this.getGrpcWatchStreamObserver().onNext(nextWatchRequest)));
      }
  
      private void processEvents(WatchResponse response) {
          EtcdWatcher watcher = this.watchers.get(response.getWatchId());
          if (watcher == null) {
              // cancel server side watcher.
              this.cancelWatcher(response.getWatchId());
              return;
          }
  
          if (response.getCompactRevision() != 0) {
              watcher.notifyWatchResponse(new WatchResponseWithError(
                  EtcdExceptionFactory
                      .newCompactedException(response.getCompactRevision())));
              return;
          }
  
          if (response.getEventsCount() == 0) {
              watcher.setRevision(response.getHeader().getRevision());
              return;
          }
  
          watcher.notifyWatchResponse(new WatchResponseWithError(new io.etcd.jetcd.watch.WatchResponse(response)));
          watcher.setRevision(
              response
                  .getEvents(response.getEventsCount() - 1)
                  .getKv().getModRevision() + 1);
      }
  
      private void resumeWatchers() {
          this.watchers.values().forEach(watcher -> {
              if (watcher.isClosed()) {
                  return;
              }
              watcher.setWatchID(-1);
              this.pendingWatchers.add(watcher);
          });
  
          this.watchers.clear();
  
          this.sendNextWatchCreateRequest();
      }
  
      private void processCanceled(WatchResponse response) {
          EtcdWatcher watcher = this.watchers.get(response.getWatchId());
          this.cancelSet.remove(response.getWatchId());
          if (watcher == null) {
              return;
          }
          String reason = response.getCancelReason();
          if (Strings.isNullOrEmpty(reason)) {
              watcher.notifyWatchResponse(new WatchResponseWithError(newEtcdException(
                  ErrorCode.OUT_OF_RANGE,
                  "etcdserver: mvcc: required revision is a future revision"))
              );
  
          } else {
              watcher.notifyWatchResponse(
                  new WatchResponseWithError(newEtcdException(ErrorCode.FAILED_PRECONDITION, reason)));
          }
      }
  
      private static boolean isNoLeaderError(Status status) {
          return status.getCode() == Code.UNAVAILABLE
              && "etcdserver: no leader".equals(status.getDescription());
      }
  
      private static boolean isHaltError(Status status) {
          // Unavailable codes mean the system will be right back.
          // (e.g., can't connect, lost leader)
          // Treat Internal codes as if something failed, leaving the
          // system in an inconsistent state, but retrying could make progress.
          // (e.g., failed in middle of send, corrupted frame)
          return status.getCode() != Code.UNAVAILABLE && status.getCode() != Code.INTERNAL;
      }
  
      private static WatchRequest toWatchCreateRequest(EtcdWatcher watcher) {
          ByteString key = UnsafeByteOperations.unsafeWrap(watcher.getKey().getBytes());
          WatchOption option = watcher.getWatchOption();
          WatchCreateRequest.Builder builder = WatchCreateRequest.newBuilder()
              .setKey(key)
              .setPrevKv(option.isPrevKV())
              .setProgressNotify(option.isProgressNotify())
              .setStartRevision(watcher.getRevision());
  
          option.getEndKey()
              .ifPresent(endKey -> builder.setRangeEnd(UnsafeByteOperations.unsafeWrap(endKey.getBytes())));
  
          if (option.isNoDelete()) {
              builder.addFilters(WatchCreateRequest.FilterType.NODELETE);
          }
  
          if (option.isNoPut()) {
              builder.addFilters(WatchCreateRequest.FilterType.NOPUT);
          }
  
          return WatchRequest.newBuilder().setCreateRequest(builder).build();
      }
  
  
  }
  