hiro/ray
1
0
Fork 0
mirror of https://github.com/vale981/ray synced 2025-03-05 10:01:43 -05:00
Code Issues Projects Releases Packages Wiki Activity Actions

[Java] Simplify Ray.wrapXxx() (#14895)

Browse source 
Co-authored-by: Qing Wang <jovany.wq@antgroup.com>
This commit is contained in:
Qing Wang 2021-03-25 10:58:12 +08:00 • committed by GitHub
parent 03afaed6e1
commit 493d15e05b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 30 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

58
java/runtime/src/main/java/io/ray/runtime/AbstractRayRuntime.java
View file

@ -219,26 +219,12 @@ public abstract class AbstractRayRuntime implements RayRuntimeInternal {
isContextSet.set(true);
}
// TODO (kfstorm): Simplify the duplicate code in wrap*** methods.
@Override
public final Runnable wrapRunnable(Runnable runnable) {
Object asyncContext = getAsyncContext();
return () -> {
boolean oldIsContextSet = isContextSet.get();
Object oldAsyncContext = null;
if (oldIsContextSet) {
oldAsyncContext = getAsyncContext();
}
setAsyncContext(asyncContext);
try {
try (RayAsyncContextUpdater updater = new RayAsyncContextUpdater(asyncContext, this)) {
runnable.run();
} finally {
if (oldIsContextSet) {
setAsyncContext(oldAsyncContext);
} else {
setIsContextSet(false);
}
}
};
}
@ -247,20 +233,8 @@ public abstract class AbstractRayRuntime implements RayRuntimeInternal {
public final <T> Callable<T> wrapCallable(Callable<T> callable) {
Object asyncContext = getAsyncContext();
return () -> {
boolean oldIsContextSet = isContextSet.get();
Object oldAsyncContext = null;
if (oldIsContextSet) {
oldAsyncContext = getAsyncContext();
}
setAsyncContext(asyncContext);
try {
try (RayAsyncContextUpdater updater = new RayAsyncContextUpdater(asyncContext, this)) {
return callable.call();
} finally {
if (oldIsContextSet) {
setAsyncContext(oldAsyncContext);
} else {
setIsContextSet(false);
}
}
};
}
@ -328,6 +302,34 @@ public abstract class AbstractRayRuntime implements RayRuntimeInternal {
return actor;
}
/// An auto closable class that is used for updating the async context when invoking Ray APIs.
private static final class RayAsyncContextUpdater implements AutoCloseable {
private AbstractRayRuntime runtime;
private boolean oldIsContextSet;
private Object oldAsyncContext = null;
public RayAsyncContextUpdater(Object asyncContext, AbstractRayRuntime runtime) {
this.runtime = runtime;
oldIsContextSet = runtime.isContextSet.get();
if (oldIsContextSet) {
oldAsyncContext = runtime.getAsyncContext();
}
runtime.setAsyncContext(asyncContext);
}
@Override
public void close() {
if (oldIsContextSet) {
runtime.setAsyncContext(oldAsyncContext);
} else {
runtime.setIsContextSet(false);
}
}
}
@Override
public WorkerContext getWorkerContext() {
return workerContext;