From de1f23d9d9fa391f24b3198d1ed6685e5d3700af Mon Sep 17 00:00:00 2001
From: david-leifker <114954101+david-leifker@users.noreply.github.com>
Date: Mon, 31 Jul 2023 09:10:40 -0500
Subject: [PATCH] feat(metrics): add metrics for aspect write and bytes (#8526)

---
 .../com/linkedin/metadata/entity/AspectDao.java | 10 ++++++++++
 .../metadata/entity/EntityServiceImpl.java      | 17 +++++++++++++++--
 .../entity/cassandra/CassandraAspectDao.java    |  4 ++++
 3 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/AspectDao.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/AspectDao.java
index bf74b1025267f..9967df9207ec7 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/entity/AspectDao.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/AspectDao.java
@@ -3,6 +3,7 @@
 import com.linkedin.common.urn.Urn;
 import com.linkedin.metadata.entity.ebean.EbeanAspectV2;
 import com.linkedin.metadata.entity.restoreindices.RestoreIndicesArgs;
+import com.linkedin.metadata.utils.metrics.MetricUtils;
 import io.ebean.PagedList;
 
 import javax.annotation.Nonnull;
@@ -28,6 +29,8 @@
  * worth looking into ways to move this responsibility inside {@link AspectDao} implementations.
  */
 public interface AspectDao {
+    String ASPECT_WRITE_COUNT_METRIC_NAME = "aspectWriteCount";
+    String ASPECT_WRITE_BYTES_METRIC_NAME = "aspectWriteBytes";
 
     @Nullable
     EntityAspect getAspect(@Nonnull final String urn, @Nonnull final String aspectName, final long version);
@@ -116,4 +119,11 @@ ListResult<String> listAspectMetadata(
 
     @Nonnull
     <T> T runInTransactionWithRetry(@Nonnull final Supplier<T> block, final int maxTransactionRetry);
+
+    default void incrementWriteMetrics(String aspectName, long count, long bytes) {
+        MetricUtils.counter(this.getClass(),
+                String.join(MetricUtils.DELIMITER, List.of(ASPECT_WRITE_COUNT_METRIC_NAME, aspectName))).inc(count);
+        MetricUtils.counter(this.getClass(),
+                String.join(MetricUtils.DELIMITER, List.of(ASPECT_WRITE_BYTES_METRIC_NAME, aspectName))).inc(bytes);
+    }
 }
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityServiceImpl.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityServiceImpl.java
index c7a9895992d90..e070944b49a05 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityServiceImpl.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityServiceImpl.java
@@ -1811,6 +1811,8 @@ public RollbackResult deleteAspect(String urn, String aspectName, @Nonnull Map<S
         latest.setCreatedBy(survivingAspect.getCreatedBy());
         latest.setCreatedFor(survivingAspect.getCreatedFor());
         _aspectDao.saveAspect(latest, false);
+        // metrics
+        _aspectDao.incrementWriteMetrics(aspectName, 1, latest.getAspect().getBytes(StandardCharsets.UTF_8).length);
         _aspectDao.deleteAspect(survivingAspect);
       } else {
         if (isKeyAspect) {
@@ -2014,6 +2016,9 @@ private UpdateAspectResult ingestAspectToLocalDBNoTransaction(
 
       _aspectDao.saveAspect(latest, false);
 
+      // metrics
+      _aspectDao.incrementWriteMetrics(aspectName, 1, latest.getAspect().getBytes(StandardCharsets.UTF_8).length);
+
       return new UpdateAspectResult(urn, oldValue, oldValue,
           EntityUtils.parseSystemMetadata(latest.getSystemMetadata()), latestSystemMetadata,
           MetadataAuditOperation.UPDATE, auditStamp, 0);
@@ -2021,13 +2026,17 @@ private UpdateAspectResult ingestAspectToLocalDBNoTransaction(
 
     // 4. Save the newValue as the latest version
     log.debug("Ingesting aspect with name {}, urn {}", aspectName, urn);
+    String newValueStr = EntityUtils.toJsonAspect(newValue);
     long versionOfOld = _aspectDao.saveLatestAspect(urn.toString(), aspectName, latest == null ? null : EntityUtils.toJsonAspect(oldValue),
         latest == null ? null : latest.getCreatedBy(), latest == null ? null : latest.getCreatedFor(),
         latest == null ? null : latest.getCreatedOn(), latest == null ? null : latest.getSystemMetadata(),
-        EntityUtils.toJsonAspect(newValue), auditStamp.getActor().toString(),
+            newValueStr, auditStamp.getActor().toString(),
         auditStamp.hasImpersonator() ? auditStamp.getImpersonator().toString() : null,
         new Timestamp(auditStamp.getTime()), EntityUtils.toJsonAspect(providedSystemMetadata), nextVersion);
 
+    // metrics
+    _aspectDao.incrementWriteMetrics(aspectName, 1, newValueStr.getBytes(StandardCharsets.UTF_8).length);
+
     return new UpdateAspectResult(urn, oldValue, newValue,
         latest == null ? null : EntityUtils.parseSystemMetadata(latest.getSystemMetadata()), providedSystemMetadata,
         MetadataAuditOperation.UPDATE, auditStamp, versionOfOld);
@@ -2072,10 +2081,14 @@ private RecordTemplate updateAspect(
       newSystemMetadata.setLastObserved(System.currentTimeMillis());
 
       log.debug("Updating aspect with name {}, urn {}", aspectName, urn);
-      _aspectDao.saveAspect(urn.toString(), aspectName, EntityUtils.toJsonAspect(value), auditStamp.getActor().toString(),
+      String aspectStr = EntityUtils.toJsonAspect(value);
+      _aspectDao.saveAspect(urn.toString(), aspectName, aspectStr, auditStamp.getActor().toString(),
           auditStamp.hasImpersonator() ? auditStamp.getImpersonator().toString() : null,
           new Timestamp(auditStamp.getTime()), EntityUtils.toJsonAspect(newSystemMetadata), version, oldAspect == null);
 
+      // metrics
+      _aspectDao.incrementWriteMetrics(aspectName, 1, aspectStr.getBytes(StandardCharsets.UTF_8).length);
+
       return new UpdateAspectResult(urn, oldValue, value, oldSystemMetadata, newSystemMetadata,
           MetadataAuditOperation.UPDATE, auditStamp, version);
     }, maxTransactionRetry);
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/cassandra/CassandraAspectDao.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/cassandra/CassandraAspectDao.java
index d75a61d6246ea..dcb58ee324a23 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/entity/cassandra/CassandraAspectDao.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/cassandra/CassandraAspectDao.java
@@ -32,6 +32,7 @@
 import com.linkedin.metadata.query.ExtraInfoArray;
 import com.linkedin.metadata.query.ListResultMetadata;
 import java.net.URISyntaxException;
+import java.nio.charset.StandardCharsets;
 import java.sql.Timestamp;
 import java.util.HashMap;
 import java.util.List;
@@ -596,6 +597,9 @@ public void saveAspect(
     );
 
     saveAspect(aspect, insert);
+
+    // metrics
+    incrementWriteMetrics(aspectName, 1, aspectMetadata.getBytes(StandardCharsets.UTF_8).length);
   }
 
   @Override