Skip to content

[AURON #2118] make native environment initialization retriable after failure #2151

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
XorSum wants to merge 2 commits into
base: master
Choose a base branch
from
Open

[AURON #2118] make native environment initialization retriable after failure #2151

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a32a036
avoid static init, retriable after failure
XorSum Apr 2, 2026
b78add6
add comment
XorSum Jun 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 37 additions & 21 deletions auron-core/src/main/java/org/apache/auron/jni/AuronCallNativeWrapper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,27 +53,7 @@ public class AuronCallNativeWrapper {
private Schema arrowSchema;
private long nativeRuntimePtr;
private Consumer<VectorSchemaRoot> batchConsumer;

// initialize native environment
static {
LOG.info("Initializing native environment (batchSize="
+ AuronAdaptor.getInstance().getAuronConfiguration().get(AuronConfiguration.BATCH_SIZE) + ", "
+ "memoryFraction="
+ AuronAdaptor.getInstance().getAuronConfiguration().get(AuronConfiguration.MEMORY_FRACTION) + ")");

// arrow configuration
System.setProperty("arrow.struct.conflict.policy", "CONFLICT_APPEND");

// preload JNI bridge classes
try {
Class.forName("org.apache.auron.jni.JniBridge");
} catch (ClassNotFoundException e) {
throw new RuntimeException("Cannot load JniBridge class", e);
}

AuronAdaptor.getInstance().loadAuronLib();
Runtime.getRuntime().addShutdownHook(new Thread(JniBridge::onExit));
}
private static volatile boolean initialized = false;

public AuronCallNativeWrapper(
BufferAllocator arrowAllocator,
Expand All @@ -90,13 +70,49 @@ public AuronCallNativeWrapper(
this.stageId = stageId;
this.taskId = taskId;

init();
LOG.warn("Start executing native plan");
this.nativeRuntimePtr = JniBridge.callNative(
nativeMemory,
AuronAdaptor.getInstance().getAuronConfiguration().get(AuronConfiguration.NATIVE_LOG_LEVEL),
this);
}

private static void init() {

@weiqingy weiqingy Jun 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The "How was this patch tested?" box is empty — and I think a faithful regression test is genuinely hard here, since it would need the real native lib plus a thread interrupted mid-Files.copy to reproduce the ClosedByInterruptException, which is racy and a process-global load. So this is a real question rather than a request: is there any lightweight check worth adding (even just asserting a second construction succeeds after init has run once), or is this effectively integration-only?

@XorSum XorSum Jun 20, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed it's effectively integration-only — any unit test would need to either expose init(), reflect-reset the static flag, or stub out the static native shutdown hook + callNative. Making the production code absorb these concessions for a mere DCL simply isn't worth the trade-off.

if (!initialized) {
synchronized (AuronCallNativeWrapper.class) {
if (!initialized) {
// Retry-safe: on failure `initialized` stays false so the next call retries.
// Every step below must be idempotent, or must not run on the failure path.
// In particular, loadAuronLib() must throw before System.load — a retried
// System.load would load the native library twice.
// initialize native environment
LOG.info("Initializing native environment (batchSize="
+ AuronAdaptor.getInstance().getAuronConfiguration().get(AuronConfiguration.BATCH_SIZE)
+ ", "
+ "memoryFraction="
+ AuronAdaptor.getInstance().getAuronConfiguration().get(AuronConfiguration.MEMORY_FRACTION)
Comment on lines +81 to +94

Copilot AI Apr 2, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

init() can run while a Spark task thread is already interrupted/cancelled (the production issue shows ClosedByInterruptException). Even though initialization is now retriable, attempting Files.copy/System.load from an interrupted task is still likely to fail and generate noisy errors. Consider short-circuiting initialization when the current task is not running (e.g., !AuronAdaptor.getInstance().isTaskRunning() or Thread.currentThread().isInterrupted()), and fail fast so a later non-cancelled task can perform the one-time init successfully.

Copilot uses AI. Check for mistakes.
+ ")");

// arrow configuration
System.setProperty("arrow.struct.conflict.policy", "CONFLICT_APPEND");

// preload JNI bridge classes
try {
Class.forName("org.apache.auron.jni.JniBridge");
} catch (ClassNotFoundException e) {
throw new RuntimeException("Cannot load JniBridge class", e);
}

AuronAdaptor.getInstance().loadAuronLib();

@weiqingy weiqingy Jun 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The retry depends on a quiet invariant here: every side effect that runs before initialized = true — System.setProperty, Class.forName, loadAuronLib, addShutdownHook — has to be either idempotent or guaranteed not to have run on the path that threw. Today that holds, since loadAuronLib() throws out of Files.copy before System.load, so on the failure path the native lib never loads and the shutdown hook never registers — a retry re-runs cleanly. That coupling lives in a different file, so an edit to loadAuronLib() that added work after System.load wouldn't have any signal at the retry site that it just broke the retry. Is it worth making that invariant visible here — a one-line note, or an idempotency check — so the next edit to the load path can't quietly break it?

@XorSum XorSum Jun 20, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! I have added a comment at the top of init() to point out the retry-safety contract.

Runtime.getRuntime().addShutdownHook(new Thread(JniBridge::onExit));

initialized = true;
Comment on lines +89 to +110

Copilot AI Apr 2, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Because initialization is now retried after failures, AuronAdaptor.getInstance().loadAuronLib() may be invoked multiple times in the same JVM. In the Spark adaptor, each call extracts to a fresh createTempFile(...).deleteOnExit(), so repeated failures can accumulate temp files until executor exit. Consider making loadAuronLib() idempotent / cached (or cleaning up temp files on failure) to avoid disk bloat on flaky/interrupt-driven init attempts.

Suggested change
// initialize native environment
LOG.info("Initializing native environment (batchSize="
+ AuronAdaptor.getInstance().getAuronConfiguration().get(AuronConfiguration.BATCH_SIZE)
+ ", "
+ "memoryFraction="
+ AuronAdaptor.getInstance().getAuronConfiguration().get(AuronConfiguration.MEMORY_FRACTION)
+ ")");
// arrow configuration
System.setProperty("arrow.struct.conflict.policy", "CONFLICT_APPEND");
// preload JNI bridge classes
try {
Class.forName("org.apache.auron.jni.JniBridge");
} catch (ClassNotFoundException e) {
throw new RuntimeException("Cannot load JniBridge class", e);
}
AuronAdaptor.getInstance().loadAuronLib();
Runtime.getRuntime().addShutdownHook(new Thread(JniBridge::onExit));
initialized = true;
try {
// initialize native environment
LOG.info("Initializing native environment (batchSize="
+ AuronAdaptor.getInstance().getAuronConfiguration().get(AuronConfiguration.BATCH_SIZE)
+ ", "
+ "memoryFraction="
+ AuronAdaptor.getInstance().getAuronConfiguration().get(AuronConfiguration.MEMORY_FRACTION)
+ ")");
// arrow configuration
System.setProperty("arrow.struct.conflict.policy", "CONFLICT_APPEND");
// preload JNI bridge classes
try {
Class.forName("org.apache.auron.jni.JniBridge");
} catch (ClassNotFoundException e) {
throw new RuntimeException("Cannot load JniBridge class", e);
}
AuronAdaptor.getInstance().loadAuronLib();
Runtime.getRuntime().addShutdownHook(new Thread(JniBridge::onExit));
} finally {
// Mark initialization as attempted to avoid repeated native library loading
initialized = true;
}

Copilot uses AI. Check for mistakes.
Comment on lines +81 to +110
}
}
}
}

/**
* Loads and processes the next batch of data from the native plan.
* <p>
Expand Down
Loading