11
11
import com .codahale .metrics .Timer ;
12
12
import org .apache .spark .SparkEnv ;
13
13
import org .apache .spark .metrics .source .FlintMetricSource ;
14
+ import org .apache .spark .metrics .source .FlintIndexMetricSource ;
14
15
import org .apache .spark .metrics .source .Source ;
15
16
import scala .collection .Seq ;
16
17
@@ -33,10 +34,20 @@ private MetricsUtil() {
33
34
* If the counter does not exist, it is created before being incremented.
34
35
*
35
36
* @param metricName The name of the metric for which the counter is incremented.
36
- * This name is used to retrieve or create the counter.
37
37
*/
38
38
public static void incrementCounter (String metricName ) {
39
- Counter counter = getOrCreateCounter (metricName );
39
+ incrementCounter (metricName , false );
40
+ }
41
+
42
+ /**
43
+ * Increments the Counter metric associated with the given metric name.
44
+ * If the counter does not exist, it is created before being incremented.
45
+ *
46
+ * @param metricName The name of the metric for which the counter is incremented.
47
+ * @param isIndexMetric Whether this metric is an index-specific metric.
48
+ */
49
+ public static void incrementCounter (String metricName , boolean isIndexMetric ) {
50
+ Counter counter = getOrCreateCounter (metricName , isIndexMetric );
40
51
if (counter != null ) {
41
52
counter .inc ();
42
53
}
@@ -48,29 +59,48 @@ public static void incrementCounter(String metricName) {
48
59
* @param metricName The name of the metric counter to be decremented.
49
60
*/
50
61
public static void decrementCounter (String metricName ) {
51
- Counter counter = getOrCreateCounter (metricName );
62
+ decrementCounter (metricName , false );
63
+ }
64
+
65
+ /**
66
+ * Decrements the value of the specified metric counter by one, if the counter exists and its current count is greater than zero.
67
+ *
68
+ * @param metricName The name of the metric counter to be decremented.
69
+ * @param isIndexMetric Whether this metric is an index-specific metric.
70
+ */
71
+ public static void decrementCounter (String metricName , boolean isIndexMetric ) {
72
+ Counter counter = getOrCreateCounter (metricName , isIndexMetric );
52
73
if (counter != null && counter .getCount () > 0 ) {
53
74
counter .dec ();
54
75
}
55
76
}
56
77
57
78
/**
58
79
* Retrieves a {@link Timer.Context} for the specified metric name, creating a new timer if one does not already exist.
59
- * This context can be used to measure the duration of a particular operation or event.
60
80
*
61
81
* @param metricName The name of the metric timer to retrieve the context for.
62
82
* @return A {@link Timer.Context} instance for timing operations, or {@code null} if the timer could not be created or retrieved.
63
83
*/
64
84
public static Timer .Context getTimerContext (String metricName ) {
65
- Timer timer = getOrCreateTimer (metricName );
85
+ return getTimerContext (metricName , false );
86
+ }
87
+
88
+ /**
89
+ * Retrieves a {@link Timer.Context} for the specified metric name, creating a new timer if one does not already exist.
90
+ *
91
+ * @param metricName The name of the metric timer to retrieve the context for.
92
+ * @param isIndexMetric Whether this metric is an index-specific metric.
93
+ * @return A {@link Timer.Context} instance for timing operations, or {@code null} if the timer could not be created or retrieved.
94
+ */
95
+ public static Timer .Context getTimerContext (String metricName , boolean isIndexMetric ) {
96
+ Timer timer = getOrCreateTimer (metricName , isIndexMetric );
66
97
return timer != null ? timer .time () : null ;
67
98
}
68
99
69
100
/**
70
- * Stops the timer associated with the given {@link Timer.Context}, effectively recording the elapsed time since the timer was started
71
- * and returning the duration. If the context is {@code null}, this method does nothing and returns {@code null}.
101
+ * Stops the timer associated with the given {@link Timer.Context}.
72
102
*
73
- * @param context The {@link Timer.Context} to stop. May be {@code null}, in which case this method has no effect and returns {@code null} .
103
+ * @param context The {@link Timer.Context} to stop. May be {@code null}.
74
104
* @return The elapsed time in nanoseconds since the timer was started, or {@code null} if the context was {@code null}.
75
105
*/
76
106
public static Long stopTimer (Timer .Context context ) {
@@ -79,53 +109,61 @@ public static Long stopTimer(Timer.Context context) {
79
109
80
110
/**
81
111
* Registers a gauge metric with the provided name and value.
82
- * The gauge will reflect the current value of the AtomicInteger provided.
83
112
*
84
113
* @param metricName The name of the gauge metric to register.
85
- * @param value The AtomicInteger whose current value should be reflected by the gauge.
114
+ * @param value The AtomicInteger whose current value should be reflected by the gauge.
86
115
*/
87
116
public static void registerGauge (String metricName , final AtomicInteger value ) {
88
- MetricRegistry metricRegistry = getMetricRegistry ();
117
+ registerGauge (metricName , value , false );
118
+ }
119
+
120
+ /**
121
+ * Registers a gauge metric with the provided name and value.
122
+ *
123
+ * @param metricName The name of the gauge metric to register.
124
+ * @param value The AtomicInteger whose current value should be reflected by the gauge.
125
+ * @param isIndexMetric Whether this metric is an index-specific metric.
126
+ */
127
+ public static void registerGauge (String metricName , final AtomicInteger value , boolean isIndexMetric ) {
128
+ MetricRegistry metricRegistry = getMetricRegistry (isIndexMetric );
89
129
if (metricRegistry == null ) {
90
130
LOG .warning ("MetricRegistry not available, cannot register gauge: " + metricName );
91
131
return ;
92
132
}
93
133
metricRegistry .register (metricName , (Gauge <Integer >) value ::get );
94
134
}
95
135
96
- // Retrieves or creates a new counter for the given metric name
97
- private static Counter getOrCreateCounter (String metricName ) {
98
- MetricRegistry metricRegistry = getMetricRegistry ();
136
+ private static Counter getOrCreateCounter (String metricName , boolean isIndexMetric ) {
137
+ MetricRegistry metricRegistry = getMetricRegistry (isIndexMetric );
99
138
return metricRegistry != null ? metricRegistry .counter (metricName ) : null ;
100
139
}
101
140
102
- // Retrieves or creates a new Timer for the given metric name
103
- private static Timer getOrCreateTimer (String metricName ) {
104
- MetricRegistry metricRegistry = getMetricRegistry ();
141
+ private static Timer getOrCreateTimer (String metricName , boolean isIndexMetric ) {
142
+ MetricRegistry metricRegistry = getMetricRegistry (isIndexMetric );
105
143
return metricRegistry != null ? metricRegistry .timer (metricName ) : null ;
106
144
}
107
145
108
- // Retrieves the MetricRegistry from the current Spark environment.
109
- private static MetricRegistry getMetricRegistry () {
146
+ private static MetricRegistry getMetricRegistry (boolean isIndexMetric ) {
110
147
SparkEnv sparkEnv = SparkEnv .get ();
111
148
if (sparkEnv == null ) {
112
149
LOG .warning ("Spark environment not available, cannot access MetricRegistry." );
113
150
return null ;
114
151
}
115
152
116
- FlintMetricSource flintMetricSource = getOrInitFlintMetricSource (sparkEnv );
117
- return flintMetricSource .metricRegistry ();
153
+ Source metricSource = isIndexMetric ?
154
+ getOrInitMetricSource (sparkEnv , FlintMetricSource .FLINT_INDEX_METRIC_SOURCE_NAME (), FlintIndexMetricSource ::new ) :
155
+ getOrInitMetricSource (sparkEnv , FlintMetricSource .FLINT_METRIC_SOURCE_NAME (), FlintMetricSource ::new );
156
+ return metricSource .metricRegistry ();
118
157
}
119
158
120
- // Gets or initializes the FlintMetricSource
121
- private static FlintMetricSource getOrInitFlintMetricSource (SparkEnv sparkEnv ) {
122
- Seq <Source > metricSourceSeq = sparkEnv .metricsSystem ().getSourcesByName (FlintMetricSource .FLINT_METRIC_SOURCE_NAME ());
159
+ private static Source getOrInitMetricSource (SparkEnv sparkEnv , String sourceName , java .util .function .Supplier <Source > sourceSupplier ) {
160
+ Seq <Source > metricSourceSeq = sparkEnv .metricsSystem ().getSourcesByName (sourceName );
123
161
124
162
if (metricSourceSeq == null || metricSourceSeq .isEmpty ()) {
125
- FlintMetricSource metricSource = new FlintMetricSource ();
163
+ Source metricSource = sourceSupplier . get ();
126
164
sparkEnv .metricsSystem ().registerSource (metricSource );
127
165
return metricSource ;
128
166
}
129
- return ( FlintMetricSource ) metricSourceSeq .head ();
167
+ return metricSourceSeq .head ();
130
168
}
131
169
}
0 commit comments