<NME> WorkerModule.java
<BEF> /*
 * Copyright 2018 The Data Transfer Project Authors.
 *
 * 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
 *
 * https://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.datatransferproject.transfer;

import static com.google.common.collect.MoreCollectors.onlyElement;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableList;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.util.concurrent.AbstractScheduledService.Scheduler;
import com.google.inject.Provides;
import com.google.inject.Singleton;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.Set;
import java.util.concurrent.TimeUnit;
import org.datatransferproject.api.launcher.ExtensionContext;
import org.datatransferproject.api.launcher.Monitor;
import org.datatransferproject.config.FlagBindingModule;
import org.datatransferproject.security.AsymmetricKeyGenerator;
import org.datatransferproject.security.SymmetricKeyGenerator;
import org.datatransferproject.spi.cloud.extension.CloudExtension;
import org.datatransferproject.api.launcher.MetricRecorder;
import org.datatransferproject.api.launcher.Monitor;
import org.datatransferproject.config.FlagBindingModule;
import org.datatransferproject.launcher.metrics.LoggingDtpInternalMetricRecorder;
import org.datatransferproject.launcher.metrics.ServiceAwareMetricRecorder;
import org.datatransferproject.security.SymmetricKeyGenerator;
import org.datatransferproject.spi.cloud.extension.CloudExtension;
import org.datatransferproject.spi.transfer.security.PublicKeySerializer;
import org.datatransferproject.types.transfer.retry.RetryStrategyLibrary;

final class WorkerModule extends FlagBindingModule {

  private final CloudExtension cloudExtension;
import org.datatransferproject.transfer.copier.InMemoryDataCopierClassLoader;
import org.datatransferproject.types.transfer.retry.RetryStrategyLibrary;
import org.datatransferproject.types.transfer.serviceconfig.TransferServiceConfig;

final class WorkerModule extends FlagBindingModule {

  private final CloudExtension cloudExtension;
  private final ExtensionContext context;
  private final List<TransferExtension> transferExtensions;
  private final SecurityExtension securityExtension;
  private final IdempotentImportExecutor idempotentImportExecutor;
  private final SymmetricKeyGenerator symmetricKeyGenerator;
  private final JobHooks jobHooks;
  private final TransferCompatibilityProvider compatibilityProvider;

  WorkerModule(
      ExtensionContext context,
      CloudExtension cloudExtension,
      List<TransferExtension> transferExtensions,
      SecurityExtension securityExtension,
      IdempotentImportExecutor idempotentImportExecutor,
      SymmetricKeyGenerator symmetricKeyGenerator,
      JobHooks jobHooks,
      TransferCompatibilityProvider transferCompatibilityProvider) {
    this.cloudExtension = cloudExtension;
    this.context = context;
    this.transferExtensions = transferExtensions;
    this.securityExtension = securityExtension;
    this.idempotentImportExecutor = idempotentImportExecutor;
    this.symmetricKeyGenerator = symmetricKeyGenerator;
    this.jobHooks = jobHooks;
    this.compatibilityProvider = transferCompatibilityProvider;
  }

  @VisibleForTesting
  static TransferExtension findTransferExtension(
      ImmutableList<TransferExtension> transferExtensions, String service) {
    try {
      return transferExtensions
          .stream()
          .filter(ext -> ext.getServiceId().toLowerCase().equals(service.toLowerCase()))
          .collect(onlyElement());
    } catch (IllegalArgumentException e) {
      throw new IllegalStateException(
          "Found multiple transfer extensions for service " + service, e);
    } catch (NoSuchElementException e) {
      throw new IllegalStateException(
          "Did not find a valid transfer extension for service " + service, e);
    }
  }

  @Override
  protected void configure() {
    // binds flags from ExtensionContext to @Named annotations
    bindFlags(context);

    bind(JobHooks.class).toInstance(jobHooks);
    bind(InMemoryDataCopier.class).to(InMemoryDataCopierClassLoader.load());
    getMonitor()
        .info(() -> "Using InMemoryDataCopier: " + InMemoryDataCopierClassLoader.load().getName());

    bind(ObjectMapper.class).toInstance(context.getTypeManager().getMapper());

    // Ensure a DtpInternalMetricRecorder exists
    LoggingDtpInternalMetricRecorder.registerRecorderIfNeeded(context);
    bind(DtpInternalMetricRecorder.class)
        .toInstance(context.getService(DtpInternalMetricRecorder.class));
  }

  @Provides
  Exporter getExporter(ImmutableList<TransferExtension> transferExtensions) {
    TransferExtension extension =
        findTransferExtension(transferExtensions, JobMetadata.getExportService());
    extension.initialize(context);
    return extension.getExporter(JobMetadata.getDataType());
  }

  AuthDataDecryptService getAuthDataDecryptService() {
    return securityExtension.getDecryptService();
  Importer getImporter(ImmutableList<TransferExtension> transferExtensions) {
    TransferExtension extension =
        findTransferExtension(transferExtensions, JobMetadata.getImportService());
    extension.initialize(context);
    return extension.getImporter(JobMetadata.getDataType());
  }


  @Provides
  @Singleton
  AppCredentialStore getBucketStore() {
    return cloudExtension.getAppCredentialStore();
  }

  @Provides
  @Singleton
  Exporter getExporter(ImmutableList<TransferExtension> transferExtensions) {
    TransferExtension extension =
        findTransferExtension(transferExtensions, JobMetadata.getExportService());
    DelegatingExtensionContext serviceSpecificContext = new DelegatingExtensionContext(context);
    serviceSpecificContext.registerOverrideService(
        MetricRecorder.class,
        new ServiceAwareMetricRecorder(
            extension.getServiceId(),
            context.getService(DtpInternalMetricRecorder.class)));
    serviceSpecificContext.registerOverrideService(
        TransferServiceConfig.class,
        getTransferServiceConfig(extension));
    extension.initialize(serviceSpecificContext);

    return compatibilityProvider.getCompatibleExporter(extension, JobMetadata.getDataType());
  }

  @Provides
  @Singleton
  Importer getImporter(ImmutableList<TransferExtension> transferExtensions) {
    TransferExtension extension =
        findTransferExtension(transferExtensions, JobMetadata.getImportService());
    DelegatingExtensionContext serviceSpecificContext = new DelegatingExtensionContext(context);
    serviceSpecificContext.registerOverrideService(
        MetricRecorder.class,
        new ServiceAwareMetricRecorder(
            extension.getServiceId(),
            context.getService(DtpInternalMetricRecorder.class)));
    serviceSpecificContext.registerOverrideService(
        TransferServiceConfig.class,
        getTransferServiceConfig(extension));
    extension.initialize(serviceSpecificContext);
    return compatibilityProvider.getCompatibleImporter(extension, JobMetadata.getDataType());
  }

  @Provides
  @Singleton
  ImmutableList<TransferExtension> getTransferExtensions() {
    return ImmutableList.copyOf(transferExtensions);
  }


  @Provides
  @Singleton
  RetryStrategyLibrary getRetryStrategyLibrary() {
    return context.getSetting("retryLibrary", null);
  }

  @Provides
  @Singleton
  Scheduler getScheduler() {
    // TODO: parse a Duration from the settings
    long interval = context.getSetting("pollInterval", 2000); // Default: poll every 2s
    return AbstractScheduledService.Scheduler.newFixedDelaySchedule(
        0, interval, TimeUnit.MILLISECONDS);
  }

  @Provides
  @Singleton
  @Annotations.CancelScheduler
  Scheduler getCancelCheckingScheduler() {
    // TODO: parse a Duration from the settings
    long interval = context.getSetting("cancelCheckPollInterval", 60000); // Default: poll every 1m
    return AbstractScheduledService.Scheduler.newFixedDelaySchedule(
        0, interval, TimeUnit.MILLISECONDS);
  }

  @Provides
  @Singleton
  Monitor getMonitor() {
    return context.getMonitor();
  }

  @Provides
  @Singleton
  ExtensionContext getContext() {
    return context;
  }

  private TransferServiceConfig getTransferServiceConfig(TransferExtension ext) {
    String configFileName = "config/" + ext.getServiceId().toLowerCase() + ".yaml";
    InputStream inputStream = this.getClass().getClassLoader().getResourceAsStream(configFileName);
    getMonitor()
        .info(
            () ->
                format(
                    "Service %s has a config file: %s", ext.getServiceId(), (inputStream != null)));
    if (inputStream == null) {
      return TransferServiceConfig.getDefaultInstance();
    } else {
      try {
        return TransferServiceConfig.create(inputStream);
      } catch (IOException e) {
        throw new RuntimeException("Couldn't create config for " + ext.getServiceId(), e);
      }
    }
  }

  @Provides
  @Singleton
  public IdempotentImportExecutor getIdempotentImportExecutor() {
    return idempotentImportExecutor;
  }
}

<MSG> Add a Stats interface and framework (#687)

* Create the interfaces and framework a metrics framework

Re: #41 add a stats/metrics framework to DTP that can be used for both internal events as well extended for integrators to add their own stats.

<DFF> @@ -15,8 +15,6 @@
  */
 package org.datatransferproject.transfer;
 
