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.stream.storage.impl.cluster;
  
  import java.time.Duration;
  import java.util.Set;
  import java.util.concurrent.Semaphore;
  import java.util.concurrent.TimeUnit;
  import lombok.AccessLevel;
  import lombok.Getter;
  import lombok.extern.slf4j.Slf4j;
  import org.apache.bookkeeper.discover.RegistrationClient;
  import org.apache.bookkeeper.discover.RegistrationClient.RegistrationListener;
  import org.apache.bookkeeper.net.BookieId;
  import org.apache.bookkeeper.stream.proto.cluster.ClusterAssignmentData;
  import org.apache.bookkeeper.stream.proto.cluster.ClusterMetadata;
  import org.apache.bookkeeper.stream.storage.api.cluster.ClusterControllerLeader;
  import org.apache.bookkeeper.stream.storage.api.cluster.ClusterMetadataStore;
  import org.apache.bookkeeper.stream.storage.impl.sc.StorageContainerController;
  import org.apache.bookkeeper.versioning.Versioned;
  
  /**
   * A default implementation of {@link ClusterControllerLeader}.
   */
  @Slf4j
  public class ClusterControllerLeaderImpl implements ClusterControllerLeader, RegistrationListener {
  
      // the metadata store for reading and writing cluster metadata.
      private final ClusterMetadataStore clusterMetadataStore;
  
      // the controller logic for assigning containers
      private final StorageContainerController scController;
  
      // permits that the controller leader can perform server changes.
      @Getter(AccessLevel.PACKAGE)
      private final Semaphore performServerChangesPermits;
  
      // registration client that watch/unwatch registered servers.
      private final RegistrationClient regClient;
  
      // keep a reference to a set of available servers
      private volatile Set<BookieId> availableServers;
  
      // variables for suspending controller
      @Getter(AccessLevel.PACKAGE)
      private final Object suspensionLock = new Object();
      private volatile boolean suspended = false;
  
      // last successful assignment happened at (timestamp)
      @Getter(AccessLevel.PACKAGE)
      private long lastSuccessfulAssignmentAt;
  
      // the min interval that controller is scheduled to assign containers
      private final Duration scheduleDuration;
  
      ClusterControllerLeaderImpl(ClusterMetadataStore clusterMetadataStore,
                                  StorageContainerController scController,
                                  RegistrationClient regClient,
                                  Duration scheduleDuration) {
          this.clusterMetadataStore = clusterMetadataStore;
          this.scController = scController;
          this.regClient = regClient;
          this.performServerChangesPermits = new Semaphore(0);
          this.lastSuccessfulAssignmentAt = -1L;
          this.scheduleDuration = scheduleDuration;
      }
  
      /**
       * Suspend the controller if the leader disconnects from zookeeper.
       */
      @Override
      public void suspend() {
          synchronized (suspensionLock) {
              suspended = true;
              suspensionLock.notifyAll();
          }
      }
  
      boolean isSuspended() {
          return suspended;
      }
  
      /**
       * Resume the controller.
       */
      @Override
      public void resume() {
          synchronized (suspensionLock) {
              suspended = false;
              suspensionLock.notifyAll();
          }
      }
  
      @Override
      public void processAsLeader() throws Exception {
          log.info("Become controller leader to monitor servers for assigning storage containers.");
  
          performServerChangesPermits.release();
  
          // monitoring the servers
          try {
              this.regClient.watchWritableBookies(this).get();
          } catch (Exception e) {
              log.warn("Controller leader fails to watch servers : {}, giving up leadership", e.getMessage());
              throw e;
          }
  
          // the leader is looping here to react to changes
          while (true) {
              try {
                  checkSuspension();
  
                  processServerChange();
              } catch (InterruptedException ie) {
                  log.warn("Controller leader is interrupted, giving up leadership");
  
                  // stop monitoring the servers
                  this.regClient.unwatchWritableBookies(this);
                  throw ie;
              } catch (Exception e) {
                  // if the leader is suspended due to losing connection to zookeeper
                  // we don't give leadership until it becomes leader again or being interrupted by curator
                  if (!suspended) {
                      log.warn("Controller leader encountered exceptions on processing server changes,"
                          + " giving up leadership");
  
                      // stop monitoring the servers
                      this.regClient.unwatchWritableBookies(this);
                      throw e;
                  }
              }
          }
      }
  
      private void checkSuspension() throws InterruptedException {
          synchronized (suspensionLock) {
              while (suspended) {
                  log.info("Controller leader is suspended, waiting for to be resumed");
                  suspensionLock.wait();
                  log.info("Controller leader is woke up from suspension");
              }
          }
      }
  
      private void processServerChange() throws InterruptedException {
          // check if the leader can perform server changes
          performServerChangesPermits.acquire();
  
          long elapsedMs = System.currentTimeMillis() - lastSuccessfulAssignmentAt;
          long remainingMs = scheduleDuration.toMillis() - elapsedMs;
          if (remainingMs > 0) {
              log.info("Waiting {} milliseconds for controller to assign containers", remainingMs);
              TimeUnit.MILLISECONDS.sleep(remainingMs);
          }
  
          // now, the controller has permits and meet the time requirement for assigning containers.
          performServerChangesPermits.drainPermits();
  
          Set<BookieId> availableServersSnapshot = availableServers;
          if (null == availableServersSnapshot || availableServersSnapshot.isEmpty()) {
              // haven't received any servers from registration service, wait for 200ms and retry.
              if (lastSuccessfulAssignmentAt < 0) {
                  log.info("No servers is alive yet. Backoff 200ms and retry.");
                  TimeUnit.MILLISECONDS.sleep(200);
                  performServerChangesPermits.release();
                  return;
              } else {
                  // there was already a successful assignment but all the servers are gone.
                  // it can be a registration service issue, so don't attempt to reassign the containers.
                  // return here direct to wait next server change
                  return;
              }
          }
  
          ClusterMetadata clusterMetadata = clusterMetadataStore.getClusterMetadata();
          ClusterAssignmentData currentState = clusterMetadataStore.getClusterAssignmentData();
  
          // servers are changed, process the change.
          ClusterAssignmentData newState = scController.computeIdealState(
              clusterMetadata,
              currentState,
              availableServersSnapshot);
  
          if (newState.equals(currentState)) {
              // no assignment state is changed, so do nothing
              if (log.isDebugEnabled()) {
                  log.debug("Assignment state is unchanged - {}", newState);
              }
          } else {
              // update the assignment state
              lastSuccessfulAssignmentAt = System.currentTimeMillis();
              clusterMetadataStore.updateClusterAssignmentData(newState);
          }
      }
  
      @Override
      public void onBookiesChanged(Versioned<Set<BookieId>> bookies) {
          log.info("Cluster topology is changed - new cluster : {}", bookies);
          // when bookies are changed, notify the leader to take actions
          this.availableServers = bookies.getValue();
          performServerChangesPermits.release();
      }
  }
  