Visitar URL original
feat(metrics): add SessionContext.memoryUsage and runtimeStats (#85) · apache/datafusion-java@c291b81 · GitHub
Skip to content

Commit c291b81

Browse files
LantaoJinLantaoJin
authored
feat(metrics): add SessionContext.memoryUsage and runtimeStats (#85)
1 parent dc9b700 commit c291b81

11 files changed

Lines changed: 1094 additions & 8 deletions

File tree

‎Makefile‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,20 @@
1515
# specific language governing permissions and limitations
1616
# under the License.
1717

18-
.PHONY: all native jvm test clean tpch-data
18+
.PHONY: all native native-runtime-metrics jvm test clean tpch-data
1919

2020
all: native jvm
2121

2222
native:
2323
cd native && cargo build
2424

25+
# Build the native crate with the `runtime-metrics` Cargo feature enabled.
26+
# Requires `--cfg tokio_unstable` because tokio-metrics gates its API there.
27+
# Default `make native` does not pull this in; callers who need
28+
# SessionContext.runtimeStats() pick this target explicitly.
29+
native-runtime-metrics:
30+
cd native && RUSTFLAGS="--cfg tokio_unstable" cargo build --features runtime-metrics
31+
2532
jvm:
2633
./mvnw package -DskipTests
2734

‎core/src/main/java/org/apache/datafusion/MemoryUsage.java‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.datafusion;
21+
22+
/**
23+
* Snapshot of session-wide memory usage from {@link SessionContext#memoryUsage()}. Returns the
24+
* total bytes currently reserved against the session's {@code MemoryPool} and the peak observed
25+
* since the session was created.
26+
*
27+
* <p>Multi-tenant attribution: place each tenant in its own {@link SessionContext}. Sessions in
28+
* DataFusion are cheap; one per tenant is the standard pattern. Within a single session, the
29+
* snapshot is the sum across all in-flight queries' operator reservations -- see {@link
30+
* SessionContext#memoryUsage()} for the precise definition of what is and is not counted.
31+
*
32+
* <p>Instances are immutable.
33+
*/
34+
public final class MemoryUsage {
35+
private final long currentBytes;
36+
private final long peakBytes;
37+
38+
public MemoryUsage(long currentBytes, long peakBytes) {
39+
this.currentBytes = currentBytes;
40+
this.peakBytes = peakBytes;
41+
}
42+
43+
/** Bytes currently reserved against this session's {@code MemoryPool}. */
44+
public long currentBytes() {
45+
return currentBytes;
46+
}
47+
48+
/**
49+
* Maximum value of {@link #currentBytes()} observed since the session was created. Monotonic:
50+
* never decreases for a given session.
51+
*/
52+
public long peakBytes() {
53+
return peakBytes;
54+
}
55+
56+
@Override
57+
public String toString() {
58+
return "MemoryUsage{currentBytes=" + currentBytes + ", peakBytes=" + peakBytes + "}";
59+
}
60+
}

‎core/src/main/java/org/apache/datafusion/RuntimeStats.java‎

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.datafusion;
21+
22+
/**
23+
* Snapshot of the underlying Tokio runtime from {@link SessionContext#runtimeStats()}. Mirrors the
24+
* subset of {@code tokio_metrics::RuntimeMetrics} fields that node-level dashboards typically
25+
* surface (worker count, busy time, queue depth, etc.).
26+
*
27+
* <p>Instances are immutable. The snapshot is process-wide rather than per-session because the JNI
28+
* library drives a single shared multi-threaded Tokio runtime; gating on the {@code SessionContext}
29+
* only ensures the caller still has a live handle.
30+
*
31+
* <p>Requires the {@code runtime-metrics} Cargo feature on the native crate, plus {@code
32+
* RUSTFLAGS="--cfg tokio_unstable"} at build time.
33+
*/
34+
public final class RuntimeStats {
35+
private final int numWorkers;
36+
private final long liveTasksCount;
37+
private final long globalQueueDepth;
38+
private final long elapsedNanos;
39+
private final long totalBusyNanos;
40+
private final long totalParkCount;
41+
private final long totalPollsCount;
42+
private final long totalNoopCount;
43+
private final long totalStealCount;
44+
private final long totalLocalScheduleCount;
45+
private final long totalOverflowCount;
46+
47+
public RuntimeStats(
48+
int numWorkers,
49+
long liveTasksCount,
50+
long globalQueueDepth,
51+
long elapsedNanos,
52+
long totalBusyNanos,
53+
long totalParkCount,
54+
long totalPollsCount,
55+
long totalNoopCount,
56+
long totalStealCount,
57+
long totalLocalScheduleCount,
58+
long totalOverflowCount) {
59+
this.numWorkers = numWorkers;
60+
this.liveTasksCount = liveTasksCount;
61+
this.globalQueueDepth = globalQueueDepth;
62+
this.elapsedNanos = elapsedNanos;
63+
this.totalBusyNanos = totalBusyNanos;
64+
this.totalParkCount = totalParkCount;
65+
this.totalPollsCount = totalPollsCount;
66+
this.totalNoopCount = totalNoopCount;
67+
this.totalStealCount = totalStealCount;
68+
this.totalLocalScheduleCount = totalLocalScheduleCount;
69+
this.totalOverflowCount = totalOverflowCount;
70+
}
71+
72+
/** Number of OS worker threads driving the multi-threaded Tokio runtime. */
73+
public int numWorkers() {
74+
return numWorkers;
75+
}
76+
77+
/** Number of tasks currently scheduled (alive) on the runtime. */
78+
public long liveTasksCount() {
79+
return liveTasksCount;
80+
}
81+
82+
/** Tasks waiting in the global injection queue, not yet picked up by a worker. */
83+
public long globalQueueDepth() {
84+
return globalQueueDepth;
85+
}
86+
87+
/** Wall-clock nanoseconds covered by this snapshot's interval (since runtime start). */
88+
public long elapsedNanos() {
89+
return elapsedNanos;
90+
}
91+
92+
/** Total nanoseconds workers spent doing work (sum across workers). */
93+
public long totalBusyNanos() {
94+
return totalBusyNanos;
95+
}
96+
97+
/** Times a worker has parked itself (gone idle waiting for work). */
98+
public long totalParkCount() {
99+
return totalParkCount;
100+
}
101+
102+
/** Total task polls completed across workers. */
103+
public long totalPollsCount() {
104+
return totalPollsCount;
105+
}
106+
107+
/** Times a worker unparked but found no work (false wakeup). */
108+
public long totalNoopCount() {
109+
return totalNoopCount;
110+
}
111+
112+
/** Times a worker successfully stole tasks from another worker. */
113+
public long totalStealCount() {
114+
return totalStealCount;
115+
}
116+
117+
/** Tasks scheduled into a worker's local queue. */
118+
public long totalLocalScheduleCount() {
119+
return totalLocalScheduleCount;
120+
}
121+
122+
/** Times a worker's local queue overflowed and pushed tasks to the injector. */
123+
public long totalOverflowCount() {
124+
return totalOverflowCount;
125+
}
126+
127+
@Override
128+
public String toString() {
129+
return "RuntimeStats{numWorkers="
130+
+ numWorkers
131+
+ ", liveTasksCount="
132+
+ liveTasksCount
133+
+ ", globalQueueDepth="
134+
+ globalQueueDepth
135+
+ ", elapsedNanos="
136+
+ elapsedNanos
137+
+ ", totalBusyNanos="
138+
+ totalBusyNanos
139+
+ ", totalParkCount="
140+
+ totalParkCount
141+
+ ", totalPollsCount="
142+
+ totalPollsCount
143+
+ ", totalNoopCount="
144+
+ totalNoopCount
145+
+ ", totalStealCount="
146+
+ totalStealCount
147+
+ ", totalLocalScheduleCount="
148+
+ totalLocalScheduleCount
149+
+ ", totalOverflowCount="
150+
+ totalOverflowCount
151+
+ "}";
152+
}
153+
}

‎core/src/main/java/org/apache/datafusion/SessionContext.java‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,73 @@ public DataFrame fromProto(byte[] planBytes) {
9999
return new DataFrame(dfHandle);
100100
}
101101

102+
/**
103+
* Snapshot the session's memory pool: bytes currently held and the peak observed since this
104+
* session was created. Thread-safe; can be polled while queries run.
105+
*
106+
* <p>For multi-tenant attribution, place each tenant in its own {@link SessionContext}. Within a
107+
* single session, the snapshot is the sum across all in-flight queries -- there is no
108+
* per-DataFrame breakdown today.
109+
*
110+
* <p>The session's {@code MemoryPool} is wrapped transparently with a tracking adapter at
111+
* construction time; the wrapper layers on top of whatever pool {@link
112+
* SessionContextBuilder#memoryLimit(long, double)} produced (or DataFusion's default unbounded
113+
* pool) and does not change pool semantics (limits, eviction, spilling).
114+
*
115+
* <p><b>What this counts:</b> bytes reserved against the {@code MemoryPool} -- operator state for
116+
* sorts, hash joins, aggregates, repartition buffers, and anything else that uses DataFusion's
117+
* {@code MemoryReservation} machinery during execution.
118+
*
119+
* <p><b>What this does <i>not</i> count:</b> memory held outside the pool, including record-batch
120+
* buffers materialised by {@link DataFrame#cache()} (stored in an in-memory {@code MemTable} as
121+
* plain {@code Vec<RecordBatch>} with no reservation), record-batch buffers that have crossed the
122+
* FFI boundary into Arrow's Java allocator, and JVM-side allocations. Operator-level reservations
123+
* are released as the plan unwinds, so a query that runs to completion typically returns {@code
124+
* currentBytes} to ~0 even if the result set is large.
125+
*
126+
* @throws IllegalStateException if this context is closed.
127+
* @throws RuntimeException if the native side has not registered a tracker for this handle
128+
* (should not happen in practice -- tracker registration is done by the constructor).
129+
*/
130+
public MemoryUsage memoryUsage() {
131+
if (nativeHandle == 0) {
132+
throw new IllegalStateException("SessionContext is closed");
133+
}
134+
long[] values = memoryUsageNative(nativeHandle);
135+
return new MemoryUsage(values[0], values[1]);
136+
}
137+
138+
/**
139+
* Snapshot operational counters from the underlying Tokio runtime: worker count, busy time, queue
140+
* depth, etc. Thread-safe; can be polled while queries run.
141+
*
142+
* <p>The runtime is process-wide rather than per-session because the JNI library drives a single
143+
* shared multi-threaded Tokio runtime. The {@link SessionContext} handle is checked only to
144+
* ensure the caller still has a live session; the values returned are not session-specific.
145+
*
146+
* <p>Requires the {@code runtime-metrics} Cargo feature on the native crate (off by default).
147+
* Rebuild with:
148+
*
149+
* <pre>{@code
150+
* RUSTFLAGS="--cfg tokio_unstable" cargo build --features runtime-metrics
151+
* }</pre>
152+
*
153+
* <p>If invoked against a native binary built without the feature, this method throws {@link
154+
* RuntimeException} with a message pointing at the rebuild command.
155+
*
156+
* @throws IllegalStateException if this context is closed.
157+
* @throws RuntimeException if the native crate was built without the {@code runtime-metrics}
158+
* feature.
159+
*/
160+
public RuntimeStats runtimeStats() {
161+
if (nativeHandle == 0) {
162+
throw new IllegalStateException("SessionContext is closed");
163+
}
164+
long[] s = runtimeStatsNative(nativeHandle);
165+
return new RuntimeStats(
166+
(int) s[0], s[1], s[2], s[3], s[4], s[5], s[6], s[7], s[8], s[9], s[10]);
167+
}
168+
102169
/**
103170
* Return the Arrow {@link Schema} of a registered table. Transferred via Arrow IPC; no {@link
104171
* org.apache.arrow.memory.BufferAllocator} is required because a schema carries no buffer data.
@@ -518,6 +585,10 @@ public void close() {
518585

519586
private static native String getOptionNative(long handle, String key);
520587

588+
private static native long[] memoryUsageNative(long handle);
589+
590+
private static native long[] runtimeStatsNative(long handle);
591+
521592
private static native void registerParquetWithOptions(
522593
long handle, String name, String path, byte[] optionsBytes, byte[] schemaIpcBytes);
523594

0 commit comments

Comments
 (0)