content: /**
   * Copyright 2019 Airbnb. Licensed under Apache-2.0. See LICENSE in the project root for license
   * information.
   */
  package com.airbnb.dynein.scheduler.modules;
  
  import com.airbnb.conveyor.async.AsyncSqsClient;
  import com.airbnb.conveyor.async.AsyncSqsClientFactory;
  import com.airbnb.conveyor.async.config.AsyncSqsClientConfiguration;
  import com.airbnb.conveyor.async.config.OverrideConfiguration;
  import com.airbnb.conveyor.async.config.RetryPolicyConfiguration;
  import com.airbnb.conveyor.async.metrics.AsyncConveyorMetrics;
  import com.airbnb.conveyor.async.metrics.NoOpConveyorMetrics;
  import com.airbnb.dynein.common.job.JacksonJobSpecTransformer;
  import com.airbnb.dynein.common.job.JobSpecTransformer;
  import com.airbnb.dynein.common.token.JacksonTokenManager;
  import com.airbnb.dynein.common.token.TokenManager;
  import com.airbnb.dynein.scheduler.Constants;
  import com.airbnb.dynein.scheduler.ManagedInboundJobQueueConsumer;
  import com.airbnb.dynein.scheduler.ScheduleManager;
  import com.airbnb.dynein.scheduler.ScheduleManagerFactory;
  import com.airbnb.dynein.scheduler.config.DynamoDBConfiguration;
  import com.airbnb.dynein.scheduler.config.SQSConfiguration;
  import com.airbnb.dynein.scheduler.config.SchedulerConfiguration;
  import com.airbnb.dynein.scheduler.config.WorkersConfiguration;
  import com.airbnb.dynein.scheduler.dynamodb.DynamoDBScheduleManagerFactory;
  import com.airbnb.dynein.scheduler.heartbeat.HeartbeatConfiguration;
  import com.airbnb.dynein.scheduler.heartbeat.HeartbeatManager;
  import com.airbnb.dynein.scheduler.metrics.Metrics;
  import com.airbnb.dynein.scheduler.metrics.NoOpMetricsImpl;
  import com.airbnb.dynein.scheduler.partition.K8SConsecutiveAllocationPolicy;
  import com.airbnb.dynein.scheduler.partition.PartitionPolicy;
  import com.airbnb.dynein.scheduler.partition.StaticAllocationPolicy;
  import com.airbnb.dynein.scheduler.worker.PartitionWorker;
  import com.airbnb.dynein.scheduler.worker.PartitionWorkerFactory;
  import com.airbnb.dynein.scheduler.worker.WorkerManager;
  import com.google.common.eventbus.EventBus;
  import com.google.common.util.concurrent.MoreExecutors;
  import com.google.common.util.concurrent.ThreadFactoryBuilder;
  import com.google.inject.AbstractModule;
  import com.google.inject.Provides;
  import com.google.inject.Singleton;
  import com.google.inject.name.Names;
  import java.time.Clock;
  import java.util.ArrayList;
  import java.util.List;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.Executors;
  import java.util.concurrent.ScheduledThreadPoolExecutor;
  import java.util.concurrent.ThreadPoolExecutor;
  import javax.inject.Named;
  import lombok.extern.slf4j.Slf4j;
  import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
  
  @Slf4j
  public class SchedulerModule extends AbstractModule {
    private final SchedulerConfiguration schedulerConfiguration;
    private final DynamoDBConfiguration dynamoDBConfiguration;
    private final SQSConfiguration sqsConfiguration;
    private final HeartbeatConfiguration heartbeatConfiguration;
    private final WorkersConfiguration workersConfiguration;
  
    public SchedulerModule(SchedulerConfiguration schedulerConfiguration) {
      this.schedulerConfiguration = schedulerConfiguration;
      this.dynamoDBConfiguration = schedulerConfiguration.getDynamoDb();
      this.sqsConfiguration = schedulerConfiguration.getSqs();
      this.heartbeatConfiguration = schedulerConfiguration.getHeartbeat();
      this.workersConfiguration = schedulerConfiguration.getWorkers();
    }
  
    @Override
    protected void configure() {
      bind(Clock.class).toInstance(Clock.systemUTC());
      bind(WorkersConfiguration.class).toInstance(workersConfiguration);
      bind(DynamoDBConfiguration.class).toInstance(dynamoDBConfiguration);
      bind(Integer.class)
          .annotatedWith(Names.named(Constants.MAX_PARTITIONS))
          .toInstance(schedulerConfiguration.getMaxPartitions());
      bind(String.class)
          .annotatedWith(Names.named(Constants.INBOUND_QUEUE_NAME))
          .toInstance(sqsConfiguration.getInboundQueueName());
  
      bind(DynamoDBScheduleManagerFactory.class).asEagerSingleton();
      bind(ScheduleManagerFactory.class).to(DynamoDBScheduleManagerFactory.class).asEagerSingleton();
      bind(ScheduleManager.class).toProvider(ScheduleManagerFactory.class).asEagerSingleton();
      bind(EventBus.class).toInstance(new EventBus());
  
      switch (workersConfiguration.getPartitionPolicy()) {
        case K8S:
          {
            bind(K8SConsecutiveAllocationPolicy.class).asEagerSingleton();
            bind(PartitionPolicy.class).to(K8SConsecutiveAllocationPolicy.class);
            break;
          }
        case STATIC:
          {
            bind(PartitionPolicy.class)
                .toInstance(
                    new StaticAllocationPolicy(workersConfiguration.getStaticPartitionList()));
            break;
          }
      }
      bind(Metrics.class).to(NoOpMetricsImpl.class).asEagerSingleton();
      bind(AsyncConveyorMetrics.class).to(NoOpConveyorMetrics.class).asEagerSingleton();
      bind(JobSpecTransformer.class).to(JacksonJobSpecTransformer.class).asEagerSingleton();
      bind(TokenManager.class).to(JacksonTokenManager.class).asEagerSingleton();
    }
  
    @Provides
    public DynamoDbAsyncClient providesDdbAsyncClient() {
      return DynamoDbAsyncClient.builder()
          .endpointOverride(dynamoDBConfiguration.getEndpoint())
          .build();
    }
  
    @Provides
    @Singleton
    public ManagedInboundJobQueueConsumer providesManagedInboundJobQueueConsumer(
        @Named(Constants.SQS_CONSUMER) AsyncSqsClient asyncClient, ScheduleManager scheduleManager) {
      return new ManagedInboundJobQueueConsumer(
          asyncClient,
          MoreExecutors.getExitingExecutorService(
              (ThreadPoolExecutor)
                  Executors.newFixedThreadPool(
                      1, new ThreadFactoryBuilder().setNameFormat("inbound-consumer-%d").build())),
          scheduleManager,
          sqsConfiguration.getInboundQueueName());
    }
  
    private AsyncSqsClientConfiguration buildBaseSqsConfig() {
      AsyncSqsClientConfiguration asyncSqsConfig =
          new AsyncSqsClientConfiguration(AsyncSqsClientConfiguration.Type.PRODUCTION);
  
      asyncSqsConfig.setRegion(sqsConfiguration.getRegion());
      asyncSqsConfig.setEndpoint(sqsConfiguration.getEndpoint());
      return asyncSqsConfig;
    }
  
    @Provides
    @Named(Constants.SQS_CONSUMER)
    AsyncSqsClientConfiguration provideConsumerSqsConfig() {
      return buildBaseSqsConfig();
    }
  
    @Provides
    @Named(Constants.SQS_PRODUCER)
    AsyncSqsClientConfiguration provideProducerSqsConfig() {
      AsyncSqsClientConfiguration asyncSqsConfig = buildBaseSqsConfig();
  
      RetryPolicyConfiguration retryConfig = new RetryPolicyConfiguration();
      retryConfig.setPolicy(RetryPolicyConfiguration.Policy.CUSTOM);
      retryConfig.setCondition(RetryPolicyConfiguration.Condition.DEFAULT);
      retryConfig.setBackOff(RetryPolicyConfiguration.BackOff.EQUAL_JITTER);
      retryConfig.setMaximumBackoffTime(sqsConfiguration.getTimeouts().getMaxBackoffTime());
      retryConfig.setBaseDelay(sqsConfiguration.getTimeouts().getBaseDelay());
  
      OverrideConfiguration overrideConfig = new OverrideConfiguration();
      overrideConfig.setApiCallAttemptTimeout(
          sqsConfiguration.getTimeouts().getApiCallAttemptTimeout());
      overrideConfig.setApiCallTimeout(sqsConfiguration.getTimeouts().getApiCallTimeout());
      overrideConfig.setRetryPolicyConfiguration(retryConfig);
      asyncSqsConfig.setOverrideConfiguration(overrideConfig);
      return asyncSqsConfig;
    }
  
    @Provides
    @Singleton
    @Named(Constants.SQS_CONSUMER)
    public AsyncSqsClient provideAsyncConsumer(
        AsyncSqsClientFactory asyncSqsClientFactory,
        @Named(Constants.SQS_CONSUMER) AsyncSqsClientConfiguration asyncSqsClientConfiguration) {
      return asyncSqsClientFactory.create(asyncSqsClientConfiguration);
    }
  
    @Provides
    @Singleton
    @Named(Constants.SQS_PRODUCER)
    public AsyncSqsClient provideAsyncProducer(
        AsyncSqsClientFactory asyncSqsClientFactory,
        @Named(Constants.SQS_PRODUCER) AsyncSqsClientConfiguration asyncSqsClientConfiguration) {
      return asyncSqsClientFactory.create(asyncSqsClientConfiguration);
    }
  
    @Provides
    @Singleton
    public HeartbeatManager provideHeartBeat(EventBus eventBus, Clock clock) {
      HeartbeatManager heartbeatManager =
          new HeartbeatManager(
              new ConcurrentHashMap<>(),
              MoreExecutors.getExitingScheduledExecutorService(
                  (ScheduledThreadPoolExecutor)
                      Executors.newScheduledThreadPool(
                          1,
                          new ThreadFactoryBuilder().setNameFormat("heartbeat-manager-%d").build())),
              eventBus,
              clock,
              heartbeatConfiguration);
      eventBus.register(heartbeatManager);
      return heartbeatManager;
    }
  
    @Provides
    @Singleton
    public WorkerManager provideWorkerManager(
        PartitionWorkerFactory factory, EventBus eventBus, Metrics metrics, PartitionPolicy policy) {
  
      List<PartitionWorker> partitionWorkers = new ArrayList<>();
      for (int i = 0; i < workersConfiguration.getNumberOfWorkers(); i++) {
        partitionWorkers.add(factory.get(i));
      }
  
      WorkerManager manager =
          new WorkerManager(
              partitionWorkers,
              factory,
              MoreExecutors.getExitingScheduledExecutorService(
                  (ScheduledThreadPoolExecutor)
                      Executors.newScheduledThreadPool(
                          1, new ThreadFactoryBuilder().setNameFormat("worker-manager-%d").build())),
              new ConcurrentHashMap<>(),
              policy,
              metrics);
  
      eventBus.register(manager);
      return manager;
    }
  }
  