/*
 * Licensed to Metamarkets Group Inc. (Metamarkets) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. Metamarkets 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 io.druid.server.coordinator;

import com.google.common.base.Preconditions;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.inject.Inject;
import com.metamx.emitter.EmittingLogger;
import io.druid.client.ServerInventoryView;
import io.druid.client.ServerView;
import io.druid.java.util.common.concurrent.Execs;
import io.druid.concurrent.LifecycleLock;
import io.druid.guice.ManageLifecycle;
import io.druid.java.util.common.ISE;
import io.druid.java.util.common.lifecycle.LifecycleStart;
import io.druid.java.util.common.lifecycle.LifecycleStop;
import io.druid.server.coordination.DruidServerMetadata;
import io.druid.server.coordinator.cost.ClusterCostCache;
import io.druid.timeline.DataSegment;

import java.util.concurrent.CancellationException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

@ManageLifecycle
public class CachingCostBalancerStrategyFactory implements BalancerStrategyFactory
{
  private static final EmittingLogger LOG = new EmittingLogger(CachingCostBalancerStrategyFactory.class);

  private final ServerInventoryView serverInventoryView;
  private final LifecycleLock lifecycleLock = new LifecycleLock();
  /** Must be single-threaded, because {@link ClusterCostCache.Builder} and downstream builders are not thread-safe */
  private final ExecutorService executor = Execs.singleThreaded("CachingCostBalancerStrategy-executor");
  private final ClusterCostCache.Builder clusterCostCacheBuilder = ClusterCostCache.builder();
  private volatile boolean initialized = false;

  @Inject
  public CachingCostBalancerStrategyFactory(ServerInventoryView serverInventoryView)
  {
    this.serverInventoryView = Preconditions.checkNotNull(serverInventoryView);
  }

  @LifecycleStart
  public void start()
  {
    if (!lifecycleLock.canStart()) {
      throw new ISE("CachingCostBalancerStrategyFactory can not be started");
    }
    try {
      serverInventoryView.registerSegmentCallback(
          executor,
          new ServerView.SegmentCallback()
          {
            @Override
            public ServerView.CallbackAction segmentAdded(DruidServerMetadata server, DataSegment segment)
            {
              clusterCostCacheBuilder.addSegment(server.getName(), segment);
              return ServerView.CallbackAction.CONTINUE;
            }

            @Override
            public ServerView.CallbackAction segmentRemoved(DruidServerMetadata server, DataSegment segment)
            {
              clusterCostCacheBuilder.removeSegment(server.getName(), segment);
              return ServerView.CallbackAction.CONTINUE;
            }

            @Override
            public ServerView.CallbackAction segmentViewInitialized()
            {
              initialized = true;
              return ServerView.CallbackAction.CONTINUE;
            }
          }
      );

      serverInventoryView.registerServerRemovedCallback(
          executor,
          server -> {
            clusterCostCacheBuilder.removeServer(server.getName());
            return ServerView.CallbackAction.CONTINUE;
          }
      );

      lifecycleLock.started();
    }
    finally {
      lifecycleLock.exitStart();
    }
  }

  @LifecycleStop
  public void stop()
  {
    if (!lifecycleLock.canStop()) {
      throw new ISE("CachingCostBalancerStrategyFactory can not be stopped");
    }
    executor.shutdownNow();
    // Not calling lifecycleLock.exitStop() because CachingCostBalancerStrategyFactory is not recycleable
  }

  @Override
  public BalancerStrategy createBalancerStrategy(final ListeningExecutorService exec)
  {
    if (!lifecycleLock.awaitStarted()) {
      throw new ISE("CachingCostBalancerStrategyFactory is not started");
    }
    if (initialized) {
      try {
        CompletableFuture<CachingCostBalancerStrategy> future = CompletableFuture.supplyAsync(
            () -> new CachingCostBalancerStrategy(clusterCostCacheBuilder.build(), exec),
            executor
        );
        try {
          return future.get(1, TimeUnit.SECONDS);
        }
        catch (CancellationException e) {
          LOG.error("CachingCostBalancerStrategy creation has been cancelled");
        }
        catch (ExecutionException e) {
          LOG.error(e, "Failed to create CachingCostBalancerStrategy");
        }
        catch (TimeoutException e) {
          LOG.error("CachingCostBalancerStrategy creation took more than 1 second!");
        }
        catch (InterruptedException e) {
          LOG.error("CachingCostBalancerStrategy creation has been interrupted");
          Thread.currentThread().interrupt();
        }
      }
      catch (RejectedExecutionException e) {
        LOG.error("CachingCostBalancerStrategy creation has been rejected");
      }
    } else {
      LOG.error("CachingCostBalancerStrategy could not be created as serverView is not initialized yet");
    }
    LOG.info("Fallback to CostBalancerStrategy");
    return new CostBalancerStrategy(exec);
  }
}

@@ -838,16 +838,124 @@
 nator;%0A%0A
+import com.fasterxml.jackson.annotation.JacksonInject;%0Aimport com.fasterxml.jackson.annotation.JsonCreator;%0A
 import c
@@ -1062,41 +1062,8 @@
 ce;%0A
-import com.google.inject.Inject;%0A
 impo
@@ -1292,29 +1292,28 @@
 uid.
-guice.ManageLifecycle
+java.util.common.ISE
 ;%0Aim
@@ -1335,35 +1335,51 @@
 ava.util.common.
-ISE
+lifecycle.Lifecycle
 ;%0Aimport io.drui
@@ -1971,25 +1971,8 @@
 n;%0A%0A
-@ManageLifecycle%0A
 publ
@@ -2654,22 +2654,27 @@
 se;%0A%0A  @
-Inject
+JsonCreator
 %0A  publi
@@ -2710,16 +2710,38 @@
 Factory(
+%0A      @JacksonInject 
 ServerIn
@@ -2771,17 +2771,79 @@
 toryView
-)
+,%0A      @JacksonInject Lifecycle lifecycle%0A  ) throws Exception
 %0A  %7B%0A   
@@ -2919,16 +2919,214 @@
 yView);%0A
+    // Adding to lifecycle dynamically because couldn't use @ManageLifecycle on the class,%0A    // see https://github.com/druid-io/druid/issues/4980%0A    lifecycle.addMaybeStartManagedInstance(this);%0A
   %7D%0A%0A  @
@@ -5023,16 +5023,35 @@
 Started(
+1, TimeUnit.MINUTES
 )) %7B%0A   

Fix CachingCostBalancerStrategyFactory's constructor (#4974)