Update tests and examples

restatedev · Feb 5, 2024 · 445ffc4 · 445ffc4
1 parent 77b3cbe
commit 445ffc4
Show file tree

Hide file tree

Showing 30 changed files with 109 additions and 107 deletions.
diff --git a/examples/src/main/java/dev/restate/sdk/examples/Counter.java b/examples/src/main/java/dev/restate/sdk/examples/Counter.java
@@ -8,7 +8,7 @@
 // https://github.com/restatedev/sdk-java/blob/main/LICENSE
 package dev.restate.sdk.examples;
 
-import dev.restate.sdk.RestateContext;
+import dev.restate.sdk.KeyedContext;
 import dev.restate.sdk.common.CoreSerdes;
 import dev.restate.sdk.common.StateKey;
 import dev.restate.sdk.examples.generated.*;
@@ -23,30 +23,28 @@ public class Counter extends CounterRestate.CounterRestateImplBase {
   private static final StateKey<Long> TOTAL = StateKey.of("total", CoreSerdes.JSON_LONG);
 
   @Override
-  public void reset(RestateContext ctx, CounterRequest request) {
-    restateContext().clear(TOTAL);
+  public void reset(KeyedContext ctx, CounterRequest request) {
+    ctx.clear(TOTAL);
   }
 
   @Override
-  public void add(RestateContext ctx, CounterAddRequest request) {
+  public void add(KeyedContext ctx, CounterAddRequest request) {
     long currentValue = ctx.get(TOTAL).orElse(0L);
     long newValue = currentValue + request.getValue();
     ctx.set(TOTAL, newValue);
   }
 
   @Override
-  public GetResponse get(RestateContext context, CounterRequest request) {
-    long currentValue = restateContext().get(TOTAL).orElse(0L);
+  public GetResponse get(KeyedContext ctx, CounterRequest request) {
+    long currentValue = ctx.get(TOTAL).orElse(0L);
 
     return GetResponse.newBuilder().setValue(currentValue).build();
   }
 
   @Override
-  public CounterUpdateResult getAndAdd(RestateContext context, CounterAddRequest request) {
+  public CounterUpdateResult getAndAdd(KeyedContext ctx, CounterAddRequest request) {
     LOG.info("Invoked get and add with " + request.getValue());
 
-    RestateContext ctx = restateContext();
-
     long currentValue = ctx.get(TOTAL).orElse(0L);
     long newValue = currentValue + request.getValue();
     ctx.set(TOTAL, newValue);

diff --git a/examples/src/main/java/dev/restate/sdk/examples/VanillaGrpcCounter.java b/examples/src/main/java/dev/restate/sdk/examples/VanillaGrpcCounter.java
@@ -9,7 +9,7 @@
 package dev.restate.sdk.examples;
 
 import com.google.protobuf.Empty;
-import dev.restate.sdk.RestateContext;
+import dev.restate.sdk.KeyedContext;
 import dev.restate.sdk.RestateService;
 import dev.restate.sdk.common.CoreSerdes;
 import dev.restate.sdk.common.StateKey;
@@ -27,15 +27,15 @@ public class VanillaGrpcCounter extends CounterGrpc.CounterImplBase implements R
 
   @Override
   public void reset(CounterRequest request, StreamObserver<Empty> responseObserver) {
-    restateContext().clear(TOTAL);
+    KeyedContext.current().clear(TOTAL);
 
     responseObserver.onNext(Empty.getDefaultInstance());
     responseObserver.onCompleted();
   }
 
   @Override
   public void add(CounterAddRequest request, StreamObserver<Empty> responseObserver) {
-    RestateContext ctx = restateContext();
+    KeyedContext ctx = KeyedContext.current();
 
     long currentValue = ctx.get(TOTAL).orElse(0L);
     long newValue = currentValue + request.getValue();
@@ -47,7 +47,7 @@ public void add(CounterAddRequest request, StreamObserver<Empty> responseObserve
 
   @Override
   public void get(CounterRequest request, StreamObserver<GetResponse> responseObserver) {
-    long currentValue = restateContext().get(TOTAL).orElse(0L);
+    long currentValue = KeyedContext.current().get(TOTAL).orElse(0L);
 
     responseObserver.onNext(GetResponse.newBuilder().setValue(currentValue).build());
     responseObserver.onCompleted();
@@ -58,7 +58,7 @@ public void getAndAdd(
       CounterAddRequest request, StreamObserver<CounterUpdateResult> responseObserver) {
     LOG.info("Invoked get and add with " + request.getValue());
 
-    RestateContext ctx = restateContext();
+    KeyedContext ctx = KeyedContext.current();
 
     long currentValue = ctx.get(TOTAL).orElse(0L);
     long newValue = currentValue + request.getValue();

diff --git a/examples/src/main/kotlin/dev/restate/sdk/examples/CounterKt.kt b/examples/src/main/kotlin/dev/restate/sdk/examples/CounterKt.kt
@@ -12,7 +12,7 @@ import dev.restate.sdk.common.CoreSerdes
 import dev.restate.sdk.common.StateKey
 import dev.restate.sdk.examples.generated.*
 import dev.restate.sdk.http.vertx.RestateHttpEndpointBuilder
-import dev.restate.sdk.kotlin.RestateContext
+import dev.restate.sdk.kotlin.KeyedContext
 import org.apache.logging.log4j.LogManager
 
 class CounterKt : CounterRestateKt.CounterRestateKtImplBase() {
@@ -21,20 +21,20 @@ class CounterKt : CounterRestateKt.CounterRestateKtImplBase() {
 
   private val TOTAL = StateKey.of("total", CoreSerdes.JSON_LONG)
 
-  override suspend fun reset(context: RestateContext, request: CounterRequest) {
+  override suspend fun reset(context: KeyedContext, request: CounterRequest) {
     context.clear(TOTAL)
   }
 
-  override suspend fun add(context: RestateContext, request: CounterAddRequest) {
+  override suspend fun add(context: KeyedContext, request: CounterAddRequest) {
     updateCounter(context, request.value)
   }
 
-  override suspend fun get(context: RestateContext, request: CounterRequest): GetResponse {
+  override suspend fun get(context: KeyedContext, request: CounterRequest): GetResponse {
     return getResponse { value = context.get(TOTAL) ?: 0L }
   }
 
   override suspend fun getAndAdd(
-      context: RestateContext,
+      context: KeyedContext,
       request: CounterAddRequest
   ): CounterUpdateResult {
     LOG.info("Invoked get and add with " + request.value)
@@ -45,7 +45,7 @@ class CounterKt : CounterRestateKt.CounterRestateKtImplBase() {
     }
   }
 
-  private suspend fun updateCounter(context: RestateContext, add: Long): Pair<Long, Long> {
+  private suspend fun updateCounter(context: KeyedContext, add: Long): Pair<Long, Long> {
     val currentValue = context.get(TOTAL) ?: 0L
     val newValue = currentValue + add
 

diff --git a/sdk-api-kotlin/src/test/kotlin/dev/restate/sdk/kotlin/AwaitableTest.kt b/sdk-api-kotlin/src/test/kotlin/dev/restate/sdk/kotlin/AwaitableTest.kt
@@ -21,7 +21,7 @@ class AwaitableTest : DeferredTestSuite() {
   private class ReverseAwaitOrder :
       GreeterGrpcKt.GreeterCoroutineImplBase(Dispatchers.Unconfined), RestateKtService {
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
-      val ctx = restateContext()
+      val ctx = KeyedContext.current()
       val a1 = ctx.callAsync(GreeterGrpcKt.greetMethod, greetingRequest { name = "Francesco" })
       val a2 = ctx.callAsync(GreeterGrpcKt.greetMethod, greetingRequest { name = "Till" })
       val a2Res = a2.await().getMessage()
@@ -38,7 +38,7 @@ class AwaitableTest : DeferredTestSuite() {
   private class AwaitTwiceTheSameAwaitable :
       GreeterGrpcKt.GreeterCoroutineImplBase(Dispatchers.Unconfined), RestateKtService {
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
-      val ctx = restateContext()
+      val ctx = KeyedContext.current()
       val a = ctx.callAsync(GreeterGrpcKt.greetMethod, greetingRequest { name = "Francesco" })
       return greetingResponse { message = a.await().getMessage() + "-" + a.await().getMessage() }
     }
@@ -51,7 +51,7 @@ class AwaitableTest : DeferredTestSuite() {
   private class AwaitAll :
       GreeterGrpcKt.GreeterCoroutineImplBase(Dispatchers.Unconfined), RestateKtService {
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
-      val ctx = restateContext()
+      val ctx = KeyedContext.current()
       val a1 = ctx.callAsync(GreeterGrpcKt.greetMethod, greetingRequest { name = "Francesco" })
       val a2 = ctx.callAsync(GreeterGrpcKt.greetMethod, greetingRequest { name = "Till" })
 
@@ -71,7 +71,7 @@ class AwaitableTest : DeferredTestSuite() {
   private class AwaitAny :
       GreeterGrpcKt.GreeterCoroutineImplBase(Dispatchers.Unconfined), RestateKtService {
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
-      val ctx = restateContext()
+      val ctx = KeyedContext.current()
       val a1 = ctx.callAsync(GreeterGrpcKt.greetMethod, greetingRequest { name = "Francesco" })
       val a2 = ctx.callAsync(GreeterGrpcKt.greetMethod, greetingRequest { name = "Till" })
       return Awaitable.any(a1, a2).await() as GreetingResponse
@@ -81,7 +81,7 @@ class AwaitableTest : DeferredTestSuite() {
   private class AwaitSelect :
       GreeterGrpcKt.GreeterCoroutineImplBase(Dispatchers.Unconfined), RestateKtService {
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
-      val ctx = restateContext()
+      val ctx = KeyedContext.current()
       val a1 = ctx.callAsync(GreeterGrpcKt.greetMethod, greetingRequest { name = "Francesco" })
       val a2 = ctx.callAsync(GreeterGrpcKt.greetMethod, greetingRequest { name = "Till" })
       return select {
@@ -98,7 +98,7 @@ class AwaitableTest : DeferredTestSuite() {
   private class CombineAnyWithAll :
       GreeterGrpcKt.GreeterCoroutineImplBase(Dispatchers.Unconfined), RestateKtService {
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
-      val ctx = restateContext()
+      val ctx = KeyedContext.current()
       val a1 = ctx.awakeable(CoreSerdes.JSON_STRING)
       val a2 = ctx.awakeable(CoreSerdes.JSON_STRING)
       val a3 = ctx.awakeable(CoreSerdes.JSON_STRING)
@@ -121,7 +121,7 @@ class AwaitableTest : DeferredTestSuite() {
   private class AwaitAnyIndex :
       GreeterGrpcKt.GreeterCoroutineImplBase(Dispatchers.Unconfined), RestateKtService {
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
-      val ctx = restateContext()
+      val ctx = KeyedContext.current()
       val a1 = ctx.awakeable(CoreSerdes.JSON_STRING)
       val a2 = ctx.awakeable(CoreSerdes.JSON_STRING)
       val a3 = ctx.awakeable(CoreSerdes.JSON_STRING)
@@ -140,7 +140,7 @@ class AwaitableTest : DeferredTestSuite() {
   private class AwaitOnAlreadyResolvedAwaitables :
       GreeterGrpcKt.GreeterCoroutineImplBase(Dispatchers.Unconfined), RestateKtService {
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
-      val ctx = restateContext()
+      val ctx = KeyedContext.current()
       val a1 = ctx.awakeable(CoreSerdes.JSON_STRING)
       val a2 = ctx.awakeable(CoreSerdes.JSON_STRING)
       val a12 = Awaitable.all(a1, a2)

diff --git a/sdk-api-kotlin/src/test/kotlin/dev/restate/sdk/kotlin/AwakeableIdTest.kt b/sdk-api-kotlin/src/test/kotlin/dev/restate/sdk/kotlin/AwakeableIdTest.kt
@@ -22,7 +22,7 @@ class AwakeableIdTest : AwakeableIdTestSuite() {
       GreeterGrpcKt.GreeterCoroutineImplBase(Dispatchers.Unconfined), RestateKtService {
 
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
-      val id: String = restateContext().awakeable(CoreSerdes.JSON_STRING).id
+      val id: String = KeyedContext.current().awakeable(CoreSerdes.JSON_STRING).id
       return greetingResponse { message = id }
     }
   }

diff --git a/sdk-api-kotlin/src/test/kotlin/dev/restate/sdk/kotlin/EagerStateTest.kt b/sdk-api-kotlin/src/test/kotlin/dev/restate/sdk/kotlin/EagerStateTest.kt
@@ -23,7 +23,7 @@ class EagerStateTest : EagerStateTestSuite() {
   private class GetEmpty :
       GreeterGrpcKt.GreeterCoroutineImplBase(Dispatchers.Unconfined), RestateKtService {
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
-      val ctx = restateContext()
+      val ctx = KeyedContext.current()
       val stateIsEmpty = ctx.get(StateKey.of("STATE", CoreSerdes.JSON_STRING)) == null
       return greetingResponse { message = stateIsEmpty.toString() }
     }
@@ -37,7 +37,7 @@ class EagerStateTest : EagerStateTestSuite() {
       GreeterGrpcKt.GreeterCoroutineImplBase(Dispatchers.Unconfined), RestateKtService {
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
       return greetingResponse {
-        message = restateContext().get(StateKey.of("STATE", CoreSerdes.JSON_STRING))!!
+        message = KeyedContext.current().get(StateKey.of("STATE", CoreSerdes.JSON_STRING))!!
       }
     }
   }
@@ -49,7 +49,7 @@ class EagerStateTest : EagerStateTestSuite() {
   private class GetAppendAndGet :
       GreeterGrpcKt.GreeterCoroutineImplBase(Dispatchers.Unconfined), RestateKtService {
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
-      val ctx = restateContext()
+      val ctx = KeyedContext.current()
       val oldState = ctx.get(StateKey.of("STATE", CoreSerdes.JSON_STRING))!!
       ctx.set(StateKey.of("STATE", CoreSerdes.JSON_STRING), oldState + request.getName())
       val newState = ctx.get(StateKey.of("STATE", CoreSerdes.JSON_STRING))!!
@@ -64,7 +64,7 @@ class EagerStateTest : EagerStateTestSuite() {
   private class GetClearAndGet :
       GreeterGrpcKt.GreeterCoroutineImplBase(Dispatchers.Unconfined), RestateKtService {
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
-      val ctx = restateContext()
+      val ctx = KeyedContext.current()
       val oldState = ctx.get(StateKey.of("STATE", CoreSerdes.JSON_STRING))!!
       ctx.clear(StateKey.of("STATE", CoreSerdes.JSON_STRING))
       assertThat(ctx.get(StateKey.of("STATE", CoreSerdes.JSON_STRING))).isNull()

diff --git a/sdk-api-kotlin/src/test/kotlin/dev/restate/sdk/kotlin/RandomTest.kt b/sdk-api-kotlin/src/test/kotlin/dev/restate/sdk/kotlin/RandomTest.kt
@@ -22,7 +22,7 @@ class RandomTest : RandomTestSuite() {
       GreeterGrpcKt.GreeterCoroutineImplBase(Dispatchers.Unconfined), RestateKtService {
 
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
-      val number = restateContext().random().nextInt()
+      val number = KeyedContext.current().random().nextInt()
       return greetingResponse { message = number.toString() }
     }
   }
@@ -34,7 +34,7 @@ class RandomTest : RandomTestSuite() {
   private class RandomInsideSideEffect :
       GreeterGrpcKt.GreeterCoroutineImplBase(Dispatchers.Unconfined), RestateKtService {
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
-      val ctx = restateContext()
+      val ctx = KeyedContext.current()
       ctx.sideEffect { ctx.random().nextInt() }
       throw IllegalStateException("This should not unreachable")
     }

diff --git a/sdk-api-kotlin/src/test/kotlin/dev/restate/sdk/kotlin/RestateCodegenTest.kt b/sdk-api-kotlin/src/test/kotlin/dev/restate/sdk/kotlin/RestateCodegenTest.kt
@@ -17,10 +17,7 @@ import kotlin.time.Duration.Companion.seconds
 
 class RestateCodegenTest : RestateCodegenTestSuite() {
   private class GreeterWithRestateClientAndServerCodegen : GreeterRestateKtImplBase() {
-    override suspend fun greet(
-        context: RestateContext,
-        request: GreetingRequest
-    ): GreetingResponse {
+    override suspend fun greet(context: KeyedContext, request: GreetingRequest): GreetingResponse {
       val client = GreeterRestateKt.newClient(context)
       client.delayed(1.seconds).greet(request)
       client.oneWay().greet(request)
@@ -33,27 +30,27 @@ class RestateCodegenTest : RestateCodegenTestSuite() {
   }
 
   private class Codegen : CodegenRestateKtImplBase() {
-    override suspend fun emptyInput(context: RestateContext): MyMessage {
+    override suspend fun emptyInput(context: UnkeyedContext): MyMessage {
       val client = CodegenRestateKt.newClient(context)
       return client.emptyInput().await()
     }
 
-    override suspend fun emptyOutput(context: RestateContext, request: MyMessage) {
+    override suspend fun emptyOutput(context: UnkeyedContext, request: MyMessage) {
       val client = CodegenRestateKt.newClient(context)
       client.emptyOutput(request).await()
     }
 
-    override suspend fun emptyInputOutput(context: RestateContext) {
+    override suspend fun emptyInputOutput(context: UnkeyedContext) {
       val client = CodegenRestateKt.newClient(context)
       client.emptyInputOutput().await()
     }
 
-    override suspend fun oneWay(context: RestateContext, request: MyMessage): MyMessage {
+    override suspend fun oneWay(context: UnkeyedContext, request: MyMessage): MyMessage {
       val client = CodegenRestateKt.newClient(context)
       return client._oneWay(request).await()
     }
 
-    override suspend fun delayed(context: RestateContext, request: MyMessage): MyMessage {
+    override suspend fun delayed(context: UnkeyedContext, request: MyMessage): MyMessage {
       val client = CodegenRestateKt.newClient(context)
       return client._delayed(request).await()
     }

diff --git a/sdk-api-kotlin/src/test/kotlin/dev/restate/sdk/kotlin/SideEffectTest.kt b/sdk-api-kotlin/src/test/kotlin/dev/restate/sdk/kotlin/SideEffectTest.kt
@@ -20,7 +20,7 @@ class SideEffectTest : SideEffectTestSuite() {
   private class SideEffect(private val sideEffectOutput: String) :
       GreeterGrpcKt.GreeterCoroutineImplBase(Dispatchers.Unconfined), RestateKtService {
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
-      val ctx: RestateContext = restateContext()
+      val ctx = KeyedContext.current()
       val result = ctx.sideEffect(CoreSerdes.JSON_STRING) { sideEffectOutput }
       return greetingResponse { message = "Hello $result" }
     }
@@ -33,7 +33,7 @@ class SideEffectTest : SideEffectTestSuite() {
   private class ConsecutiveSideEffect(private val sideEffectOutput: String) :
       GreeterGrpcKt.GreeterCoroutineImplBase(Dispatchers.Unconfined), RestateKtService {
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
-      val ctx: RestateContext = restateContext()
+      val ctx = KeyedContext.current()
       val firstResult = ctx.sideEffect(CoreSerdes.JSON_STRING) { sideEffectOutput }
       val secondResult =
           ctx.sideEffect(CoreSerdes.JSON_STRING) { firstResult.uppercase(Locale.getDefault()) }
@@ -52,7 +52,7 @@ class SideEffectTest : SideEffectTestSuite() {
 
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
       val sideEffectThread =
-          restateContext().sideEffect(CoreSerdes.JSON_STRING) { Thread.currentThread().name }
+          KeyedContext.current().sideEffect(CoreSerdes.JSON_STRING) { Thread.currentThread().name }
       check(sideEffectThread.contains("CheckContextSwitchingTestCoroutine")) {
         "Side effect thread is not running within the same coroutine context of the handler method: $sideEffectThread"
       }
@@ -67,7 +67,7 @@ class SideEffectTest : SideEffectTestSuite() {
   private class SideEffectGuard :
       GreeterGrpcKt.GreeterCoroutineImplBase(Dispatchers.Unconfined), RestateKtService {
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
-      val ctx = restateContext()
+      val ctx = KeyedContext.current()
       ctx.sideEffect {
         ctx.oneWayCall(GreeterGrpcKt.greetMethod, greetingRequest { name = "something" })
       }

diff --git a/sdk-api-kotlin/src/test/kotlin/dev/restate/sdk/kotlin/SleepTest.kt b/sdk-api-kotlin/src/test/kotlin/dev/restate/sdk/kotlin/SleepTest.kt
@@ -21,7 +21,7 @@ class SleepTest : SleepTestSuite() {
   private class SleepGreeter :
       GreeterGrpcKt.GreeterCoroutineImplBase(Dispatchers.Unconfined), RestateKtService {
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
-      val ctx = restateContext()
+      val ctx = KeyedContext.current()
       ctx.sleep(1000.milliseconds)
       return greetingResponse { message = "Hello" }
     }
@@ -34,7 +34,7 @@ class SleepTest : SleepTestSuite() {
   private class ManySleeps :
       GreeterGrpcKt.GreeterCoroutineImplBase(Dispatchers.Unconfined), RestateKtService {
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
-      val ctx = restateContext()
+      val ctx = KeyedContext.current()
       val awaitables = mutableListOf<Awaitable<Unit>>()
       for (i in 0..9) {
         awaitables.add(ctx.timer(1000.milliseconds))

diff --git a/sdk-api-kotlin/src/test/kotlin/dev/restate/sdk/kotlin/StateMachineFailuresTest.kt b/sdk-api-kotlin/src/test/kotlin/dev/restate/sdk/kotlin/StateMachineFailuresTest.kt
@@ -27,7 +27,7 @@ class StateMachineFailuresTest : StateMachineFailuresTestSuite() {
       GreeterGrpcKt.GreeterCoroutineImplBase(Dispatchers.Unconfined), RestateKtService {
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
       try {
-        restateContext().get(STATE)
+        KeyedContext.current().get(STATE)
       } catch (e: Throwable) {
         // A user should never catch Throwable!!!
         if (e !is CancellationException && e !is TerminalException) {
@@ -57,7 +57,7 @@ class StateMachineFailuresTest : StateMachineFailuresTestSuite() {
   private class SideEffectFailure(private val serde: Serde<Int>) :
       GreeterGrpcKt.GreeterCoroutineImplBase(Dispatchers.Unconfined), RestateKtService {
     override suspend fun greet(request: GreetingRequest): GreetingResponse {
-      restateContext().sideEffect(serde) { 0 }
+      KeyedContext.current().sideEffect(serde) { 0 }
       return greetingResponse { message = "Francesco" }
     }
   }