-import static com.google.common.collect.MoreCollectors.onlyElement;
-
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
@@ -25,13 +23,13 @@ import com.google.common.util.concurrent.AbstractScheduledService;
 import com.google.common.util.concurrent.AbstractScheduledService.Scheduler;
 import com.google.inject.Provides;
 import com.google.inject.Singleton;
-import java.util.List;
-import java.util.NoSuchElementException;
-import java.util.Set;
-import java.util.concurrent.TimeUnit;
+import org.datatransferproject.api.launcher.DelegatingExtensionContext;
+import org.datatransferproject.api.launcher.DtpInternalMetricRecorder;
 import org.datatransferproject.api.launcher.ExtensionContext;
+import org.datatransferproject.api.launcher.MetricRecorder;
 import org.datatransferproject.api.launcher.Monitor;
 import org.datatransferproject.config.FlagBindingModule;
+import org.datatransferproject.launcher.metrics.ServiceAwareMetricRecorder;
 import org.datatransferproject.security.AsymmetricKeyGenerator;
 import org.datatransferproject.security.SymmetricKeyGenerator;
 import org.datatransferproject.spi.cloud.extension.CloudExtension;
@@ -45,6 +43,13 @@ import org.datatransferproject.spi.transfer.security.AuthDataDecryptService;
 import org.datatransferproject.spi.transfer.security.PublicKeySerializer;
 import org.datatransferproject.types.transfer.retry.RetryStrategyLibrary;
 
