Add image version to dataproc provisioner metric

cdapio · Jan 12, 2024 · 2bdbc8a · 2bdbc8a
1 parent 489ba81
commit 2bdbc8a
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 15 deletions.
diff --git a/cdap-runtime-ext-dataproc/src/main/java/io/cdap/cdap/runtime/spi/common/DataprocMetric.java b/cdap-runtime-ext-dataproc/src/main/java/io/cdap/cdap/runtime/spi/common/DataprocMetric.java
@@ -29,19 +29,27 @@ public class DataprocMetric {
   private final Exception exception;
   @Nullable
   private final LaunchMode launchMode;
+  @Nullable
+  private final String imageVersion;
 
   private DataprocMetric(String metricName, String region, @Nullable Exception exception,
-      @Nullable LaunchMode launchMode) {
+      @Nullable LaunchMode launchMode, @Nullable String imageVersion) {
     this.metricName = metricName;
     this.region = region;
     this.exception = exception;
     this.launchMode = launchMode;
+    this.imageVersion = imageVersion;
   }
 
   public String getMetricName() {
     return metricName;
   }
 
+  @Nullable
+  public String getImageVersion() {
+    return imageVersion;
+  }
+
   public String getRegion() {
     return region;
   }
@@ -72,6 +80,7 @@ public static Builder builder(String metricName) {
   public static class Builder {
     private final String metricName;
     private String region;
+    private String imageVersion;
     @Nullable
     private Exception exception;
     @Nullable
@@ -96,6 +105,11 @@ public Builder setLaunchMode(@Nullable LaunchMode launchMode) {
       return this;
     }
 
+    public Builder setImageVersion(String imageVersion) {
+      this.imageVersion = imageVersion;
+      return this;
+    }
+
     /**
      * Returns a DataprocMetric.
      *
@@ -106,7 +120,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);
+      return new DataprocMetric(metricName, region, exception, launchMode, imageVersion);
     }
   }
 }
diff --git a/cdap-runtime-ext-dataproc/src/main/java/io/cdap/cdap/runtime/spi/common/DataprocUtils.java b/cdap-runtime-ext-dataproc/src/main/java/io/cdap/cdap/runtime/spi/common/DataprocUtils.java
@@ -303,13 +303,14 @@ public static String getSystemProjectId() {
    * Emit a dataproc metric.
    **/
   public static void emitMetric(ProvisionerContext context, String region,
-      String metricName, @Nullable Exception e) {
+      String metricName, @Nullable String imageVersion, @Nullable Exception e) {
     emitMetric(context,
-        DataprocMetric.builder(metricName).setRegion(region).setException(e).build());
+        DataprocMetric.builder(metricName).setRegion(region).setException(e).setImageVersion(imageVersion).build());
   }
 
-  public static void emitMetric(ProvisionerContext context, String region, String metricName) {
-    emitMetric(context, region, metricName, null);
+  public static void emitMetric(ProvisionerContext context, String region,
+      @Nullable String imageVersion, String metricName) {
+    emitMetric(context, region, metricName, imageVersion, null);
   }
 
   /**
@@ -335,6 +336,9 @@ public static void emitMetric(ProvisionerContext context, DataprocMetric datapro
     if (dataprocMetric.getLaunchMode() != null) {
       tags.put("lchmode", dataprocMetric.getLaunchMode().name());
     }
+    if (!Strings.isNullOrEmpty(dataprocMetric.getImageVersion())) {
+      tags.put("imgVer", dataprocMetric.getImageVersion());
+    }
     ProvisionerMetrics metrics = context.getMetrics(tags.build());
     metrics.count(dataprocMetric.getMetricName(), 1);
   }

diff --git a/...proc/src/main/java/io/cdap/cdap/runtime/spi/provisioner/dataproc/DataprocProvisioner.java b/...proc/src/main/java/io/cdap/cdap/runtime/spi/provisioner/dataproc/DataprocProvisioner.java
@@ -172,7 +172,7 @@ public Cluster createCluster(ProvisionerContext context) throws Exception {
     try (DataprocClient client = clientFactory.create(conf, sshPublicKey != null)) {
       Cluster reused = tryReuseCluster(client, context, conf);
       if (reused != null) {
-        DataprocUtils.emitMetric(context, conf.getRegion(),
+        DataprocUtils.emitMetric(context, conf.getRegion(), getImageVersion(context, conf),
             "provisioner.createCluster.reuse.count");
         return reused;
       }
@@ -224,12 +224,12 @@ public Cluster createCluster(ProvisionerContext context) throws Exception {
             numWarnings, numWarnings > 1 ? "s" : "",
             String.join("\n", createOperationMeta.getWarningsList()));
       }
-      DataprocUtils.emitMetric(context, conf.getRegion(),
+      DataprocUtils.emitMetric(context, conf.getRegion(), getImageVersion(context, conf),
           "provisioner.createCluster.response.count");
       return new Cluster(clusterName, ClusterStatus.CREATING, Collections.emptyList(),
           Collections.emptyMap());
     } catch (Exception e) {
-      DataprocUtils.emitMetric(context, conf.getRegion(),
+      DataprocUtils.emitMetric(context, conf.getRegion(), getImageVersion(context, conf),
           "provisioner.createCluster.response.count", e);
       throw e;
     }
@@ -448,11 +448,11 @@ public ClusterStatus getClusterStatus(ProvisionerContext context, Cluster cluste
 
     try (DataprocClient client = clientFactory.create(conf)) {
       status = client.getClusterStatus(clusterName);
-      DataprocUtils.emitMetric(context, conf.getRegion(),
+      DataprocUtils.emitMetric(context, conf.getRegion(), getImageVersion(context, conf),
           "provisioner.clusterStatus.response.count");
       return status;
     } catch (Exception e) {
-      DataprocUtils.emitMetric(context, conf.getRegion(),
+      DataprocUtils.emitMetric(context, conf.getRegion(), getImageVersion(context, conf),
           "provisioner.clusterStatus.response.count", e);
       throw e;
     }
@@ -474,11 +474,11 @@ public Cluster getClusterDetail(ProvisionerContext context, Cluster cluster) thr
     String clusterName = cluster.getName();
     try (DataprocClient client = clientFactory.create(conf, shouldUseSsh(context, conf))) {
       Optional<Cluster> existing = client.getCluster(clusterName);
-      DataprocUtils.emitMetric(context, conf.getRegion(),
+      DataprocUtils.emitMetric(context, conf.getRegion(), getImageVersion(context, conf),
           "provisioner.clusterDetail.response.count");
       return existing.orElseGet(() -> new Cluster(cluster, ClusterStatus.NOT_EXISTS));
     } catch (Exception e) {
-      DataprocUtils.emitMetric(context, conf.getRegion(),
+      DataprocUtils.emitMetric(context, conf.getRegion(), getImageVersion(context, conf),
           "provisioner.clusterDetail.response.count", e);
       throw e;
     }
@@ -507,10 +507,10 @@ protected void doDeleteCluster(ProvisionerContext context, Cluster cluster, Data
       } else {
         client.deleteCluster(clusterName);
       }
-      DataprocUtils.emitMetric(context, conf.getRegion(),
+      DataprocUtils.emitMetric(context, conf.getRegion(), getImageVersion(context, conf),
           "provisioner.deleteCluster.response.count");
     } catch (Exception e) {
-      DataprocUtils.emitMetric(context, conf.getRegion(),
+      DataprocUtils.emitMetric(context, conf.getRegion(), getImageVersion(context, conf),
           "provisioner.deleteCluster.response.count", e);
       throw e;
     }