Skip to content

Commit 4bb4f24

Browse files
david-leifkerdavid-leifker
committed
feat(system-metrics): track api usage by user, client, api
1 parent f507e2c commit 4bb4f24

File tree

6 files changed

+456
-15
lines changed

6 files changed

+456
-15
lines changed

‎metadata-io/build.gradle

-2
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,6 @@ dependencies {
3131

3232
implementation externalDependency.guava
3333
implementation externalDependency.reflections
34-
// https://mvnrepository.com/artifact/nl.basjes.parse.useragent/yauaa
35-
implementation 'nl.basjes.parse.useragent:yauaa:7.27.0'
3634

3735
api(externalDependency.dgraph4j) {
3836
exclude group: 'com.google.guava', module: 'guava'

‎metadata-io/src/main/java/com/linkedin/metadata/dao/throttle/APIThrottle.java

+2-12
Original file line numberDiff line numberDiff line change
@@ -12,17 +12,9 @@
1212
import java.util.stream.Collectors;
1313
import javax.annotation.Nonnull;
1414
import javax.annotation.Nullable;
15-
import nl.basjes.parse.useragent.UserAgent;
16-
import nl.basjes.parse.useragent.UserAgentAnalyzer;
1715

1816
public class APIThrottle {
1917
private static final Set<String> AGENT_EXEMPTIONS = Set.of("Browser");
20-
private static final UserAgentAnalyzer UAA =
21-
UserAgentAnalyzer.newBuilder()
22-
.hideMatcherLoadStats()
23-
.withField(UserAgent.AGENT_CLASS)
24-
.withCache(1000)
25-
.build();
2618

2719
private APIThrottle() {}
2820

@@ -56,13 +48,11 @@ public static void evaluate(
5648
private static boolean isExempt(@Nullable RequestContext requestContext) {
5749
// Exclude internal calls
5850
if (requestContext == null
59-
|| requestContext.getUserAgent() == null
51+
|| requestContext.getAgentClass() == null
6052
|| requestContext.getUserAgent().isEmpty()) {
6153
return true;
6254
}
63-
64-
UserAgent ua = UAA.parse(requestContext.getUserAgent());
65-
return AGENT_EXEMPTIONS.contains(ua.get(UserAgent.AGENT_CLASS).getValue());
55+
return AGENT_EXEMPTIONS.contains(requestContext.getAgentClass());
6656
}
6757

6858
private static Set<Long> eventMatchMaxWaitMs(

‎metadata-io/src/test/java/com/linkedin/metadata/dao/throttle/APIThrottleTest.java

+18
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
import java.util.List;
1313
import java.util.Map;
1414
import java.util.Set;
15+
import nl.basjes.parse.useragent.UserAgent;
1516
import org.testng.Assert;
1617
import org.testng.annotations.BeforeMethod;
1718
import org.testng.annotations.Test;
@@ -61,6 +62,7 @@ public void testExemptions() {
6162

6263
for (ThrottleEvent event : ALL_EVENTS) {
6364
when(mockRequestContext.getUserAgent()).thenReturn(null);
65+
when(mockRequestContext.getAgentClass()).thenReturn(null);
6466
try {
6567
APIThrottle.evaluate(opContext, Set.of(event), false);
6668
} catch (Exception ex) {
@@ -76,12 +78,16 @@ public void testExemptions() {
7678
for (String ua : exemptions) {
7779
try {
7880
when(mockRequestContext.getUserAgent()).thenReturn(ua);
81+
when(mockRequestContext.getAgentClass())
82+
.thenReturn(RequestContext.UAA.parse(ua).get(UserAgent.AGENT_CLASS).getValue());
7983
APIThrottle.evaluate(opContext, Set.of(event), true);
8084
} catch (Exception ex) {
8185
Assert.fail("Exception was thrown and NOT expected! " + event);
8286
}
8387
try {
8488
when(mockRequestContext.getUserAgent()).thenReturn(ua);
89+
when(mockRequestContext.getAgentClass())
90+
.thenReturn(RequestContext.UAA.parse(ua).get(UserAgent.AGENT_CLASS).getValue());
8591
APIThrottle.evaluate(opContext, Set.of(event), false);
8692
} catch (Exception ex) {
8793
Assert.fail("Exception was thrown and NOT expected! " + event);
@@ -106,6 +112,8 @@ public void testThrottleException() {
106112
&& !event.getActiveThrottles().contains(MANUAL)) {
107113
try {
108114
when(mockRequestContext.getUserAgent()).thenReturn(ua);
115+
when(mockRequestContext.getAgentClass())
116+
.thenReturn(RequestContext.UAA.parse(ua).get(UserAgent.AGENT_CLASS).getValue());
109117
APIThrottle.evaluate(opContext, Set.of(event), true);
110118
Assert.fail(String.format("Exception WAS expected! %s %s", ua, event));
111119
} catch (Exception ignored) {
@@ -115,6 +123,8 @@ public void testThrottleException() {
115123
&& !event.getActiveThrottles().contains(MANUAL)) {
116124
try {
117125
when(mockRequestContext.getUserAgent()).thenReturn(ua);
126+
when(mockRequestContext.getAgentClass())
127+
.thenReturn(RequestContext.UAA.parse(ua).get(UserAgent.AGENT_CLASS).getValue());
118128
APIThrottle.evaluate(opContext, Set.of(event), true);
119129
} catch (Exception ex) {
120130
Assert.fail(String.format("Exception was thrown and NOT expected! %s %s", ua, event));
@@ -126,6 +136,8 @@ public void testThrottleException() {
126136
&& !event.getActiveThrottles().contains(MANUAL)) {
127137
try {
128138
when(mockRequestContext.getUserAgent()).thenReturn(ua);
139+
when(mockRequestContext.getAgentClass())
140+
.thenReturn(RequestContext.UAA.parse(ua).get(UserAgent.AGENT_CLASS).getValue());
129141
APIThrottle.evaluate(opContext, Set.of(event), false);
130142
Assert.fail(String.format("Exception WAS expected! %s %s", ua, event));
131143
} catch (Exception ignored) {
@@ -135,6 +147,8 @@ public void testThrottleException() {
135147
&& !event.getActiveThrottles().contains(MANUAL)) {
136148
try {
137149
when(mockRequestContext.getUserAgent()).thenReturn(ua);
150+
when(mockRequestContext.getAgentClass())
151+
.thenReturn(RequestContext.UAA.parse(ua).get(UserAgent.AGENT_CLASS).getValue());
138152
APIThrottle.evaluate(opContext, Set.of(event), false);
139153
} catch (Exception ex) {
140154
Assert.fail(String.format("Exception was thrown and NOT expected! %s %s", ua, event));
@@ -145,12 +159,16 @@ public void testThrottleException() {
145159
if (event.getActiveThrottles().contains(MANUAL)) {
146160
try {
147161
when(mockRequestContext.getUserAgent()).thenReturn(ua);
162+
when(mockRequestContext.getAgentClass())
163+
.thenReturn(RequestContext.UAA.parse(ua).get(UserAgent.AGENT_CLASS).getValue());
148164
APIThrottle.evaluate(opContext, Set.of(event), true);
149165
Assert.fail(String.format("Exception WAS expected! %s %s", ua, event));
150166
} catch (Exception ignored) {
151167
}
152168
try {
153169
when(mockRequestContext.getUserAgent()).thenReturn(ua);
170+
when(mockRequestContext.getAgentClass())
171+
.thenReturn(RequestContext.UAA.parse(ua).get(UserAgent.AGENT_CLASS).getValue());
154172
APIThrottle.evaluate(opContext, Set.of(event), false);
155173
Assert.fail(String.format("Exception WAS expected! %s %s", ua, event));
156174
} catch (Exception ignored) {

‎metadata-operation-context/build.gradle

+3
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,9 @@ dependencies {
77
api project(':metadata-utils')
88
api project(':metadata-auth:auth-api')
99

10+
// https://mvnrepository.com/artifact/nl.basjes.parse.useragent/yauaa
11+
api 'nl.basjes.parse.useragent:yauaa:7.30.0'
12+
1013
implementation externalDependency.slf4jApi
1114
implementation externalDependency.servletApi
1215
implementation spec.product.pegasus.restliServer

‎metadata-operation-context/src/main/java/io/datahubproject/metadata/context/RequestContext.java

+58-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
11
package io.datahubproject.metadata.context;
22

3+
import static com.linkedin.metadata.Constants.DATAHUB_ACTOR;
4+
import static com.linkedin.metadata.Constants.SYSTEM_ACTOR;
5+
36
import com.google.common.net.HttpHeaders;
7+
import com.linkedin.metadata.utils.metrics.MetricUtils;
48
import com.linkedin.restli.server.ResourceContext;
59
import io.opentelemetry.api.trace.Span;
610
import jakarta.servlet.http.HttpServletRequest;
@@ -16,11 +20,20 @@
1620
import lombok.Builder;
1721
import lombok.Getter;
1822
import lombok.extern.slf4j.Slf4j;
23+
import nl.basjes.parse.useragent.UserAgent;
24+
import nl.basjes.parse.useragent.UserAgentAnalyzer;
1925

2026
@Slf4j
2127
@Getter
2228
@Builder
2329
public class RequestContext implements ContextInterface {
30+
public static final UserAgentAnalyzer UAA =
31+
UserAgentAnalyzer.newBuilder()
32+
.hideMatcherLoadStats()
33+
.withField(UserAgent.AGENT_CLASS)
34+
.withCache(1000)
35+
.build();
36+
2437
@Nonnull
2538
public static final RequestContext TEST =
2639
RequestContext.builder().requestID("test").requestAPI(RequestAPI.TEST).build();
@@ -36,7 +49,7 @@ public class RequestContext implements ContextInterface {
3649
@Nonnull private final String requestID;
3750

3851
@Nonnull private final String userAgent;
39-
@Builder.Default private boolean validated = true;
52+
@Nonnull private final String agentClass;
4053

4154
public RequestContext(
4255
@Nonnull String actorUrn,
@@ -49,8 +62,28 @@ public RequestContext(
4962
this.requestAPI = requestAPI;
5063
this.requestID = requestID;
5164
this.userAgent = userAgent;
65+
66+
/*
67+
* "Browser",
68+
* "Robot",
69+
* "Crawler",
70+
* "Mobile App",
71+
* "Email Client",
72+
* "Library",
73+
* "Hacker",
74+
* "Unknown"
75+
*/
76+
if (this.userAgent != null && !this.userAgent.isEmpty()) {
77+
UserAgent ua = UAA.parse(this.userAgent);
78+
this.agentClass = ua.get(UserAgent.AGENT_CLASS).getValue();
79+
} else {
80+
this.agentClass = "Unknown";
81+
}
82+
5283
// Uniform common logging of requests across APIs
5384
log.info(toString());
85+
// API metrics
86+
captureAPIMetrics(this);
5487
}
5588

5689
@Override
@@ -59,6 +92,7 @@ public Optional<Integer> getCacheKeyComponent() {
5992
}
6093

6194
public static class RequestContextBuilder {
95+
6296
private RequestContext build() {
6397

6498
// Add context for tracing
@@ -174,6 +208,26 @@ private static String extractSourceIP(@Nonnull ResourceContext resourceContext)
174208
}
175209
}
176210

211+
private static void captureAPIMetrics(RequestContext requestContext) {
212+
// System user?
213+
final String userCategory;
214+
if (SYSTEM_ACTOR.equals(requestContext.actorUrn)) {
215+
userCategory = "system";
216+
} else if (DATAHUB_ACTOR.equals(requestContext.actorUrn)) {
217+
userCategory = "admin";
218+
} else {
219+
userCategory = "regular";
220+
}
221+
222+
MetricUtils.counter(
223+
String.format(
224+
"requestContext_%s_%s_%s",
225+
userCategory,
226+
requestContext.getAgentClass().toLowerCase().replaceAll("\\s+", ""),
227+
requestContext.getRequestAPI().toString().toLowerCase()))
228+
.inc();
229+
}
230+
177231
@Override
178232
public String toString() {
179233
return "RequestContext{"
@@ -191,6 +245,9 @@ public String toString() {
191245
+ ", userAgent='"
192246
+ userAgent
193247
+ '\''
248+
+ ", agentClass='"
249+
+ agentClass
250+
+ '\''
194251
+ '}';
195252
}
196253

0 commit comments

Comments
 (0)