dev.restate.sdk.ContextImpl Maven / Gradle / Ivy
The newest version!
// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
//
// This file is part of the Restate Java SDK,
// which is released under the MIT license.
//
// You can find a copy of the license in file LICENSE in the root
// directory of this repository or package, or at
// https://github.com/restatedev/sdk-java/blob/main/LICENSE
package dev.restate.sdk;
import dev.restate.sdk.common.*;
import dev.restate.sdk.common.function.ThrowingSupplier;
import dev.restate.sdk.common.syscalls.Deferred;
import dev.restate.sdk.common.syscalls.EnterSideEffectSyscallCallback;
import dev.restate.sdk.common.syscalls.ExitSideEffectSyscallCallback;
import dev.restate.sdk.common.syscalls.Syscalls;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.util.Collection;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import org.jspecify.annotations.NonNull;
import org.jspecify.annotations.Nullable;
class ContextImpl implements ObjectContext, WorkflowContext {
final Syscalls syscalls;
ContextImpl(Syscalls syscalls) {
this.syscalls = syscalls;
}
@Override
public String key() {
return syscalls.objectKey();
}
@Override
public Request request() {
return syscalls.request();
}
@Override
public Optional get(StateKey key) {
Deferred deferred = Util.blockOnSyscall(cb -> syscalls.get(key.name(), cb));
if (!deferred.isCompleted()) {
Util.blockOnSyscall(cb -> syscalls.resolveDeferred(deferred, cb));
}
return Util.unwrapOptionalReadyResult(deferred.toResult())
.map(bs -> Util.deserializeWrappingException(syscalls, key.serde(), bs));
}
@Override
public Collection stateKeys() {
Deferred> deferred = Util.blockOnSyscall(syscalls::getKeys);
if (!deferred.isCompleted()) {
Util.blockOnSyscall(cb -> syscalls.resolveDeferred(deferred, cb));
}
return Util.unwrapResult(deferred.toResult());
}
@Override
public void clear(StateKey key) {
Util.blockOnSyscall(cb -> syscalls.clear(key.name(), cb));
}
@Override
public void clearAll() {
Util.blockOnSyscall(syscalls::clearAll);
}
@Override
public void set(StateKey key, @NonNull T value) {
Util.blockOnSyscall(
cb ->
syscalls.set(
key.name(), Util.serializeWrappingException(syscalls, key.serde(), value), cb));
}
@Override
public Awaitable timer(Duration duration) {
Deferred result = Util.blockOnSyscall(cb -> syscalls.sleep(duration, cb));
return Awaitable.single(syscalls, result);
}
@Override
public Awaitable call(
Target target, Serde inputSerde, Serde outputSerde, T parameter) {
ByteBuffer input = Util.serializeWrappingException(syscalls, inputSerde, parameter);
Deferred result = Util.blockOnSyscall(cb -> syscalls.call(target, input, cb));
return Awaitable.single(syscalls, result)
.map(bs -> Util.deserializeWrappingException(syscalls, outputSerde, bs));
}
@Override
public void send(Target target, Serde inputSerde, T parameter) {
ByteBuffer input = Util.serializeWrappingException(syscalls, inputSerde, parameter);
Util.blockOnSyscall(cb -> syscalls.send(target, input, null, cb));
}
@Override
public void send(Target target, Serde inputSerde, T parameter, Duration delay) {
ByteBuffer input = Util.serializeWrappingException(syscalls, inputSerde, parameter);
Util.blockOnSyscall(cb -> syscalls.send(target, input, delay, cb));
}
@Override
public T run(String name, Serde serde, ThrowingSupplier action) {
CompletableFuture> enterFut = new CompletableFuture<>();
syscalls.enterSideEffectBlock(
name,
new EnterSideEffectSyscallCallback() {
@Override
public void onNotExecuted() {
enterFut.complete(new CompletableFuture<>());
}
@Override
public void onSuccess(ByteBuffer result) {
enterFut.complete(CompletableFuture.completedFuture(result));
}
@Override
public void onFailure(TerminalException t) {
enterFut.complete(CompletableFuture.failedFuture(t));
}
@Override
public void onCancel(Throwable t) {
enterFut.cancel(true);
}
});
// If a failure was stored, it's simply thrown here
CompletableFuture exitFut = Util.awaitCompletableFuture(enterFut);
if (exitFut.isDone()) {
// We already have a result, we don't need to execute the action
return Util.deserializeWrappingException(
syscalls, serde, Util.awaitCompletableFuture(exitFut));
}
ExitSideEffectSyscallCallback exitCallback =
new ExitSideEffectSyscallCallback() {
@Override
public void onSuccess(ByteBuffer result) {
exitFut.complete(result);
}
@Override
public void onFailure(TerminalException t) {
exitFut.completeExceptionally(t);
}
@Override
public void onCancel(@Nullable Throwable t) {
exitFut.cancel(true);
}
};
T res = null;
Throwable failure = null;
try {
res = action.get();
} catch (Throwable e) {
failure = e;
}
if (failure != null) {
syscalls.exitSideEffectBlockWithException(failure, null, exitCallback);
} else {
syscalls.exitSideEffectBlock(
Util.serializeWrappingException(syscalls, serde, res), exitCallback);
}
return Util.deserializeWrappingException(syscalls, serde, Util.awaitCompletableFuture(exitFut));
}
@Override
public Awakeable awakeable(Serde serde) throws TerminalException {
// Retrieve the awakeable
Map.Entry> awakeable = Util.blockOnSyscall(syscalls::awakeable);
return new Awakeable<>(syscalls, awakeable.getValue(), serde, awakeable.getKey());
}
@Override
public AwakeableHandle awakeableHandle(String id) {
return new AwakeableHandle() {
@Override
public void resolve(Serde serde, @NonNull T payload) {
Util.blockOnSyscall(
cb ->
syscalls.resolveAwakeable(
id, Util.serializeWrappingException(syscalls, serde, payload), cb));
}
@Override
public void reject(String reason) {
Util.blockOnSyscall(cb -> syscalls.rejectAwakeable(id, reason, cb));
}
};
}
@Override
public RestateRandom random() {
return new RestateRandom(this.request().invocationId().toRandomSeed(), this.syscalls);
}
@Override
public DurablePromise promise(DurablePromiseKey key) {
return new DurablePromise<>() {
@Override
public Awaitable awaitable() {
Deferred result = Util.blockOnSyscall(cb -> syscalls.promise(key.name(), cb));
return Awaitable.single(syscalls, result)
.map(bs -> Util.deserializeWrappingException(syscalls, key.serde(), bs));
}
@Override
public Output peek() {
Deferred deferred =
Util.blockOnSyscall(cb -> syscalls.peekPromise(key.name(), cb));
if (!deferred.isCompleted()) {
Util.blockOnSyscall(cb -> syscalls.resolveDeferred(deferred, cb));
}
return Util.unwrapOutputReadyResult(deferred.toResult())
.map(bs -> Util.deserializeWrappingException(syscalls, key.serde(), bs));
}
};
}
@Override
public DurablePromiseHandle promiseHandle(DurablePromiseKey key) {
return new DurablePromiseHandle<>() {
@Override
public void resolve(T payload) throws IllegalStateException {
Deferred deferred =
Util.blockOnSyscall(
cb ->
syscalls.resolvePromise(
key.name(),
Util.serializeWrappingException(syscalls, key.serde(), payload),
cb));
if (!deferred.isCompleted()) {
Util.blockOnSyscall(cb -> syscalls.resolveDeferred(deferred, cb));
}
Util.unwrapResult(deferred.toResult());
}
@Override
public void reject(String reason) throws IllegalStateException {
Deferred deferred =
Util.blockOnSyscall(cb -> syscalls.rejectPromise(key.name(), reason, cb));
if (!deferred.isCompleted()) {
Util.blockOnSyscall(cb -> syscalls.resolveDeferred(deferred, cb));
}
Util.unwrapResult(deferred.toResult());
}
};
}
}