Skip to content

Commit 6f70884

Browse files
noChargernoCharger
authored
Add OpenSearch metrics (#229)
* Add flint opensearch metrics Signed-off-by: Louis Chu <[email protected]> * Refactor to follow the principle of programming to an interface Signed-off-by: Louis Chu <[email protected]> * Resolve comments from Vamsi Signed-off-by: Louis Chu <[email protected]> * Address comment from Peng Signed-off-by: Louis Chu <[email protected]> --------- Signed-off-by: Louis Chu <[email protected]>
1 parent 48dda0a commit 6f70884

File tree

18 files changed

+502
-90
lines changed

18 files changed

+502
-90
lines changed

‎build.sbt

+2-1
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,8 @@ lazy val flintCore = (project in file("flint-core"))
6767
"org.scalatest" %% "scalatest-flatspec" % "3.2.15" % "test",
6868
"org.scalatestplus" %% "mockito-4-6" % "3.2.15.0" % "test",
6969
"com.stephenn" %% "scalatest-json-jsonassert" % "0.2.5" % "test",
70-
"org.mockito" % "mockito-core" % "2.23.0" % "test",
70+
"org.mockito" % "mockito-core" % "4.6.1" % "test",
71+
"org.mockito" % "mockito-inline" % "4.6.1" % "test",
7172
"org.mockito" % "mockito-junit-jupiter" % "3.12.4" % "test",
7273
"org.junit.jupiter" % "junit-jupiter-api" % "5.9.0" % "test",
7374
"org.junit.jupiter" % "junit-jupiter-engine" % "5.9.0" % "test",

‎flint-core/src/main/java/org/opensearch/flint/core/IRestHighLevelClient.java

+62
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.flint.core;
7+
8+
import org.opensearch.action.bulk.BulkRequest;
9+
import org.opensearch.action.bulk.BulkResponse;
10+
import org.opensearch.action.delete.DeleteRequest;
11+
import org.opensearch.action.delete.DeleteResponse;
12+
import org.opensearch.action.get.GetRequest;
13+
import org.opensearch.action.get.GetResponse;
14+
import org.opensearch.action.index.IndexRequest;
15+
import org.opensearch.action.index.IndexResponse;
16+
import org.opensearch.action.search.ClearScrollRequest;
17+
import org.opensearch.action.search.ClearScrollResponse;
18+
import org.opensearch.action.search.SearchRequest;
19+
import org.opensearch.action.search.SearchResponse;
20+
import org.opensearch.action.search.SearchScrollRequest;
21+
import org.opensearch.action.update.UpdateRequest;
22+
import org.opensearch.action.DocWriteResponse;
23+
import org.opensearch.client.indices.CreateIndexRequest;
24+
import org.opensearch.client.indices.CreateIndexResponse;
25+
import org.opensearch.client.indices.GetIndexRequest;
26+
import org.opensearch.client.indices.GetIndexResponse;
27+
import org.opensearch.action.admin.indices.delete.DeleteIndexRequest;
28+
import org.opensearch.client.RequestOptions;
29+
30+
import java.io.Closeable;
31+
import java.io.IOException;
32+
33+
/**
34+
* Interface for wrapping the OpenSearch High Level REST Client with additional functionality,
35+
* such as metrics tracking.
36+
*/
37+
public interface IRestHighLevelClient extends Closeable {
38+
39+
BulkResponse bulk(BulkRequest bulkRequest, RequestOptions options) throws IOException;
40+
41+
ClearScrollResponse clearScroll(ClearScrollRequest clearScrollRequest, RequestOptions options) throws IOException;
42+
43+
CreateIndexResponse createIndex(CreateIndexRequest createIndexRequest, RequestOptions options) throws IOException;
44+
45+
void deleteIndex(DeleteIndexRequest deleteIndexRequest, RequestOptions options) throws IOException;
46+
47+
DeleteResponse delete(DeleteRequest deleteRequest, RequestOptions options) throws IOException;
48+
49+
GetResponse get(GetRequest getRequest, RequestOptions options) throws IOException;
50+
51+
GetIndexResponse getIndex(GetIndexRequest getIndexRequest, RequestOptions options) throws IOException;
52+
53+
IndexResponse index(IndexRequest indexRequest, RequestOptions options) throws IOException;
54+
55+
Boolean isIndexExists(GetIndexRequest getIndexRequest, RequestOptions options) throws IOException;
56+
57+
SearchResponse search(SearchRequest searchRequest, RequestOptions options) throws IOException;
58+
59+
SearchResponse scroll(SearchScrollRequest searchScrollRequest, RequestOptions options) throws IOException;
60+
61+
DocWriteResponse update(UpdateRequest updateRequest, RequestOptions options) throws IOException;
62+
}

‎flint-core/src/main/java/org/opensearch/flint/core/RestHighLevelClientWrapper.java

+197
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.flint.core;
7+
8+
import org.opensearch.action.admin.indices.delete.DeleteIndexRequest;
9+
import org.opensearch.action.bulk.BulkRequest;
10+
import org.opensearch.action.bulk.BulkResponse;
11+
import org.opensearch.action.delete.DeleteRequest;
12+
import org.opensearch.action.delete.DeleteResponse;
13+
import org.opensearch.action.get.GetRequest;
14+
import org.opensearch.action.get.GetResponse;
15+
import org.opensearch.action.index.IndexRequest;
16+
import org.opensearch.action.index.IndexResponse;
17+
import org.opensearch.action.search.ClearScrollRequest;
18+
import org.opensearch.action.search.ClearScrollResponse;
19+
import org.opensearch.action.search.SearchRequest;
20+
import org.opensearch.action.search.SearchResponse;
21+
import org.opensearch.action.search.SearchScrollRequest;
22+
import org.opensearch.action.update.UpdateRequest;
23+
import org.opensearch.OpenSearchException;
24+
import org.opensearch.action.update.UpdateResponse;
25+
import org.opensearch.client.RequestOptions;
26+
import org.opensearch.client.RestHighLevelClient;
27+
import org.opensearch.client.indices.CreateIndexRequest;
28+
import org.opensearch.client.indices.CreateIndexResponse;
29+
import org.opensearch.client.indices.GetIndexRequest;
30+
import org.opensearch.client.indices.GetIndexResponse;
31+
import org.opensearch.flint.core.metrics.MetricsUtil;
32+
33+
import java.io.IOException;
34+
35+
import static org.opensearch.flint.core.metrics.MetricConstants.*;
36+
37+
/**
38+
* A wrapper class for RestHighLevelClient to facilitate OpenSearch operations
39+
* with integrated metrics tracking.
40+
*/
41+
public class RestHighLevelClientWrapper implements IRestHighLevelClient {
42+
private final RestHighLevelClient client;
43+
44+
/**
45+
* Constructs a new RestHighLevelClientWrapper.
46+
*
47+
* @param client the RestHighLevelClient instance to wrap
48+
*/
49+
public RestHighLevelClientWrapper(RestHighLevelClient client) {
50+
this.client = client;
51+
}
52+
53+
@Override
54+
public BulkResponse bulk(BulkRequest bulkRequest, RequestOptions options) throws IOException {
55+
return execute(OS_WRITE_OP_METRIC_PREFIX, () -> client.bulk(bulkRequest, options));
56+
}
57+
58+
@Override
59+
public ClearScrollResponse clearScroll(ClearScrollRequest clearScrollRequest, RequestOptions options) throws IOException {
60+
return execute(OS_READ_OP_METRIC_PREFIX, () -> client.clearScroll(clearScrollRequest, options));
61+
}
62+
63+
@Override
64+
public CreateIndexResponse createIndex(CreateIndexRequest createIndexRequest, RequestOptions options) throws IOException {
65+
return execute(OS_WRITE_OP_METRIC_PREFIX, () -> client.indices().create(createIndexRequest, options));
66+
}
67+
68+
@Override
69+
public void deleteIndex(DeleteIndexRequest deleteIndexRequest, RequestOptions options) throws IOException {
70+
execute(OS_WRITE_OP_METRIC_PREFIX, () -> client.indices().delete(deleteIndexRequest, options));
71+
}
72+
73+
@Override
74+
public DeleteResponse delete(DeleteRequest deleteRequest, RequestOptions options) throws IOException {
75+
return execute(OS_WRITE_OP_METRIC_PREFIX, () -> client.delete(deleteRequest, options));
76+
}
77+
78+
@Override
79+
public GetResponse get(GetRequest getRequest, RequestOptions options) throws IOException {
80+
return execute(OS_READ_OP_METRIC_PREFIX, () -> client.get(getRequest, options));
81+
}
82+
83+
@Override
84+
public GetIndexResponse getIndex(GetIndexRequest getIndexRequest, RequestOptions options) throws IOException {
85+
return execute(OS_READ_OP_METRIC_PREFIX, () -> client.indices().get(getIndexRequest, options));
86+
}
87+
88+
@Override
89+
public IndexResponse index(IndexRequest indexRequest, RequestOptions options) throws IOException {
90+
return execute(OS_WRITE_OP_METRIC_PREFIX, () -> client.index(indexRequest, options));
91+
}
92+
93+
@Override
94+
public Boolean isIndexExists(GetIndexRequest getIndexRequest, RequestOptions options) throws IOException {
95+
return execute(OS_READ_OP_METRIC_PREFIX, () -> client.indices().exists(getIndexRequest, options));
96+
}
97+
98+
@Override
99+
public SearchResponse search(SearchRequest searchRequest, RequestOptions options) throws IOException {
100+
return execute(OS_READ_OP_METRIC_PREFIX, () -> client.search(searchRequest, options));
101+
}
102+
103+
@Override
104+
public SearchResponse scroll(SearchScrollRequest searchScrollRequest, RequestOptions options) throws IOException {
105+
return execute(OS_READ_OP_METRIC_PREFIX, () -> client.scroll(searchScrollRequest, options));
106+
}
107+
108+
@Override
109+
public UpdateResponse update(UpdateRequest updateRequest, RequestOptions options) throws IOException {
110+
return execute(OS_WRITE_OP_METRIC_PREFIX, () -> client.update(updateRequest, options));
111+
}
112+
113+
/**
114+
* Executes a given operation, tracks metrics, and handles exceptions.
115+
*
116+
* @param metricNamePrefix the prefix for the metric name
117+
* @param operation the operation to execute
118+
* @param <T> the return type of the operation
119+
* @return the result of the operation
120+
* @throws IOException if an I/O exception occurs
121+
*/
122+
private <T> T execute(String metricNamePrefix, IOCallable<T> operation) throws IOException {
123+
try {
124+
T result = operation.call();
125+
recordOperationSuccess(metricNamePrefix);
126+
return result;
127+
} catch (Exception e) {
128+
recordOperationFailure(metricNamePrefix, e);
129+
throw e;
130+
}
131+
}
132+
133+
/**
134+
* Records the success of an OpenSearch operation by incrementing the corresponding metric counter.
135+
* This method constructs the metric name by appending ".200.count" to the provided metric name prefix.
136+
* The metric name is then used to increment the counter, indicating a successful operation.
137+
*
138+
* @param metricNamePrefix the prefix for the metric name which is used to construct the full metric name for success
139+
*/
140+
private void recordOperationSuccess(String metricNamePrefix) {
141+
String successMetricName = metricNamePrefix + ".2xx.count";
142+
MetricsUtil.incrementCounter(successMetricName);
143+
}
144+
145+
/**
146+
* Records the failure of an OpenSearch operation by incrementing the corresponding metric counter.
147+
* If the exception is an OpenSearchException with a specific status code (e.g., 403),
148+
* it increments a metric specifically for that status code.
149+
* Otherwise, it increments a general failure metric counter based on the status code category (e.g., 4xx, 5xx).
150+
*
151+
* @param metricNamePrefix the prefix for the metric name which is used to construct the full metric name for failure
152+
* @param e the exception encountered during the operation, used to determine the type of failure
153+
*/
154+
private void recordOperationFailure(String metricNamePrefix, Exception e) {
155+
OpenSearchException openSearchException = extractOpenSearchException(e);
156+
int statusCode = openSearchException != null ? openSearchException.status().getStatus() : 500;
157+
158+
if (statusCode == 403) {
159+
String forbiddenErrorMetricName = metricNamePrefix + ".403.count";
160+
MetricsUtil.incrementCounter(forbiddenErrorMetricName);
161+
}
162+
163+
String failureMetricName = metricNamePrefix + "." + (statusCode / 100) + "xx.count";
164+
MetricsUtil.incrementCounter(failureMetricName);
165+
}
166+
167+
/**
168+
* Extracts an OpenSearchException from the given Throwable.
169+
* Checks if the Throwable is an instance of OpenSearchException or caused by one.
170+
*
171+
* @param ex the exception to be checked
172+
* @return the extracted OpenSearchException, or null if not found
173+
*/
174+
private OpenSearchException extractOpenSearchException(Throwable ex) {
175+
if (ex instanceof OpenSearchException) {
176+
return (OpenSearchException) ex;
177+
} else if (ex.getCause() instanceof OpenSearchException) {
178+
return (OpenSearchException) ex.getCause();
179+
}
180+
return null;
181+
}
182+
183+
/**
184+
* Functional interface for operations that can throw IOException.
185+
*
186+
* @param <T> the return type of the operation
187+
*/
188+
@FunctionalInterface
189+
private interface IOCallable<T> {
190+
T call() throws IOException;
191+
}
192+
193+
@Override
194+
public void close() throws IOException {
195+
client.close();
196+
}
197+
}

‎flint-core/src/main/java/org/opensearch/flint/core/metrics/MetricConstants.java

+24
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.flint.core.metrics;
7+
8+
/**
9+
* This class defines custom metric constants used for monitoring flint operations.
10+
*/
11+
public class MetricConstants {
12+
13+
/**
14+
* The prefix for all read-related metrics in OpenSearch.
15+
* This constant is used as a part of metric names to categorize and identify metrics related to read operations.
16+
*/
17+
public static final String OS_READ_OP_METRIC_PREFIX = "opensearch.read";
18+
19+
/**
20+
* The prefix for all write-related metrics in OpenSearch.
21+
* Similar to OS_READ_METRIC_PREFIX, this constant is used for categorizing and identifying metrics that pertain to write operations.
22+
*/
23+
public static final String OS_WRITE_OP_METRIC_PREFIX = "opensearch.write";
24+
}

‎flint-core/src/main/java/org/opensearch/flint/core/metrics/MetricsUtil.java

+68
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.flint.core.metrics;
7+
8+
import com.codahale.metrics.Counter;
9+
import org.apache.spark.SparkEnv;
10+
import org.apache.spark.metrics.source.FlintMetricSource;
11+
import org.apache.spark.metrics.source.Source;
12+
import scala.collection.Seq;
13+
14+
import java.util.logging.Logger;
15+
16+
/**
17+
* Utility class for managing metrics in the OpenSearch Flint context.
18+
*/
19+
public final class MetricsUtil {
20+
21+
private static final Logger LOG = Logger.getLogger(MetricsUtil.class.getName());
22+
23+
// Private constructor to prevent instantiation
24+
private MetricsUtil() {
25+
}
26+
27+
/**
28+
* Increments the Counter metric associated with the given metric name.
29+
* If the counter does not exist, it is created before being incremented.
30+
*
31+
* @param metricName The name of the metric for which the counter is incremented.
32+
* This name is used to retrieve or create the counter.
33+
*/
34+
public static void incrementCounter(String metricName) {
35+
Counter counter = getOrCreateCounter(metricName);
36+
if (counter != null) {
37+
counter.inc();
38+
}
39+
}
40+
41+
// Retrieves or creates a new counter for the given metric name
42+
private static Counter getOrCreateCounter(String metricName) {
43+
SparkEnv sparkEnv = SparkEnv.get();
44+
if (sparkEnv == null) {
45+
LOG.warning("Spark environment not available, cannot instrument metric: " + metricName);
46+
return null;
47+
}
48+
49+
FlintMetricSource flintMetricSource = getOrInitFlintMetricSource(sparkEnv);
50+
Counter counter = flintMetricSource.metricRegistry().getCounters().get(metricName);
51+
if (counter == null) {
52+
counter = flintMetricSource.metricRegistry().counter(metricName);
53+
}
54+
return counter;
55+
}
56+
57+
// Gets or initializes the FlintMetricSource
58+
private static FlintMetricSource getOrInitFlintMetricSource(SparkEnv sparkEnv) {
59+
Seq<Source> metricSourceSeq = sparkEnv.metricsSystem().getSourcesByName(FlintMetricSource.FLINT_METRIC_SOURCE_NAME());
60+
61+
if (metricSourceSeq == null || metricSourceSeq.isEmpty()) {
62+
FlintMetricSource metricSource = new FlintMetricSource();
63+
sparkEnv.metricsSystem().registerSource(metricSource);
64+
return metricSource;
65+
}
66+
return (FlintMetricSource) metricSourceSeq.head();
67+
}
68+
}

0 commit comments

Comments
 (0)