Add opt-in per-submission-mode breakdown for user metrics

Author: Binbing Hou

genie-web/src/integTest/java/com/netflix/genie/web/data/services/impl/jpa/JpaPersistenceServiceImplJobsIntegrationTest.java

@@ -1067,15 +1067,27 @@ void canGetJobMetadata() throws GenieException {
     @Test
     @DatabaseSetup("persistence/jobs/search.xml")
     void canGetUserResourceSummaries() {
-        final Map<String, UserResourcesSummary> summaries = this.service.getUserResourcesSummaries(
+        // API-submitted (api=true) jobs.
+        final Map<String, UserResourcesSummary> apiSummaries = this.service.getUserResourcesSummaries(
             JobStatus.getActiveStatuses(),
             true
         );
-        Assertions.assertThat(summaries.keySet()).contains("tgianos");
-        final UserResourcesSummary userResourcesSummary = summaries.get("tgianos");
-        Assertions.assertThat(userResourcesSummary.getUser()).isEqualTo("tgianos");
-        Assertions.assertThat(userResourcesSummary.getRunningJobsCount()).isEqualTo(2L);
-        Assertions.assertThat(userResourcesSummary.getUsedMemory()).isEqualTo(4096L);
+        Assertions.assertThat(apiSummaries.keySet()).contains("tgianos");
+        final UserResourcesSummary apiSummary = apiSummaries.get("tgianos");
+        Assertions.assertThat(apiSummary.getUser()).isEqualTo("tgianos");
+        Assertions.assertThat(apiSummary.getRunningJobsCount()).isEqualTo(2L);
+        Assertions.assertThat(apiSummary.getUsedMemory()).isEqualTo(4096L);
+
+        // Agent/CLI-submitted (api=false) jobs are reported separately, with memory populated.
+        final Map<String, UserResourcesSummary> agentSummaries = this.service.getUserResourcesSummaries(
+            JobStatus.getActiveStatuses(),
+            false
+        );
+        Assertions.assertThat(agentSummaries.keySet()).contains("tgianos");
+        final UserResourcesSummary agentSummary = agentSummaries.get("tgianos");
+        Assertions.assertThat(agentSummary.getUser()).isEqualTo("tgianos");
+        Assertions.assertThat(agentSummary.getRunningJobsCount()).isEqualTo(2L);
+        Assertions.assertThat(agentSummary.getUsedMemory()).isEqualTo(4096L);
     }
 
     @Test

genie-web/src/main/java/com/netflix/genie/web/properties/UserMetricsProperties.java

@@ -47,4 +47,13 @@ public class UserMetricsProperties {
     private boolean enabled = true;
 
     private long refreshInterval = 30_000;
+
+    /**
+     * Whether per-user metrics are additionally published broken down by job submission mode (tagged
+     * {@code submissionMode=api} for REST vs {@code submissionMode=agent} for agent/CLI), under the
+     * {@code *.by-submission-mode.gauge} metric names. The existing
+     * {@code genie.user.active-jobs.gauge}/{@code genie.user.active-memory.gauge} gauges are unaffected
+     * regardless of this setting.
+     */
+    private boolean splitBySubmissionMode;
 }

genie-web/src/main/java/com/netflix/genie/web/tasks/leader/UserMetricsTask.java

@@ -31,6 +31,7 @@
 import io.micrometer.core.instrument.MeterRegistry;
 import lombok.extern.slf4j.Slf4j;
 
+import java.util.Collections;
 import java.util.Map;
 import java.util.Set;
 
@@ -46,12 +47,19 @@ public class UserMetricsTask extends LeaderTask {
     private static final String USER_ACTIVE_JOBS_METRIC_NAME = "genie.user.active-jobs.gauge";
     private static final String USER_ACTIVE_MEMORY_METRIC_NAME = "genie.user.active-memory.gauge";
     private static final String USER_ACTIVE_USERS_METRIC_NAME = "genie.user.active-users.gauge";
+    private static final String USER_ACTIVE_JOBS_BY_SUBMISSION_MODE_METRIC_NAME =
+        "genie.user.active-jobs.by-submission-mode.gauge";
+    private static final String USER_ACTIVE_MEMORY_BY_SUBMISSION_MODE_METRIC_NAME =
+        "genie.user.active-memory.by-submission-mode.gauge";
     private static final UserResourcesRecord USER_RECORD_PLACEHOLDER = new UserResourcesRecord("nobody");
     private final MeterRegistry registry;
     private final PersistenceService persistenceService;
     private final UserMetricsProperties userMetricsProperties;
 
     private final Map<String, UserResourcesRecord> userResourcesRecordMap = Maps.newHashMap();
+    // Records for the optional per-submission-mode resource breakdown: api (true/false) -> user -> record.
+    private final Map<Boolean, Map<String, UserResourcesRecord>> userResourcesBySubmissionModeRecordMap =
+        Maps.newHashMap();
     private final AtomicDouble activeUsersCount;
 
     /**
@@ -153,6 +161,13 @@ public void run() {
             ).update(jobs, memory);
         }
 
+        // Optionally publish a finer-grained breakdown tagged by submission mode (submissionMode=api REST vs
+        // submissionMode=agent agent/CLI). Emitted under separate metric names so the gauges above are unchanged.
+        // Disabled by default.
+        if (this.userMetricsProperties.isSplitBySubmissionMode()) {
+            this.publishResourcesBySubmissionMode(summaries);
+        }
+
         log.debug("Done publishing user metrics");
     }
 
@@ -166,6 +181,7 @@ public void cleanup() {
 
         // Reset all users
         this.userResourcesRecordMap.clear();
+        this.userResourcesBySubmissionModeRecordMap.clear();
 
         // Reset active users count
         this.activeUsersCount.set(Double.NaN);
@@ -189,6 +205,70 @@ private Number getUsersCount() {
         return activeUsersCount.get();
     }
 
+    private void publishResourcesBySubmissionMode(final Map<String, UserResourcesSummary> apiSummaries) {
+        // Reuse the api=true summaries already fetched above; only the agent (api=false) slice is
+        // an additional query.
+        this.updateSubmissionModeRecords(true, apiSummaries);
+        this.updateSubmissionModeRecords(
+            false,
+            this.persistenceService.getUserResourcesSummaries(JobStatus.getActiveStatuses(), false)
+        );
+    }
+
+    private void updateSubmissionModeRecords(final boolean api, final Map<String, UserResourcesSummary> summaries) {
+        final Map<String, UserResourcesRecord> recordsForMode =
+            this.userResourcesBySubmissionModeRecordMap.computeIfAbsent(api, mode -> Maps.newHashMap());
+
+        final String submissionModeTagValue =
+            api ? MetricsConstants.TagValues.SUBMISSION_MODE_API : MetricsConstants.TagValues.SUBMISSION_MODE_AGENT;
+
+        // Drop users no longer active in this mode so their gauges fall back to NaN.
+        recordsForMode.keySet().retainAll(summaries.keySet());
+
+        for (final UserResourcesSummary summary : summaries.values()) {
+            final String user = summary.getUser();
+            recordsForMode.computeIfAbsent(
+                user,
+                userName -> {
+                    // Gauge creation is idempotent so re-registration on reappearance is harmless.
+                    Gauge.builder(
+                            USER_ACTIVE_JOBS_BY_SUBMISSION_MODE_METRIC_NAME,
+                            () -> this.getSubmissionModeJobCount(api, userName)
+                        )
+                        .tags(
+                            MetricsConstants.TagKeys.USER, userName,
+                            MetricsConstants.TagKeys.SUBMISSION_MODE, submissionModeTagValue
+                        )
+                        .register(this.registry);
+                    Gauge.builder(
+                            USER_ACTIVE_MEMORY_BY_SUBMISSION_MODE_METRIC_NAME,
+                            () -> this.getSubmissionModeMemoryAmount(api, userName)
+                        )
+                        .tags(
+                            MetricsConstants.TagKeys.USER, userName,
+                            MetricsConstants.TagKeys.SUBMISSION_MODE, submissionModeTagValue
+                        )
+                        .register(this.registry);
+                    return new UserResourcesRecord(userName);
+                }
+            ).update(summary.getRunningJobsCount(), summary.getUsedMemory());
+        }
+    }
+
+    private Number getSubmissionModeJobCount(final boolean api, final String user) {
+        return this.userResourcesBySubmissionModeRecordMap
+            .getOrDefault(api, Collections.emptyMap())
+            .getOrDefault(user, USER_RECORD_PLACEHOLDER)
+            .jobCount.get();
+    }
+
+    private Number getSubmissionModeMemoryAmount(final boolean api, final String user) {
+        return this.userResourcesBySubmissionModeRecordMap
+            .getOrDefault(api, Collections.emptyMap())
+            .getOrDefault(user, USER_RECORD_PLACEHOLDER)
+            .memoryAmount.get();
+    }
+
     private static class UserResourcesRecord {
         private final String userName;
         private final AtomicDouble jobCount = new AtomicDouble(Double.NaN);

genie-web/src/main/java/com/netflix/genie/web/util/MetricsConstants.java

@@ -81,6 +81,11 @@ public static final class TagKeys {
          */
         public static final String USER = "user";
 
+        /**
+         * Key to tag the job submission mode ({@code api} for REST API, {@code agent} for agent/CLI).
+         */
+        public static final String SUBMISSION_MODE = "submissionMode";
+
         /**
          * Key to tag the user concurrent job limit.
          */
@@ -122,6 +127,16 @@ public static final class TagValues {
          */
         public static final String FAILURE = "failure";
 
+        /**
+         * Tag value for jobs submitted through the REST API (used with TagKeys.SUBMISSION_MODE).
+         */
+        public static final String SUBMISSION_MODE_API = "api";
+
+        /**
+         * Tag value for jobs submitted through the agent/CLI (used with TagKeys.SUBMISSION_MODE).
+         */
+        public static final String SUBMISSION_MODE_AGENT = "agent";
+
         /**
          * Utility class private constructor.
          */

genie-web/src/test/groovy/com/netflix/genie/web/tasks/leader/UserMetricsTaskSpec.groovy

@@ -172,9 +172,77 @@ class UserMetricsTaskSpec extends Specification {
         measureMemory("boo") == Double.NaN
     }
 
+    def "Run publishes a per-submission-mode breakdown when split-by-submission-mode is enabled"() {
+        setup:
+        Map<String, UserResourcesSummary> apiSummaries = [
+            "foo": new UserResourcesSummary("foo", 10, 1024),
+            "bar": new UserResourcesSummary("bar", 20, 2048)
+        ]
+        Map<String, UserResourcesSummary> agentSummaries = [
+            "foo": new UserResourcesSummary("foo", 3, 256),
+            "baz": new UserResourcesSummary("baz", 5, 512)
+        ]
+
+        when:
+        this.task = new UserMetricsTask(this.registry, this.dataServices, this.userMetricProperties)
+
+        then:
+        1 * registry.gauge(_ as Meter.Id, _, _ as ToDoubleFunction) >> {
+            args -> return captureGauge(args[0] as Meter.Id, args[1] as Object, args[2] as ToDoubleFunction<Object>)
+        }
+
+        when:
+        this.task.run()
+
+        then:
+        _ * userMetricProperties.isSplitBySubmissionMode() >> true
+        1 * persistenceService.getUserResourcesSummaries(JobStatus.activeStatuses, true) >> apiSummaries
+        1 * persistenceService.getUserResourcesSummaries(JobStatus.activeStatuses, false) >> agentSummaries
+        _ * registry.gauge(_ as Meter.Id, _, _ as ToDoubleFunction) >> {
+            args -> return captureGauge(args[0] as Meter.Id, args[1] as Object, args[2] as ToDoubleFunction<Object>)
+        }
+
+        // The existing api=true-only gauges are still published, unchanged.
+        measureJobs("foo") == 10
+        measureMemory("bar") == 2048
+
+        // New per-submission-mode gauges: api (REST) slice.
+        measureJobsBySubmissionMode("foo", "api") == 10
+        measureJobsBySubmissionMode("bar", "api") == 20
+        measureMemoryBySubmissionMode("foo", "api") == 1024
+        measureMemoryBySubmissionMode("bar", "api") == 2048
+
+        // New per-submission-mode gauges: agent slice.
+        measureJobsBySubmissionMode("foo", "agent") == 3
+        measureJobsBySubmissionMode("baz", "agent") == 5
+        measureMemoryBySubmissionMode("foo", "agent") == 256
+        measureMemoryBySubmissionMode("baz", "agent") == 512
+
+        when: "a user drops out of a mode on the next cycle"
+        this.task.run()
+
+        then:
+        _ * userMetricProperties.isSplitBySubmissionMode() >> true
+        1 * persistenceService.getUserResourcesSummaries(JobStatus.activeStatuses, true) >> apiSummaries
+        1 * persistenceService.getUserResourcesSummaries(JobStatus.activeStatuses, false) >>
+            (["foo": new UserResourcesSummary("foo", 3, 256)] as Map<String, UserResourcesSummary>)
+        _ * registry.gauge(_ as Meter.Id, _, _ as ToDoubleFunction) >> {
+            args -> return captureGauge(args[0] as Meter.Id, args[1] as Object, args[2] as ToDoubleFunction<Object>)
+        }
+
+        // baz no longer has agent jobs, so its by-submission-mode gauges report NaN.
+        measureJobsBySubmissionMode("baz", "agent") == Double.NaN
+        measureMemoryBySubmissionMode("baz", "agent") == Double.NaN
+        measureJobsBySubmissionMode("foo", "agent") == 3
+        measureMemoryBySubmissionMode("foo", "agent") == 256
+    }
+
     Gauge captureGauge(final Meter.Id id, final Object obj, final ToDoubleFunction<Object> f) {
         String userTagValue = id.getTag(MetricsConstants.TagKeys.USER)
-        String gaugeKey = id.getName() + (userTagValue == null ? "" : ("-" + userTagValue))
+        String submissionModeTagValue = id.getTag(MetricsConstants.TagKeys.SUBMISSION_MODE)
+        String gaugeKey = id.getName() +
+            (userTagValue == null ? "" : ("-" + userTagValue)) +
+            (submissionModeTagValue == null ? "" : ("-" + submissionModeTagValue))
         this.gaugesFunctions.put(gaugeKey, { -> f.applyAsDouble(obj) })
         return Mock(Gauge)
     }
@@ -192,4 +260,16 @@ class UserMetricsTaskSpec extends Specification {
         String gaugeKey = UserMetricsTask.USER_ACTIVE_MEMORY_METRIC_NAME + "-" + user
         return gaugesFunctions.get(gaugeKey).call()
     }
+
+    double measureJobsBySubmissionMode(String user, String submissionMode) {
+        String gaugeKey =
+            UserMetricsTask.USER_ACTIVE_JOBS_BY_SUBMISSION_MODE_METRIC_NAME + "-" + user + "-" + submissionMode
+        return gaugesFunctions.get(gaugeKey).call()
+    }
+
+    double measureMemoryBySubmissionMode(String user, String submissionMode) {
+        String gaugeKey =
+            UserMetricsTask.USER_ACTIVE_MEMORY_BY_SUBMISSION_MODE_METRIC_NAME + "-" + user + "-" + submissionMode
+        return gaugesFunctions.get(gaugeKey).call()
+    }
 }