Merge pull request #3768 from atlanhq/client-tag

feat: add client origin tag in prometheus
atlanhq · Nov 19, 2024 · b53fa05 · b53fa05
2 parents b2af14d + b3f24c4
commit b53fa05
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 4 deletions.
diff --git a/common/src/main/java/org/apache/atlas/service/metrics/MetricsRegistry.java b/common/src/main/java/org/apache/atlas/service/metrics/MetricsRegistry.java
@@ -8,7 +8,7 @@
 
 public interface MetricsRegistry {
 
-    void collect(String requestId, String requestUri, AtlasPerfMetrics metrics);
+    void collect(String requestId, String requestUri, AtlasPerfMetrics metrics, String clientOrigin);
 
     void collectApplicationMetrics(String requestId, String requestUri, List<AtlasPerfMetrics.Metric> applicationMetrics);
 

diff --git a/common/src/main/java/org/apache/atlas/service/metrics/MetricsRegistryServiceImpl.java b/common/src/main/java/org/apache/atlas/service/metrics/MetricsRegistryServiceImpl.java
@@ -30,6 +30,7 @@ public class MetricsRegistryServiceImpl implements MetricsRegistry {
 
     private static final String NAME = "name";
     private static final String URI = "uri";
+    private static final String ORIGIN = "origin";
     private static final String METHOD_DIST_SUMMARY = "method_dist_summary";
     private static final String APPLICATION_LEVEL_METRICS_SUMMARY = "application_level_metrics_summary";
     private static final double[] PERCENTILES = {0.99};
@@ -43,7 +44,7 @@ public MetricsRegistryServiceImpl() throws AtlasException {
     }
 
     @Override
-    public void collect(String requestId, String requestUri, AtlasPerfMetrics metrics) {
+    public void collect(String requestId, String requestUri, AtlasPerfMetrics metrics, String clientOrigin) {
         try {
             if (!ApplicationProperties.get().getBoolean(METHOD_LEVEL_METRICS_ENABLE, false)) {
                 return;
@@ -52,7 +53,7 @@ public void collect(String requestId, String requestUri, AtlasPerfMetrics metric
             for (String name : this.filteredMethods) {
                 if(metrics.hasMetric(name)) {
                     AtlasPerfMetrics.Metric metric = metrics.getMetric(name);
-                    Timer.builder(METHOD_DIST_SUMMARY).tags(Tags.of(NAME, metric.getName(), URI, requestUri)).publishPercentiles(PERCENTILES)
+                    Timer.builder(METHOD_DIST_SUMMARY).tags(Tags.of(NAME, metric.getName(), URI, requestUri, ORIGIN, clientOrigin)).publishPercentiles(PERCENTILES)
                             .register(getMeterRegistry()).record(metric.getTotalTimeMSecs(), TimeUnit.MILLISECONDS);
                 }
             }

diff --git a/server-api/src/main/java/org/apache/atlas/RequestContext.java b/server-api/src/main/java/org/apache/atlas/RequestContext.java
@@ -80,6 +80,7 @@ public class RequestContext {
     private Set<String> userGroups;
     private String clientIPAddress;
     private List<String> forwardedAddresses;
+    private String clientOrigin;
     private DeleteType deleteType = DeleteType.DEFAULT;
     private boolean isPurgeRequested = false;
     private int maxAttempts = 1;
@@ -178,7 +179,7 @@ public void clearCache() {
         if (metrics != null && !metrics.isEmpty()) {
             METRICS.debug(metrics.toString());
             if (Objects.nonNull(this.metricsRegistry)){
-                this.metricsRegistry.collect(traceId, this.requestUri, metrics);
+                this.metricsRegistry.collect(traceId, this.requestUri, metrics, this.getClientOrigin());
             }
             metrics.clear();
         }
@@ -744,6 +745,14 @@ public void setIncludeClassificationNames(boolean includeClassificationNames) {
         this.includeClassificationNames = includeClassificationNames;
     }
 
+    public String getClientOrigin() {
+        return clientOrigin;
+    }
+
+    public void setClientOrigin(String clientOrigin) {
+        this.clientOrigin = StringUtils.isEmpty(this.clientOrigin) ? "other" :clientOrigin;
+    }
+
     public class EntityGuidPair {
         private final Object entity;
         private final String guid;

diff --git a/webapp/src/main/java/org/apache/atlas/web/filters/AuditFilter.java b/webapp/src/main/java/org/apache/atlas/web/filters/AuditFilter.java
@@ -60,6 +60,7 @@ public class AuditFilter implements Filter {
     private static final Logger AUDIT_LOG = LoggerFactory.getLogger("AUDIT");
     public static final String TRACE_ID   = "trace_id";
     public static final String X_ATLAN_REQUEST_ID = "X-Atlan-Request-Id";
+    public static final String X_ATLAN_CLIENT_ORIGIN = "X-Atlan-Client-Origin";
     private boolean deleteTypeOverrideEnabled                = false;
     private boolean createShellEntityForNonExistingReference = false;
 
@@ -103,6 +104,7 @@ public void doFilter(ServletRequest request, ServletResponse response, FilterCha
             requestContext.setCreateShellEntityForNonExistingReference(createShellEntityForNonExistingReference);
             requestContext.setForwardedAddresses(AtlasAuthorizationUtils.getForwardedAddressesFromRequest(httpRequest));
             requestContext.setSkipFailedEntities(skipFailedEntities);
+            requestContext.setClientOrigin(httpRequest.getHeader(X_ATLAN_CLIENT_ORIGIN));
             requestContext.setMetricRegistry(metricsRegistry);
             MDC.put(TRACE_ID, internalRequestId);
             MDC.put(X_ATLAN_REQUEST_ID, ofNullable(httpRequest.getHeader(X_ATLAN_REQUEST_ID)).orElse(EMPTY));