This is an automated email from the ASF dual-hosted git repository.
SteNicholas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new 50323e1f3 [CELEBORN-2316] Introduce metadata operation metrics
50323e1f3 is described below
commit 50323e1f323c9432692fcc65bc703be107395288
Author: AmandeepSingh285 <[email protected]>
AuthorDate: Wed May 13 14:47:38 2026 +0800
[CELEBORN-2316] Introduce metadata operation metrics
## What changes were proposed in this pull request?
Introduce metrics around RocksDB operations. Metrics to have success and
failure count for metadata operations for RocksDB observability.
### Why are the changes needed?
Introduce metrics around RocksDB metadata operations. Current
implementation, metadata operations do not have any observability added.
RocksDB goes into a read only mode when any critical errors are encountered
which results in all write operations failing. Observability around metadata
operations is critical and failures could result in metadata entering an
inconsistent state.
### Does this PR resolve a correctness bug?
No.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Tested in local staging setup -
<img width="2166" height="314" alt="Screenshot 2026-05-11 at 3 06 05 PM"
src="https://github.com/user-attachments/assets/f6f75515-f937-436d-9ee3-de6ca83bfcdf";
/>
<img width="1113" height="362" alt="image"
src="https://github.com/user-attachments/assets/89cc0b40-e96c-420e-a5aa-36849fb15fd4";
/>
Closes #3673 from AmandeepSingh285/adding-metadata-metrics.
Lead-authored-by: AmandeepSingh285 <[email protected]>
Co-authored-by: amandeeps.28 <[email protected]>
Signed-off-by: SteNicholas <[email protected]>
---
assets/grafana/celeborn-dashboard.json | 176 ++++++++++++
docs/monitoring.md | 1 +
.../service/deploy/worker/shuffledb/DB.java | 55 +++-
.../deploy/worker/shuffledb/DBProvider.java | 9 +-
.../service/deploy/worker/shuffledb/LevelDB.java | 23 +-
.../deploy/worker/shuffledb/LevelDBIterator.java | 19 +-
.../deploy/worker/shuffledb/MetadataMetrics.java | 118 +++++++++
.../service/deploy/worker/shuffledb/RocksDB.java | 56 ++--
.../deploy/worker/shuffledb/RocksDBIterator.java | 29 +-
.../worker/storage/PartitionFilesSorter.java | 2 +-
.../service/deploy/worker/WorkerSource.scala | 28 ++
.../deploy/worker/storage/StorageManager.scala | 2 +-
.../deploy/worker/shuffledb/DBProviderSuiteJ.java | 9 +-
.../worker/shuffledb/MetadataMetricsSuiteJ.java | 294 +++++++++++++++++++++
14 files changed, 750 insertions(+), 71 deletions(-)
diff --git a/assets/grafana/celeborn-dashboard.json
b/assets/grafana/celeborn-dashboard.json
index cff188291..65a984d67 100644
--- a/assets/grafana/celeborn-dashboard.json
+++ b/assets/grafana/celeborn-dashboard.json
@@ -4063,6 +4063,182 @@
],
"title": "metrics_PartitionFileSizeBytes_Max",
"type": "timeseries"
+ },
+ {
+ "datasource": {
+ "type": "prometheus",
+ "uid": "${DS_PROMETHEUS}"
+ },
+ "description": "Cumulative count of metadata DB (RocksDB) operations
performed by the worker, broken down by operation (read|write) and status
(success|fail).",
+ "fieldConfig": {
+ "defaults": {
+ "color": {
+ "mode": "palette-classic"
+ },
+ "custom": {
+ "axisCenteredZero": false,
+ "axisColorMode": "text",
+ "axisLabel": "",
+ "axisPlacement": "auto",
+ "barAlignment": 0,
+ "drawStyle": "line",
+ "fillOpacity": 0,
+ "gradientMode": "none",
+ "hideFrom": {
+ "legend": false,
+ "tooltip": false,
+ "viz": false
+ },
+ "lineInterpolation": "linear",
+ "lineWidth": 1,
+ "pointSize": 5,
+ "scaleDistribution": {
+ "type": "linear"
+ },
+ "showPoints": "auto",
+ "spanNulls": false,
+ "stacking": {
+ "group": "A",
+ "mode": "none"
+ },
+ "thresholdsStyle": {
+ "mode": "off"
+ }
+ },
+ "mappings": [],
+ "thresholds": {
+ "mode": "absolute",
+ "steps": [
+ {
+ "color": "green"
+ }
+ ]
+ }
+ },
+ "overrides": []
+ },
+ "gridPos": {
+ "h": 8,
+ "w": 12,
+ "x": 0,
+ "y": 216
+ },
+ "id": 270,
+ "options": {
+ "legend": {
+ "calcs": [],
+ "displayMode": "list",
+ "placement": "bottom",
+ "showLegend": true
+ },
+ "tooltip": {
+ "mode": "single",
+ "sort": "none"
+ }
+ },
+ "targets": [
+ {
+ "datasource": {
+ "type": "prometheus",
+ "uid": "${DS_PROMETHEUS}"
+ },
+ "editorMode": "code",
+ "expr":
"metrics_MetadataOperationStatusCount_Count{role=\"Worker\",
instance=~\"${instance}\"}",
+ "legendFormat": "{{dbBackend}} {{operation}}/{{status}}
${baseLegend}",
+ "range": true,
+ "refId": "A"
+ }
+ ],
+ "title": "metrics_MetadataOperationStatusCount_Count",
+ "type": "timeseries"
+ },
+ {
+ "datasource": {
+ "type": "prometheus",
+ "uid": "${DS_PROMETHEUS}"
+ },
+ "description": "Number of metadata DB (RocksDB) operations on the
worker over the last 1 minute, by operation (read|write) and status
(success|fail). Useful for spotting RocksDB error spikes that the cumulative
chart visually flattens.",
+ "fieldConfig": {
+ "defaults": {
+ "color": {
+ "mode": "palette-classic"
+ },
+ "custom": {
+ "axisCenteredZero": false,
+ "axisColorMode": "text",
+ "axisLabel": "",
+ "axisPlacement": "auto",
+ "barAlignment": 0,
+ "drawStyle": "line",
+ "fillOpacity": 0,
+ "gradientMode": "none",
+ "hideFrom": {
+ "legend": false,
+ "tooltip": false,
+ "viz": false
+ },
+ "lineInterpolation": "linear",
+ "lineWidth": 1,
+ "pointSize": 5,
+ "scaleDistribution": {
+ "type": "linear"
+ },
+ "showPoints": "auto",
+ "spanNulls": false,
+ "stacking": {
+ "group": "A",
+ "mode": "none"
+ },
+ "thresholdsStyle": {
+ "mode": "off"
+ }
+ },
+ "mappings": [],
+ "thresholds": {
+ "mode": "absolute",
+ "steps": [
+ {
+ "color": "green"
+ }
+ ]
+ }
+ },
+ "overrides": []
+ },
+ "gridPos": {
+ "h": 8,
+ "w": 12,
+ "x": 12,
+ "y": 216
+ },
+ "id": 271,
+ "options": {
+ "legend": {
+ "calcs": [],
+ "displayMode": "list",
+ "placement": "bottom",
+ "showLegend": true
+ },
+ "tooltip": {
+ "mode": "single",
+ "sort": "none"
+ }
+ },
+ "targets": [
+ {
+ "datasource": {
+ "type": "prometheus",
+ "uid": "${DS_PROMETHEUS}"
+ },
+ "editorMode": "code",
+ "expr":
"increase(metrics_MetadataOperationStatusCount_Count{role=\"Worker\",
instance=~\"${instance}\"}[1m])",
+ "legendFormat": "{{dbBackend}} {{operation}}/{{status}}
${baseLegend}",
+ "range": true,
+ "refId": "A"
+ }
+ ],
+ "title": "metrics_MetadataOperationStatusCount_increase_1m",
+ "type": "timeseries"
}
],
"title": "Worker",
diff --git a/docs/monitoring.md b/docs/monitoring.md
index cef191cd2..aa845640a 100644
--- a/docs/monitoring.md
+++ b/docs/monitoring.md
@@ -281,6 +281,7 @@ These metrics are exposed by Celeborn worker.
| OssFlushSize | The size of data flushed to
OSS.
|
| S3FlushCount | The amount of data flushed to
S3.
|
| S3FlushSize | The size of data flushed to S3.
|
+ | MetadataOperationStatusCount | The count of metadata DB
operations performed by the worker |
| push_usedHeapMemory |
|
| push_usedDirectMemory |
|
| push_numHeapArenas |
|
diff --git
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/DB.java
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/DB.java
index 20b275d8d..c199c0645 100644
---
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/DB.java
+++
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/DB.java
@@ -19,23 +19,62 @@ package org.apache.celeborn.service.deploy.worker.shuffledb;
import java.io.Closeable;
-/** Note: code copied from Apache Spark. */
-public interface DB extends Closeable {
+import org.apache.celeborn.common.metrics.source.AbstractSource;
+
+/**
+ * Common base class for the metadata KV store.
+ *
+ * <p>This class owns the {@link MetadataMetrics} recorder so that every
backend (RocksDB, LevelDB,
+ * ...) gets read/write success/fail counters for free, without having to
repeat the metric plumbing
+ * in each implementation. Subclasses only implement the raw backend-specific
primitives ({@code
+ * putInternal}, {@code getInternal}, {@code deleteInternal}, {@code
newIterator}); the public API
+ * methods below are {@code final} and wrap those primitives with the metric
recording.
+ *
+ * <p>Note: code copied from Apache Spark.
+ */
+public abstract class DB implements Closeable {
+
+ private final MetadataMetrics metrics;
+
+ protected DB(AbstractSource source, DBBackend dbBackend) {
+ this.metrics = new MetadataMetrics(source, dbBackend);
+ }
+
/** Set the DB entry for "key" to "value". */
- void put(byte[] key, byte[] value);
+ public final void put(byte[] key, byte[] value) {
+ metrics.onWrite(() -> putInternal(key, value));
+ }
- /** Set the DB entry for "key" to "value". Support Sync option */
- void put(byte[] key, byte[] value, boolean sync);
+ /** Set the DB entry for "key" to "value". Support Sync option. */
+ public final void put(byte[] key, byte[] value, boolean sync) {
+ metrics.onWrite(() -> putInternal(key, value, sync));
+ }
/**
* Get which returns a new byte array storing the value associated with the
specified input key if
* any.
*/
- byte[] get(byte[] key);
+ public final byte[] get(byte[] key) {
+ return metrics.onRead(() -> getInternal(key));
+ }
/** Delete the DB entry (if any) for "key". */
- void delete(byte[] key);
+ public final void delete(byte[] key) {
+ metrics.onWrite(() -> deleteInternal(key));
+ }
/** Return an iterator over the contents of the DB. */
- DBIterator iterator();
+ public final DBIterator iterator() {
+ return newIterator(metrics);
+ }
+
+ protected abstract void putInternal(byte[] key, byte[] value) throws
Exception;
+
+ protected abstract void putInternal(byte[] key, byte[] value, boolean sync)
throws Exception;
+
+ protected abstract byte[] getInternal(byte[] key) throws Exception;
+
+ protected abstract void deleteInternal(byte[] key) throws Exception;
+
+ protected abstract DBIterator newIterator(MetadataMetrics metrics);
}
diff --git
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/DBProvider.java
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/DBProvider.java
index 41df2a4dc..8022a62e3 100644
---
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/DBProvider.java
+++
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/DBProvider.java
@@ -23,21 +23,24 @@ import java.io.IOException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import org.apache.celeborn.common.metrics.source.AbstractSource;
+
/** Note: code copied from Apache Spark. */
public class DBProvider {
private static final Logger logger =
LoggerFactory.getLogger(DBProvider.class);
- public static DB initDB(DBBackend dbBackend, File dbFile, StoreVersion
version)
+ public static DB initDB(
+ DBBackend dbBackend, File dbFile, StoreVersion version, AbstractSource
source)
throws IOException {
if (dbFile != null) {
switch (dbBackend) {
case LEVELDB:
org.iq80.leveldb.DB levelDB = LevelDBProvider.initLevelDB(dbFile,
version);
logger.warn("The LEVELDB is deprecated. Please use ROCKSDB
instead.");
- return levelDB != null ? new LevelDB(levelDB) : null;
+ return levelDB != null ? new LevelDB(levelDB, source, dbBackend) :
null;
case ROCKSDB:
org.rocksdb.RocksDB rocksDB = RocksDBProvider.initRockDB(dbFile,
version);
- return rocksDB != null ? new RocksDB(rocksDB) : null;
+ return rocksDB != null ? new RocksDB(rocksDB, source, dbBackend) :
null;
default:
throw new IllegalArgumentException("Unsupported DBBackend: " +
dbBackend);
}
diff --git
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/LevelDB.java
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/LevelDB.java
index 2dfeadb46..c4430398d 100644
---
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/LevelDB.java
+++
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/LevelDB.java
@@ -21,22 +21,25 @@ import java.io.IOException;
import org.iq80.leveldb.WriteOptions;
+import org.apache.celeborn.common.metrics.source.AbstractSource;
+
/** Note: code copied from Apache Spark. */
-public class LevelDB implements DB {
+public class LevelDB extends DB {
private final org.iq80.leveldb.DB db;
private final WriteOptions SYNC_WRITE_OPTIONS = new
WriteOptions().sync(true);
- public LevelDB(org.iq80.leveldb.DB db) {
+ public LevelDB(org.iq80.leveldb.DB db, AbstractSource source, DBBackend
dbBackend) {
+ super(source, dbBackend);
this.db = db;
}
@Override
- public void put(byte[] key, byte[] value) {
+ protected void putInternal(byte[] key, byte[] value) {
db.put(key, value);
}
@Override
- public void put(byte[] key, byte[] value, boolean sync) {
+ protected void putInternal(byte[] key, byte[] value, boolean sync) {
if (sync) {
db.put(key, value, SYNC_WRITE_OPTIONS);
} else {
@@ -45,22 +48,22 @@ public class LevelDB implements DB {
}
@Override
- public byte[] get(byte[] key) {
+ protected byte[] getInternal(byte[] key) {
return db.get(key);
}
@Override
- public void delete(byte[] key) {
+ protected void deleteInternal(byte[] key) {
db.delete(key);
}
@Override
- public void close() throws IOException {
- db.close();
+ protected DBIterator newIterator(MetadataMetrics metrics) {
+ return new LevelDBIterator(db.iterator(), metrics);
}
@Override
- public DBIterator iterator() {
- return new LevelDBIterator(db.iterator());
+ public void close() throws IOException {
+ db.close();
}
}
diff --git
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/LevelDBIterator.java
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/LevelDBIterator.java
index 4f136e500..7a534b94f 100644
---
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/LevelDBIterator.java
+++
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/LevelDBIterator.java
@@ -26,14 +26,17 @@ public class LevelDBIterator implements DBIterator {
private final org.iq80.leveldb.DBIterator it;
+ private final MetadataMetrics metrics;
+
private boolean checkedNext;
private boolean closed;
private Map.Entry<byte[], byte[]> next;
- public LevelDBIterator(org.iq80.leveldb.DBIterator it) {
+ public LevelDBIterator(org.iq80.leveldb.DBIterator it, MetadataMetrics
metrics) {
this.it = it;
+ this.metrics = metrics;
}
@Override
@@ -74,14 +77,16 @@ public class LevelDBIterator implements DBIterator {
@Override
public void seek(byte[] key) {
- it.seek(key);
+ metrics.onRead(() -> it.seek(key));
}
private Map.Entry<byte[], byte[]> loadNext() {
- boolean hasNext = it.hasNext();
- if (!hasNext) {
- return null;
- }
- return it.next();
+ return metrics.onRead(
+ () -> {
+ if (!it.hasNext()) {
+ return null;
+ }
+ return it.next();
+ });
}
}
diff --git
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/MetadataMetrics.java
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/MetadataMetrics.java
new file mode 100644
index 000000000..f972f5cbf
--- /dev/null
+++
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/MetadataMetrics.java
@@ -0,0 +1,118 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.celeborn.service.deploy.worker.shuffledb;
+
+import scala.collection.immutable.Map;
+
+import org.apache.celeborn.common.metrics.source.AbstractSource;
+import org.apache.celeborn.service.deploy.worker.WorkerSource;
+
+/**
+ * Records success/fail counters for metadata DB operations.
+ *
+ * <p>Callers express the operation kind (read or write) and the action; the
recorder owns the
+ * metric name, label selection, and the try/catch ceremony.
+ *
+ * <p>The functional interfaces declare {@code throws Exception} so that any
backend (RocksDB,
+ * LevelDB, ...) can hand off its raw operation without leaking
backend-specific exception types
+ * into this class.
+ */
+class MetadataMetrics {
+
+ @FunctionalInterface
+ interface ThrowingRunnable {
+ void run() throws Exception;
+ }
+
+ @FunctionalInterface
+ interface ThrowingSupplier<T> {
+ T get() throws Exception;
+ }
+
+ private final AbstractSource source;
+ private final Map<String, String> writeSuccessLabels;
+ private final Map<String, String> writeFailLabels;
+ private final Map<String, String> readSuccessLabels;
+ private final Map<String, String> readFailLabels;
+
+ MetadataMetrics(AbstractSource source, DBBackend dbBackend) {
+ // Register the four label combinations so callers see the counters even
before the first op
+ // (addCounter is putIfAbsent, so the call is safe even if something else
registered them).
+ this.source = source;
+ String backend = dbBackend.name();
+
+ this.writeSuccessLabels = WorkerSource.WRITE_SUCCESS_COUNT_LABELS(backend);
+ this.writeFailLabels = WorkerSource.WRITE_FAIL_COUNT_LABELS(backend);
+ this.readSuccessLabels = WorkerSource.READ_SUCCESS_COUNT_LABELS(backend);
+ this.readFailLabels = WorkerSource.READ_FAIL_COUNT_LABELS(backend);
+ if (source != null) {
+ String name = WorkerSource.METADATA_OPERATION_STATUS_COUNT();
+ source.addCounter(name, writeSuccessLabels);
+ source.addCounter(name, writeFailLabels);
+ source.addCounter(name, readSuccessLabels);
+ source.addCounter(name, readFailLabels);
+ }
+ }
+
+ void onWrite(ThrowingRunnable action) {
+ record(
+ () -> {
+ action.run();
+ return null;
+ },
+ writeSuccessLabels,
+ writeFailLabels);
+ }
+
+ void onRead(ThrowingRunnable action) {
+ record(
+ () -> {
+ action.run();
+ return null;
+ },
+ readSuccessLabels,
+ readFailLabels);
+ }
+
+ <T> T onRead(ThrowingSupplier<T> action) {
+ return record(action, readSuccessLabels, readFailLabels);
+ }
+
+ private <T> T record(
+ ThrowingSupplier<T> action,
+ Map<String, String> successLabels,
+ Map<String, String> failLabels) {
+ try {
+ T result = action.get();
+ incCounter(successLabels);
+ return result;
+ } catch (Exception e) {
+ incCounter(failLabels);
+ if (e instanceof RuntimeException) {
+ throw (RuntimeException) e;
+ }
+ throw new RuntimeException(e);
+ }
+ }
+
+ private void incCounter(Map<String, String> labels) {
+ if (source != null) {
+ source.incCounter(WorkerSource.METADATA_OPERATION_STATUS_COUNT(), 1,
labels);
+ }
+ }
+}
diff --git
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDB.java
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDB.java
index 5e262393e..b45988ca0 100644
---
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDB.java
+++
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDB.java
@@ -22,66 +22,58 @@ import java.io.IOException;
import org.rocksdb.RocksDBException;
import org.rocksdb.WriteOptions;
+import org.apache.celeborn.common.metrics.source.AbstractSource;
+
/**
* RocksDB implementation of the local KV storage used to persist the shuffle
state.
*
* <p>Note: code copied from Apache Spark.
*/
-public class RocksDB implements DB {
+public class RocksDB extends DB {
private final org.rocksdb.RocksDB db;
private final WriteOptions SYNC_WRITE_OPTIONS = new
WriteOptions().setSync(true);
- public RocksDB(org.rocksdb.RocksDB db) {
+ public RocksDB(org.rocksdb.RocksDB db, AbstractSource source, DBBackend
dbBackend) {
+ super(source, dbBackend);
this.db = db;
}
@Override
- public void put(byte[] key, byte[] value) {
- try {
- db.put(key, value);
- } catch (RocksDBException e) {
- throw new RuntimeException(e);
- }
+ protected void putInternal(byte[] key, byte[] value) throws RocksDBException
{
+ db.put(key, value);
}
@Override
- public void put(byte[] key, byte[] value, boolean sync) {
- try {
- if (sync) {
- db.put(SYNC_WRITE_OPTIONS, key, value);
- } else {
- db.put(key, value);
- }
- } catch (RocksDBException e) {
- throw new RuntimeException(e);
+ protected void putInternal(byte[] key, byte[] value, boolean sync) throws
RocksDBException {
+ if (sync) {
+ db.put(SYNC_WRITE_OPTIONS, key, value);
+ } else {
+ db.put(key, value);
}
}
@Override
- public byte[] get(byte[] key) {
- try {
- return db.get(key);
- } catch (RocksDBException e) {
- throw new RuntimeException(e);
- }
+ protected byte[] getInternal(byte[] key) throws RocksDBException {
+ return db.get(key);
}
@Override
- public void delete(byte[] key) {
- try {
- db.delete(key);
- } catch (RocksDBException e) {
- throw new RuntimeException(e);
- }
+ protected void deleteInternal(byte[] key) throws RocksDBException {
+ db.delete(key);
}
@Override
- public DBIterator iterator() {
- return new RocksDBIterator(db.newIterator());
+ protected DBIterator newIterator(MetadataMetrics metrics) {
+ return new RocksDBIterator(db.newIterator(), metrics);
}
@Override
public void close() throws IOException {
- db.close();
+ try {
+ db.close();
+ } finally {
+ // WriteOptions is a native handle; release it even if db.close() throws.
+ SYNC_WRITE_OPTIONS.close();
+ }
}
}
diff --git
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBIterator.java
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBIterator.java
index 1d0001530..cce181dce 100644
---
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBIterator.java
+++
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/shuffledb/RocksDBIterator.java
@@ -32,6 +32,7 @@ import org.rocksdb.RocksIterator;
public class RocksDBIterator implements DBIterator {
private final RocksIterator it;
+ private final MetadataMetrics metrics;
private boolean checkedNext;
@@ -39,8 +40,9 @@ public class RocksDBIterator implements DBIterator {
private Map.Entry<byte[], byte[]> next;
- public RocksDBIterator(RocksIterator it) {
+ public RocksDBIterator(RocksIterator it, MetadataMetrics metrics) {
this.it = it;
+ this.metrics = metrics;
}
@Override
@@ -81,15 +83,28 @@ public class RocksDBIterator implements DBIterator {
@Override
public void seek(byte[] key) {
- it.seek(key);
+ metrics.onRead(
+ () -> {
+ it.seek(key);
+ it.status();
+ });
}
private Map.Entry<byte[], byte[]> loadNext() {
- if (it.isValid()) {
- Map.Entry<byte[], byte[]> nextEntry = new
AbstractMap.SimpleEntry<>(it.key(), it.value());
- it.next();
- return nextEntry;
+ if (!it.isValid()) {
+ metrics.onRead(
+ () -> {
+ it.status();
+ return null;
+ });
+ return null;
}
- return null;
+ return metrics.onRead(
+ () -> {
+ Map.Entry<byte[], byte[]> nextEntry = new
AbstractMap.SimpleEntry<>(it.key(), it.value());
+ it.next();
+ it.status();
+ return nextEntry;
+ });
}
}
diff --git
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java
index cdfb95c1e..5979a8434 100644
---
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java
+++
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java
@@ -124,7 +124,7 @@ public class PartitionFilesSorter extends
ShuffleRecoverHelper {
String recoverySortedFilesFileName =
dbBackend.fileName(RECOVERY_SORTED_FILES_FILE_NAME_PREFIX);
this.recoverFile = new File(recoverPath, recoverySortedFilesFileName);
- this.sortedFilesDb = DBProvider.initDB(dbBackend, recoverFile,
CURRENT_VERSION);
+ this.sortedFilesDb = DBProvider.initDB(dbBackend, recoverFile,
CURRENT_VERSION, source);
reloadAndCleanSortedShuffleFiles(this.sortedFilesDb);
} catch (Exception e) {
throw new IllegalStateException(
diff --git
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/WorkerSource.scala
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/WorkerSource.scala
index 15eb6a24e..f93858051 100644
---
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/WorkerSource.scala
+++
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/WorkerSource.scala
@@ -18,6 +18,7 @@
package org.apache.celeborn.service.deploy.worker
import java.util
+import java.util.Locale
import java.util.concurrent.ConcurrentHashMap
import scala.collection.JavaConverters._
@@ -74,6 +75,13 @@ class WorkerSource(conf: CelebornConf) extends
AbstractSource(conf, Role.WORKER)
addCounter(S3_FLUSH_COUNT)
addCounter(S3_FLUSH_SIZE)
+ // dbBackendLabel normalizes the case, so we can pass the configured value
as-is.
+ private val configuredDbBackend = conf.workerGracefulShutdownRecoverDbBackend
+ addCounter(METADATA_OPERATION_STATUS_COUNT,
WRITE_FAIL_COUNT_LABELS(configuredDbBackend))
+ addCounter(METADATA_OPERATION_STATUS_COUNT,
WRITE_SUCCESS_COUNT_LABELS(configuredDbBackend))
+ addCounter(METADATA_OPERATION_STATUS_COUNT,
READ_FAIL_COUNT_LABELS(configuredDbBackend))
+ addCounter(METADATA_OPERATION_STATUS_COUNT,
READ_SUCCESS_COUNT_LABELS(configuredDbBackend))
+
// add timers
addTimer(COMMIT_FILES_TIME)
addTimer(RESERVE_SLOTS_TIME)
@@ -290,4 +298,24 @@ object WorkerSource {
// clean
val CLEAN_TASK_QUEUE_SIZE = "CleanTaskQueueSize"
val CLEAN_EXPIRED_SHUFFLE_KEYS_TIME = "CleanExpiredShuffleKeysTime"
+
+ // metadata operations
+ val METADATA_OPERATION_STATUS_COUNT = "MetadataOperationStatusCount"
+ val WRITE_OPERATION_LABEL = Map("operation" -> "write")
+ val READ_OPERATION_LABEL = Map("operation" -> "read")
+
+ val SUCCESS_STATUS_LABEL = Map("status" -> "success")
+ val FAIL_STATUS_LABEL = Map("status" -> "fail")
+
+ def dbBackendLabel(dbBackend: String): Map[String, String] =
+ Map("dbBackend" -> dbBackend.toLowerCase(Locale.ROOT))
+
+ def WRITE_FAIL_COUNT_LABELS(dbBackend: String): Map[String, String] =
+ WRITE_OPERATION_LABEL ++ FAIL_STATUS_LABEL ++ dbBackendLabel(dbBackend)
+ def WRITE_SUCCESS_COUNT_LABELS(dbBackend: String): Map[String, String] =
+ WRITE_OPERATION_LABEL ++ SUCCESS_STATUS_LABEL ++ dbBackendLabel(dbBackend)
+ def READ_FAIL_COUNT_LABELS(dbBackend: String): Map[String, String] =
+ READ_OPERATION_LABEL ++ FAIL_STATUS_LABEL ++ dbBackendLabel(dbBackend)
+ def READ_SUCCESS_COUNT_LABELS(dbBackend: String): Map[String, String] =
+ READ_OPERATION_LABEL ++ SUCCESS_STATUS_LABEL ++ dbBackendLabel(dbBackend)
}
diff --git
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
index 6d375932b..25014470a 100644
---
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
+++
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
@@ -297,7 +297,7 @@ final private[worker] class StorageManager(conf:
CelebornConf, workerSource: Abs
val dbBackend =
DBBackend.byName(conf.workerGracefulShutdownRecoverDbBackend)
RECOVERY_FILE_NAME = dbBackend.fileName(RECOVERY_FILE_NAME_PREFIX)
val recoverFile = new File(conf.workerGracefulShutdownRecoverPath,
RECOVERY_FILE_NAME)
- this.db = DBProvider.initDB(dbBackend, recoverFile, CURRENT_VERSION)
+ this.db = DBProvider.initDB(dbBackend, recoverFile, CURRENT_VERSION,
workerSource)
reloadAndCleanFileInfos(this.db)
} catch {
case e: Exception =>
diff --git
a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/shuffledb/DBProviderSuiteJ.java
b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/shuffledb/DBProviderSuiteJ.java
index 69f0aa190..d5208e5d1 100644
---
a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/shuffledb/DBProviderSuiteJ.java
+++
b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/shuffledb/DBProviderSuiteJ.java
@@ -34,8 +34,10 @@ import org.junit.runners.Parameterized;
import org.junit.runners.Parameterized.Parameter;
import org.junit.runners.Parameterized.Parameters;
+import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.Utils;
+import org.apache.celeborn.service.deploy.worker.WorkerSource;
@RunWith(Parameterized.class)
public class DBProviderSuiteJ {
@@ -64,14 +66,17 @@ public class DBProviderSuiteJ {
createDirectory
? Utils.createDirectory(dbDir.getPath(), namePrefix)
: new File(dbDir.getPath(), String.format("%s-%s", namePrefix,
UUID.randomUUID()));
+ WorkerSource workerSource = new WorkerSource(new CelebornConf());
try {
StoreVersion v1 = new StoreVersion(1, 0);
- DBProvider.initDB(dbBackend, dbFile, v1).close();
+ DBProvider.initDB(dbBackend, dbFile, v1, workerSource).close();
StoreVersion v2 = new StoreVersion(2, 0);
IOException ioe =
- assertThrows(IOException.class, () -> DBProvider.initDB(dbBackend,
dbFile, v2));
+ assertThrows(
+ IOException.class, () -> DBProvider.initDB(dbBackend, dbFile,
v2, workerSource));
assertTrue(ioe.getMessage().contains("incompatible with current version
StoreVersion[2.0]"));
} finally {
+ workerSource.destroy();
JavaUtils.deleteRecursively(dbDir);
}
}
diff --git
a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/shuffledb/MetadataMetricsSuiteJ.java
b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/shuffledb/MetadataMetricsSuiteJ.java
new file mode 100644
index 000000000..ac48f219f
--- /dev/null
+++
b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/shuffledb/MetadataMetricsSuiteJ.java
@@ -0,0 +1,294 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.celeborn.service.deploy.worker.shuffledb;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.fail;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import java.util.concurrent.atomic.AtomicInteger;
+
+import scala.collection.JavaConverters;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.Mockito;
+import org.rocksdb.RocksDBException;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.metrics.source.NamedCounter;
+import org.apache.celeborn.service.deploy.worker.WorkerSource;
+
+public class MetadataMetricsSuiteJ {
+
+ private static final String METRIC =
WorkerSource.METADATA_OPERATION_STATUS_COUNT();
+ private static final String BACKEND = "rocksdb";
+
+ private WorkerSource mockSource;
+
+ @Before
+ public void setUp() {
+ mockSource = Mockito.mock(WorkerSource.class);
+ }
+
+ @Test
+ public void nullSourceConstructorDoesNotThrow() {
+ new MetadataMetrics(null, DBBackend.ROCKSDB);
+ }
+
+ @Test
+ public void nullSourceOnWriteStillExecutesAction() {
+ MetadataMetrics metrics = new MetadataMetrics(null, DBBackend.ROCKSDB);
+ AtomicInteger invocations = new AtomicInteger();
+
+ metrics.onWrite(invocations::incrementAndGet);
+
+ assertEquals(1, invocations.get());
+ }
+
+ @Test
+ public void nullSourceOnReadSupplierStillReturnsValue() {
+ MetadataMetrics metrics = new MetadataMetrics(null, DBBackend.ROCKSDB);
+
+ byte[] result = metrics.onRead((MetadataMetrics.ThrowingSupplier<byte[]>)
() -> new byte[] {7});
+
+ assertEquals(7, result[0]);
+ }
+
+ @Test
+ public void nullSourceFailureStillRethrows() {
+ MetadataMetrics metrics = new MetadataMetrics(null, DBBackend.ROCKSDB);
+ RocksDBException cause = new RocksDBException("test");
+
+ RuntimeException thrown =
+ assertThrows(
+ RuntimeException.class,
+ () ->
+ metrics.onWrite(
+ () -> {
+ throw cause;
+ }));
+
+ assertSame(cause, thrown.getCause());
+ }
+
+ @Test
+ public void constructorRegistersAllFourLabelCombinations() {
+ new MetadataMetrics(mockSource, DBBackend.ROCKSDB);
+
+ verify(mockSource).addCounter(METRIC,
WorkerSource.WRITE_SUCCESS_COUNT_LABELS(BACKEND));
+ verify(mockSource).addCounter(METRIC,
WorkerSource.WRITE_FAIL_COUNT_LABELS(BACKEND));
+ verify(mockSource).addCounter(METRIC,
WorkerSource.READ_SUCCESS_COUNT_LABELS(BACKEND));
+ verify(mockSource).addCounter(METRIC,
WorkerSource.READ_FAIL_COUNT_LABELS(BACKEND));
+ }
+
+ @Test
+ public void constructorUsesProvidedBackendForLabelValue() {
+ new MetadataMetrics(mockSource, DBBackend.LEVELDB);
+
+ verify(mockSource).addCounter(METRIC,
WorkerSource.WRITE_SUCCESS_COUNT_LABELS("leveldb"));
+ verify(mockSource).addCounter(METRIC,
WorkerSource.WRITE_FAIL_COUNT_LABELS("leveldb"));
+ verify(mockSource).addCounter(METRIC,
WorkerSource.READ_SUCCESS_COUNT_LABELS("leveldb"));
+ verify(mockSource).addCounter(METRIC,
WorkerSource.READ_FAIL_COUNT_LABELS("leveldb"));
+ }
+
+ @Test
+ public void onWriteRunnableSuccessIncrementsWriteSuccessCounter() {
+ MetadataMetrics metrics = new MetadataMetrics(mockSource,
DBBackend.ROCKSDB);
+ AtomicInteger invocations = new AtomicInteger();
+
+ metrics.onWrite(invocations::incrementAndGet);
+
+ assertEquals(1, invocations.get());
+ verify(mockSource).incCounter(METRIC, 1L,
WorkerSource.WRITE_SUCCESS_COUNT_LABELS(BACKEND));
+ verify(mockSource, never())
+ .incCounter(METRIC, 1L, WorkerSource.WRITE_FAIL_COUNT_LABELS(BACKEND));
+ }
+
+ @Test
+ public void onWriteRunnableRocksDBExceptionIsWrappedAndFailCounted() {
+ MetadataMetrics metrics = new MetadataMetrics(mockSource,
DBBackend.ROCKSDB);
+ RocksDBException cause = new RocksDBException("test");
+
+ RuntimeException thrown =
+ assertThrows(
+ RuntimeException.class,
+ () ->
+ metrics.onWrite(
+ () -> {
+ throw cause;
+ }));
+
+ assertSame(cause, thrown.getCause());
+ verify(mockSource).incCounter(METRIC, 1L,
WorkerSource.WRITE_FAIL_COUNT_LABELS(BACKEND));
+ verify(mockSource, never())
+ .incCounter(METRIC, 1L,
WorkerSource.WRITE_SUCCESS_COUNT_LABELS(BACKEND));
+ }
+
+ @Test
+ public void onWriteRunnableRuntimeExceptionPropagatesUnwrapped() {
+ MetadataMetrics metrics = new MetadataMetrics(mockSource,
DBBackend.ROCKSDB);
+ IllegalStateException original = new IllegalStateException("test");
+
+ IllegalStateException thrown =
+ assertThrows(
+ IllegalStateException.class,
+ () ->
+ metrics.onWrite(
+ () -> {
+ throw original;
+ }));
+
+ assertSame(original, thrown);
+ verify(mockSource).incCounter(METRIC, 1L,
WorkerSource.WRITE_FAIL_COUNT_LABELS(BACKEND));
+ }
+
+ @Test
+ public void onReadRunnableSuccessIncrementsReadSuccessCounter() {
+ MetadataMetrics metrics = new MetadataMetrics(mockSource,
DBBackend.ROCKSDB);
+
+ metrics.onRead((MetadataMetrics.ThrowingRunnable) () -> {});
+
+ verify(mockSource).incCounter(METRIC, 1L,
WorkerSource.READ_SUCCESS_COUNT_LABELS(BACKEND));
+ }
+
+ @Test
+ public void onReadRunnableFailureIncrementsReadFailCounter() {
+ MetadataMetrics metrics = new MetadataMetrics(mockSource,
DBBackend.ROCKSDB);
+ RocksDBException cause = new RocksDBException("test");
+
+ RuntimeException thrown =
+ assertThrows(
+ RuntimeException.class,
+ () ->
+ metrics.onRead(
+ (MetadataMetrics.ThrowingRunnable)
+ () -> {
+ throw cause;
+ }));
+
+ assertSame(cause, thrown.getCause());
+ verify(mockSource).incCounter(METRIC, 1L,
WorkerSource.READ_FAIL_COUNT_LABELS(BACKEND));
+ }
+
+ @Test
+ public void onReadSupplierReturnsValueOnSuccess() {
+ MetadataMetrics metrics = new MetadataMetrics(mockSource,
DBBackend.ROCKSDB);
+ byte[] expected = new byte[] {1, 2, 3};
+
+ byte[] result = metrics.onRead((MetadataMetrics.ThrowingSupplier<byte[]>)
() -> expected);
+
+ assertSame(expected, result);
+ verify(mockSource).incCounter(METRIC, 1L,
WorkerSource.READ_SUCCESS_COUNT_LABELS(BACKEND));
+ }
+
+ @Test
+ public void onReadSupplierFailureIncrementsReadFailCounter() {
+ MetadataMetrics metrics = new MetadataMetrics(mockSource,
DBBackend.ROCKSDB);
+ RocksDBException cause = new RocksDBException("test");
+
+ RuntimeException thrown =
+ assertThrows(
+ RuntimeException.class,
+ () ->
+ metrics.onRead(
+ (MetadataMetrics.ThrowingSupplier<Object>)
+ () -> {
+ throw cause;
+ }));
+
+ assertSame(cause, thrown.getCause());
+ verify(mockSource).incCounter(METRIC, 1L,
WorkerSource.READ_FAIL_COUNT_LABELS(BACKEND));
+ }
+
+ @Test
+ public void multipleOpsAccumulateCounterIncrements() {
+ MetadataMetrics metrics = new MetadataMetrics(mockSource,
DBBackend.ROCKSDB);
+
+ metrics.onWrite(() -> {});
+ metrics.onWrite(() -> {});
+ metrics.onRead((MetadataMetrics.ThrowingRunnable) () -> {});
+
+ verify(mockSource, times(2))
+ .incCounter(METRIC, 1L,
WorkerSource.WRITE_SUCCESS_COUNT_LABELS(BACKEND));
+ verify(mockSource, times(1))
+ .incCounter(METRIC, 1L,
WorkerSource.READ_SUCCESS_COUNT_LABELS(BACKEND));
+ }
+
+ /**
+ * End-to-end check using a real WorkerSource. This exercises the actual
counter registration +
+ * increment path without any mocks, and confirms the value-level effect on
the counter.
+ */
+ @Test
+ public void endToEndRealWorkerSourceTracksCounts() {
+ WorkerSource source = new WorkerSource(new CelebornConf());
+ try {
+ MetadataMetrics metrics = new MetadataMetrics(source, DBBackend.ROCKSDB);
+
+ metrics.onWrite(() -> {});
+ metrics.onWrite(() -> {});
+ metrics.onRead((MetadataMetrics.ThrowingRunnable) () -> {});
+ assertThrows(
+ RuntimeException.class,
+ () ->
+ metrics.onWrite(
+ () -> {
+ throw new RocksDBException("fail");
+ }));
+
+ assertEquals(2L, counterCount(source,
WorkerSource.WRITE_SUCCESS_COUNT_LABELS(BACKEND)));
+ assertEquals(1L, counterCount(source,
WorkerSource.WRITE_FAIL_COUNT_LABELS(BACKEND)));
+ assertEquals(1L, counterCount(source,
WorkerSource.READ_SUCCESS_COUNT_LABELS(BACKEND)));
+ assertEquals(0L, counterCount(source,
WorkerSource.READ_FAIL_COUNT_LABELS(BACKEND)));
+ } finally {
+ source.destroy();
+ }
+ }
+
+ private static long counterCount(
+ WorkerSource source, scala.collection.immutable.Map<String, String>
labels) {
+ String name = WorkerSource.METADATA_OPERATION_STATUS_COUNT();
+ java.util.Map<String, String> wanted =
JavaConverters.mapAsJavaMapConverter(labels).asJava();
+ for (NamedCounter c :
JavaConverters.seqAsJavaListConverter(source.counters()).asJava()) {
+ if (!c.name().equals(name)) {
+ continue;
+ }
+ java.util.Map<String, String> actual =
+ JavaConverters.mapAsJavaMapConverter(c.labels()).asJava();
+ if (containsAll(actual, wanted)) {
+ return c.counter().getCount();
+ }
+ }
+ fail("counter not found: " + name + " " + wanted);
+ return -1L;
+ }
+
+ private static boolean containsAll(
+ java.util.Map<String, String> actual, java.util.Map<String, String>
required) {
+ for (java.util.Map.Entry<String, String> e : required.entrySet()) {
+ if (!e.getValue().equals(actual.get(e.getKey()))) {
+ return false;
+ }
+ }
+ return true;
+ }
+}