+import java.util.List;
+import java.util.NoSuchElementException;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+
+import static com.google.common.collect.MoreCollectors.onlyElement;
+
 final class WorkerModule extends FlagBindingModule {
 
   private final CloudExtension cloudExtension;
@@ -121,7 +126,13 @@ final class WorkerModule extends FlagBindingModule {
   Exporter getExporter(ImmutableList<TransferExtension> transferExtensions) {
     TransferExtension extension =
         findTransferExtension(transferExtensions, JobMetadata.getExportService());
-    extension.initialize(context);
+    DelegatingExtensionContext serviceSpecificContext = new DelegatingExtensionContext(context);
+    serviceSpecificContext.registerOverrideService(
+        MetricRecorder.class,
+        new ServiceAwareMetricRecorder(
+            extension.getServiceId(),
+            context.getService(DtpInternalMetricRecorder.class)));
+    extension.initialize(serviceSpecificContext);
     return extension.getExporter(JobMetadata.getDataType());
   }
 
@@ -130,7 +141,13 @@ final class WorkerModule extends FlagBindingModule {
   Importer getImporter(ImmutableList<TransferExtension> transferExtensions) {
     TransferExtension extension =
         findTransferExtension(transferExtensions, JobMetadata.getImportService());
-    extension.initialize(context);
+    DelegatingExtensionContext serviceSpecificContext = new DelegatingExtensionContext(context);
+    serviceSpecificContext.registerOverrideService(
+        MetricRecorder.class,
+        new ServiceAwareMetricRecorder(
+            extension.getServiceId(),
+            context.getService(DtpInternalMetricRecorder.class)));
+    extension.initialize(serviceSpecificContext);
     return extension.getImporter(JobMetadata.getDataType());
   }
 

{'addition_count': 25, 'commit_subject': 'Add a Stats interface and framework (#687)', 'deletion_count': 8, 'file_extension': '.java', 'lang': 'java', 'license': 'apache-2.0', 'repo_name': 'google/data-transfer-project'}