Collecting response code for Dataproc operations

Updated

Updated

Updated

Updated

Updated

Updated

Updated

Author: Komal Yadav

cdap-runtime-ext-dataproc/src/main/java/io/cdap/cdap/runtime/spi/common/DataprocMetric.java

@@ -32,14 +32,17 @@ public class DataprocMetric {
   private final LaunchMode launchMode;
   @Nullable
   private final String imageVersion;
+  @Nullable
+  private final String opType;
 
   private DataprocMetric(String metricName, String region, @Nullable Exception exception,
-      @Nullable LaunchMode launchMode, @Nullable String imageVersion) {
+      @Nullable LaunchMode launchMode, @Nullable String imageVersion, @Nullable String opType) {
     this.metricName = metricName;
     this.region = region;
     this.exception = exception;
     this.launchMode = launchMode;
     this.imageVersion = imageVersion;
+    this.opType = opType;
   }
 
   public String getMetricName() {
@@ -72,6 +75,11 @@ public LaunchMode getLaunchMode() {
     return launchMode;
   }
 
+  @Nullable
+  public String getOpType() {
+    return opType;
+  }
+
   /**
    * Returns a builder to create a DataprocMetric.
    *
@@ -93,6 +101,8 @@ public static class Builder {
     private Exception exception;
     @Nullable
     private LaunchMode launchMode;
+    @Nullable
+    private String opType;
 
     private Builder(String metricName) {
       this.metricName = metricName;
@@ -118,6 +128,11 @@ public Builder setImageVersion(String imageVersion) {
       return this;
     }
 
+    public Builder setOpType(@Nullable String opType) {
+      this.opType = opType;
+      return this;
+    }
+
     /**
      * Returns a DataprocMetric.
      *
@@ -128,7 +143,7 @@ public DataprocMetric build() {
         // region should always be set unless there is a bug in the code
         throw new IllegalStateException("Dataproc metric is missing the region");
       }
-      return new DataprocMetric(metricName, region, exception, launchMode, imageVersion);
+      return new DataprocMetric(metricName, region, exception, launchMode, imageVersion, opType);
     }
   }
 }

cdap-runtime-ext-dataproc/src/main/java/io/cdap/cdap/runtime/spi/common/DataprocUtils.java

@@ -397,6 +397,9 @@ public static void emitMetric(ProvisionerContext context, DataprocMetric datapro
     if (!Strings.isNullOrEmpty(dataprocMetric.getImageVersion())) {
       tags.put("imgVer", dataprocMetric.getImageVersion());
     }
+    if (!Strings.isNullOrEmpty(dataprocMetric.getOpType())) {
+      tags.put("method", dataprocMetric.getOpType());
+    }
     ProvisionerMetrics metrics = context.getMetrics(tags.build());
     metrics.count(dataprocMetric.getMetricName(), 1);
   }

cdap-runtime-ext-dataproc/src/main/java/io/cdap/cdap/runtime/spi/provisioner/dataproc/DataprocClient.java

@@ -73,6 +73,7 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -822,6 +823,22 @@ String getClusterFailureMsg(String name) throws RetryableProvisionException {
     return "";
   }
 
+  /**
+   * Get the latest operation by cluster name and operation type.
+   */
+  public Optional<Operation> getLatestOperation(String clusterName, String operationType) throws ApiException {
+    String projectId = conf.getProjectId();
+    String region = conf.getRegion();
+    String resourceName = String.format("projects/%s/regions/%s/operations", projectId, region);
+    String filter = String.format("clusterName=\"%s\" AND operationType=\"%s\"", clusterName, operationType);
+
+    OperationsClient.ListOperationsPagedResponse response =
+      client.getOperationsClient().listOperations(resourceName, filter);
+
+    Iterator<Operation> iterator = response.getPage().getValues().iterator();
+    return iterator.hasNext() ? Optional.of(iterator.next()) : Optional.empty();
+  }
+
   /**
    * Get information about the specified cluster. The cluster will not be present if it could not be
    * found.

cdap-runtime-ext-dataproc/src/main/java/io/cdap/cdap/runtime/spi/provisioner/dataproc/DataprocProvisioner.java

@@ -16,16 +16,22 @@
 
 package io.cdap.cdap.runtime.spi.provisioner.dataproc;
 
+import com.google.api.gax.grpc.GrpcStatusCode;
+import com.google.api.gax.rpc.ApiException;
+import com.google.api.gax.rpc.StatusCode;
 import com.google.cloud.dataproc.v1.ClusterOperationMetadata;
 import com.google.cloud.dataproc.v1.ClusterStatus.State;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Stopwatch;
 import com.google.common.collect.ImmutableSet;
+import com.google.longrunning.Operation;
+import com.google.rpc.Status;
 import io.cdap.cdap.api.exception.ErrorCategory;
 import io.cdap.cdap.api.exception.ErrorType;
 import io.cdap.cdap.runtime.spi.ProgramRunInfo;
 import io.cdap.cdap.runtime.spi.RuntimeMonitorType;
 import io.cdap.cdap.runtime.spi.common.DataprocImageVersion;
+import io.cdap.cdap.runtime.spi.common.DataprocMetric;
 import io.cdap.cdap.runtime.spi.common.DataprocUtils;
 import io.cdap.cdap.runtime.spi.provisioner.Cluster;
 import io.cdap.cdap.runtime.spi.provisioner.ClusterStatus;
@@ -37,18 +43,22 @@
 import io.cdap.cdap.runtime.spi.ssh.SSHContext;
 import io.cdap.cdap.runtime.spi.ssh.SSHKeyPair;
 import io.cdap.cdap.runtime.spi.ssh.SSHPublicKey;
+import io.grpc.Status.Code;
 import java.util.Collections;
+import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.locks.Lock;
 import java.util.regex.Pattern;
 import javax.annotation.Nullable;
+
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -73,6 +83,12 @@ public class DataprocProvisioner extends AbstractDataprocProvisioner {
   //A lock to use for cluster reuse
   private static final String REUSE_LOCK = "reuse";
 
+  private static final Set<ClusterStatus> TRANSITIONING_STATES =
+    EnumSet.of(ClusterStatus.CREATING, ClusterStatus.DELETING);
+
+  private static final Set<ClusterStatus> TERMINAL_STATES =
+    EnumSet.of(ClusterStatus.RUNNING, ClusterStatus.FAILED, ClusterStatus.NOT_EXISTS);
+
   private final DataprocClientFactory clientFactory;
 
   @SuppressWarnings("WeakerAccess")
@@ -471,16 +487,21 @@ public ClusterStatus getClusterStatus(ProvisionerContext context, Cluster cluste
     DataprocConf conf = DataprocConf.create(createContextProperties(context));
     String clusterName = cluster.getName();
 
-    ClusterStatus status = cluster.getStatus();
+    ClusterStatus previousStatus = cluster.getStatus();
     // if we are skipping the delete, need to avoid checking the real cluster status and pretend like it is deleted.
-    if (conf.isSkipDelete() && status == ClusterStatus.DELETING) {
+    if (conf.isSkipDelete() && previousStatus == ClusterStatus.DELETING) {
       return ClusterStatus.NOT_EXISTS;
     }
 
+    ClusterStatus status;
     try (DataprocClient client = clientFactory.create(conf, context.getErrorCategory())) {
       status = client.getClusterStatus(clusterName);
       DataprocUtils.emitMetric(context, conf.getRegion(), getImageVersion(context, conf),
           "provisioner.clusterStatus.response.count");
+
+      if (isEndState(previousStatus, status)) {
+        emitLroMetric(context, client, conf, previousStatus, clusterName);
+      }
       return status;
     } catch (Exception e) {
       DataprocUtils.emitMetric(context, conf.getRegion(), getImageVersion(context, conf),
@@ -489,6 +510,84 @@ public ClusterStatus getClusterStatus(ProvisionerContext context, Cluster cluste
     }
   }
 
+  private boolean isEndState(ClusterStatus previous, ClusterStatus current) {
+    return TRANSITIONING_STATES.contains(previous) && TERMINAL_STATES.contains(current);
+  }
+
+  /**
+   * Emits a metric for the outcome of a Long Running Operation (LRO) when a cluster
+   * transitions to a terminal status.
+   *
+   * @param context the provisioner context
+   * @param client the dataproc client
+   * @param conf the dataproc configuration
+   * @param previousStatus the status of the cluster before the transition
+   * @param clusterName the name of the cluster
+   */
+  private void emitLroMetric(ProvisionerContext context, DataprocClient client, DataprocConf conf,
+                             ClusterStatus previousStatus, String clusterName) {
+    String opType = getOpType(previousStatus);
+    if (opType == null) {
+      return;
+    }
+
+    try {
+      Optional<Operation> opOpt =
+        client.getLatestOperation(clusterName, opType);
+
+      if (!opOpt.isPresent() || !opOpt.get().getDone()) {
+        return;
+      }
+
+      Operation op = opOpt.get();
+      String metricName = "provisioner.operation.response.count";
+      String imageVersion = getImageVersion(context, conf);
+
+      DataprocMetric.Builder builder = DataprocMetric.builder(metricName)
+        .setRegion(conf.getRegion())
+        .setImageVersion(imageVersion)
+        .setOpType(opType);
+
+      if (op.hasError()) {
+        builder.setException(createApiException(op.getError()));
+      }
+
+      DataprocUtils.emitMetric(context, builder.build());
+
+    } catch (Exception ex) {
+      LOG.warn("Failed to retrieve LRO operation metadata for metric emission on cluster {}", clusterName, ex);
+    }
+  }
+
+  @Nullable
+  private String getOpType(ClusterStatus status) {
+    if (status == ClusterStatus.CREATING) {
+      return "CREATE";
+    }
+    if (status == ClusterStatus.DELETING) {
+      return "DELETE";
+    }
+    return null;
+  }
+
+  private Exception createApiException(Status error) {
+    int rpcCode = error.getCode();
+
+    Code grpcCode = Code.INTERNAL;
+
+    if (rpcCode >= 0 && rpcCode < io.grpc.Status.Code.values().length) {
+      grpcCode = Code.values()[rpcCode];
+    }
+
+    StatusCode gaxStatusCode =
+      GrpcStatusCode.of(grpcCode);
+
+    ApiException apiException = new ApiException(
+      new Exception(error.getMessage()), gaxStatusCode, false);
+
+    return new Exception(apiException);
+  }
+
   @Override
   public String getClusterFailureMsg(ProvisionerContext context, Cluster cluster) throws Exception {
     DataprocConf conf = DataprocConf.create(createContextProperties(context));

cdap-runtime-ext-dataproc/src/test/java/io/cdap/cdap/runtime/spi/provisioner/dataproc/DataprocClientTest.java

@@ -59,6 +59,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.ExecutionException;
 import org.hamcrest.core.IsInstanceOf;
 import org.junit.Assert;
@@ -378,4 +379,51 @@ public void testGetClusterStatusCapturesErrorMessage() throws GeneralSecurityExc
     String errorMsgRet = client.getClusterFailureMsg(cdapCluster.getName());
     Assert.assertTrue(errorMsgRet.contains(errorMsg));
   }
+
+  @Test
+  public void testGetLatestOperation() throws Exception {
+    DataprocClient client = sshDataprocClientFactory.create(dataprocConf,
+                                                            new ErrorCategory(ErrorCategoryEnum.OTHERS));
+
+    OperationsClient operationsClient = PowerMockito.mock(OperationsClient.class);
+    PowerMockito.when(clusterControllerClientMock.getOperationsClient()).thenReturn(operationsClient);
+
+    OperationsClient.ListOperationsPagedResponse listOperationsPagedResponse =
+      PowerMockito.mock(OperationsClient.ListOperationsPagedResponse.class);
+    PowerMockito.when(operationsClient.listOperations(
+        Mockito.eq("projects/dummy-project/regions/us-test1/operations"),
+        Mockito.eq("clusterName=\"mycluster\" AND operationType=\"CREATE\"")))
+      .thenReturn(listOperationsPagedResponse);
+
+    OperationsClient.ListOperationsPage page = Mockito.mock(OperationsClient.ListOperationsPage.class);
+    Mockito.when(listOperationsPagedResponse.getPage()).thenReturn(page);
+
+    Operation operation = Operation.newBuilder().setName("op-name").setDone(true).build();
+    List<Operation> operations = Collections.singletonList(operation);
+    Mockito.when(page.getValues()).thenReturn(operations);
+
+    Optional<Operation> result = client.getLatestOperation("mycluster", "CREATE");
+    Assert.assertTrue(result.isPresent());
+    Assert.assertEquals("op-name", result.get().getName());
+  }
+
+  @Test
+  public void testGetLatestOperationThrowsException() throws Exception {
+    DataprocClient client = sshDataprocClientFactory.create(dataprocConf,
+                                                            new ErrorCategory(ErrorCategoryEnum.OTHERS));
+
+    OperationsClient operationsClient = PowerMockito.mock(OperationsClient.class);
+    PowerMockito.when(clusterControllerClientMock.getOperationsClient()).thenReturn(operationsClient);
+
+    ApiException mockException = PowerMockito.mock(ApiException.class);
+    PowerMockito.when(operationsClient.listOperations(Mockito.anyString(), Mockito.anyString()))
+      .thenThrow(mockException);
+
+    try {
+      client.getLatestOperation("mycluster", "CREATE");
+      Assert.fail("Expected ApiException to be thrown");
+    } catch (ApiException e) {
+      Assert.assertEquals(mockException, e);
+    }
+  }
 }