Skip to content

Commit 64b4d7a

Browse files
slinkydeveloperslinkydeveloper
authored
Split RestateContext interface in KeyedContext/UnkeyedContext (#213)
* [java] Split RestateContext interface in KeyedContext/UnkeyedContext * [kotlin] Split RestateContext interface in KeyedContext/UnkeyedContext * Make sure the protoc code generator injects the correct interface * Update tests and examples
1 parent 008868c commit 64b4d7a

File tree

42 files changed

+329
-271
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

42 files changed

+329
-271
lines changed

‎examples/src/main/java/dev/restate/sdk/examples/Counter.java

Lines changed: 7 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
// https://github.com/restatedev/sdk-java/blob/main/LICENSE
99
package dev.restate.sdk.examples;
1010

11-
import dev.restate.sdk.RestateContext;
11+
import dev.restate.sdk.KeyedContext;
1212
import dev.restate.sdk.common.CoreSerdes;
1313
import dev.restate.sdk.common.StateKey;
1414
import dev.restate.sdk.examples.generated.*;
@@ -23,30 +23,28 @@ public class Counter extends CounterRestate.CounterRestateImplBase {
2323
private static final StateKey<Long> TOTAL = StateKey.of("total", CoreSerdes.JSON_LONG);
2424

2525
@Override
26-
public void reset(RestateContext ctx, CounterRequest request) {
27-
restateContext().clear(TOTAL);
26+
public void reset(KeyedContext ctx, CounterRequest request) {
27+
ctx.clear(TOTAL);
2828
}
2929

3030
@Override
31-
public void add(RestateContext ctx, CounterAddRequest request) {
31+
public void add(KeyedContext ctx, CounterAddRequest request) {
3232
long currentValue = ctx.get(TOTAL).orElse(0L);
3333
long newValue = currentValue + request.getValue();
3434
ctx.set(TOTAL, newValue);
3535
}
3636

3737
@Override
38-
public GetResponse get(RestateContext context, CounterRequest request) {
39-
long currentValue = restateContext().get(TOTAL).orElse(0L);
38+
public GetResponse get(KeyedContext ctx, CounterRequest request) {
39+
long currentValue = ctx.get(TOTAL).orElse(0L);
4040

4141
return GetResponse.newBuilder().setValue(currentValue).build();
4242
}
4343

4444
@Override
45-
public CounterUpdateResult getAndAdd(RestateContext context, CounterAddRequest request) {
45+
public CounterUpdateResult getAndAdd(KeyedContext ctx, CounterAddRequest request) {
4646
LOG.info("Invoked get and add with " + request.getValue());
4747

48-
RestateContext ctx = restateContext();
49-
5048
long currentValue = ctx.get(TOTAL).orElse(0L);
5149
long newValue = currentValue + request.getValue();
5250
ctx.set(TOTAL, newValue);

‎examples/src/main/java/dev/restate/sdk/examples/VanillaGrpcCounter.java

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
package dev.restate.sdk.examples;
1010

1111
import com.google.protobuf.Empty;
12-
import dev.restate.sdk.RestateContext;
12+
import dev.restate.sdk.KeyedContext;
1313
import dev.restate.sdk.RestateService;
1414
import dev.restate.sdk.common.CoreSerdes;
1515
import dev.restate.sdk.common.StateKey;
@@ -27,15 +27,15 @@ public class VanillaGrpcCounter extends CounterGrpc.CounterImplBase implements R
2727

2828
@Override
2929
public void reset(CounterRequest request, StreamObserver<Empty> responseObserver) {
30-
restateContext().clear(TOTAL);
30+
KeyedContext.current().clear(TOTAL);
3131

3232
responseObserver.onNext(Empty.getDefaultInstance());
3333
responseObserver.onCompleted();
3434
}
3535

3636
@Override
3737
public void add(CounterAddRequest request, StreamObserver<Empty> responseObserver) {
38-
RestateContext ctx = restateContext();
38+
KeyedContext ctx = KeyedContext.current();
3939

4040
long currentValue = ctx.get(TOTAL).orElse(0L);
4141
long newValue = currentValue + request.getValue();
@@ -47,7 +47,7 @@ public void add(CounterAddRequest request, StreamObserver<Empty> responseObserve
4747

4848
@Override
4949
public void get(CounterRequest request, StreamObserver<GetResponse> responseObserver) {
50-
long currentValue = restateContext().get(TOTAL).orElse(0L);
50+
long currentValue = KeyedContext.current().get(TOTAL).orElse(0L);
5151

5252
responseObserver.onNext(GetResponse.newBuilder().setValue(currentValue).build());
5353
responseObserver.onCompleted();
@@ -58,7 +58,7 @@ public void getAndAdd(
5858
CounterAddRequest request, StreamObserver<CounterUpdateResult> responseObserver) {
5959
LOG.info("Invoked get and add with " + request.getValue());
6060

61-
RestateContext ctx = restateContext();
61+
KeyedContext ctx = KeyedContext.current();
6262

6363
long currentValue = ctx.get(TOTAL).orElse(0L);
6464
long newValue = currentValue + request.getValue();

‎examples/src/main/kotlin/dev/restate/sdk/examples/CounterKt.kt

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ import dev.restate.sdk.common.CoreSerdes
1212
import dev.restate.sdk.common.StateKey
1313
import dev.restate.sdk.examples.generated.*
1414
import dev.restate.sdk.http.vertx.RestateHttpEndpointBuilder
15-
import dev.restate.sdk.kotlin.RestateContext
15+
import dev.restate.sdk.kotlin.KeyedContext
1616
import org.apache.logging.log4j.LogManager
1717

1818
class CounterKt : CounterRestateKt.CounterRestateKtImplBase() {
@@ -21,20 +21,20 @@ class CounterKt : CounterRestateKt.CounterRestateKtImplBase() {
2121

2222
private val TOTAL = StateKey.of("total", CoreSerdes.JSON_LONG)
2323

24-
override suspend fun reset(context: RestateContext, request: CounterRequest) {
24+
override suspend fun reset(context: KeyedContext, request: CounterRequest) {
2525
context.clear(TOTAL)
2626
}
2727

28-
override suspend fun add(context: RestateContext, request: CounterAddRequest) {
28+
override suspend fun add(context: KeyedContext, request: CounterAddRequest) {
2929
updateCounter(context, request.value)
3030
}
3131

32-
override suspend fun get(context: RestateContext, request: CounterRequest): GetResponse {
32+
override suspend fun get(context: KeyedContext, request: CounterRequest): GetResponse {
3333
return getResponse { value = context.get(TOTAL) ?: 0L }
3434
}
3535

3636
override suspend fun getAndAdd(
37-
context: RestateContext,
37+
context: KeyedContext,
3838
request: CounterAddRequest
3939
): CounterUpdateResult {
4040
LOG.info("Invoked get and add with " + request.value)
@@ -45,7 +45,7 @@ class CounterKt : CounterRestateKt.CounterRestateKtImplBase() {
4545
}
4646
}
4747

48-
private suspend fun updateCounter(context: RestateContext, add: Long): Pair<Long, Long> {
48+
private suspend fun updateCounter(context: KeyedContext, add: Long): Pair<Long, Long> {
4949
val currentValue = context.get(TOTAL) ?: 0L
5050
val newValue = currentValue + add
5151

‎protoc-gen-restate/src/main/java/dev/restate/sdk/gen/RestateGen.java

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import com.salesforce.jprotoc.ProtoTypeMap;
1616
import com.salesforce.jprotoc.ProtocPlugin;
1717
import dev.restate.generated.ext.Ext;
18+
import dev.restate.generated.ext.ServiceType;
1819
import java.util.ArrayList;
1920
import java.util.Collections;
2021
import java.util.List;
@@ -120,6 +121,12 @@ private ServiceContext buildServiceContext(
120121
serviceContext.serviceName = serviceProto.getName();
121122
serviceContext.deprecated = serviceProto.getOptions().getDeprecated();
122123

124+
// Resolve context type
125+
serviceContext.contextType =
126+
serviceProto.getOptions().getExtension(Ext.serviceType) == ServiceType.UNKEYED
127+
? "UnkeyedContext"
128+
: "KeyedContext";
129+
123130
// Resolve javadoc
124131
DescriptorProtos.SourceCodeInfo.Location serviceLocation =
125132
locations.stream()
@@ -215,6 +222,7 @@ private static class ServiceContext {
215222
public String packageName;
216223
public String className;
217224
public String serviceName;
225+
public String contextType;
218226
public String apidoc;
219227
public boolean deprecated;
220228
public final List<MethodContext> methods = new ArrayList<>();

‎protoc-gen-restate/src/main/resources/javaStub.mustache

Lines changed: 18 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,8 @@
22
package {{packageName}};
33
{{/packageName}}
44

5-
import dev.restate.sdk.RestateContext;
5+
import dev.restate.sdk.UnkeyedContext;
6+
import dev.restate.sdk.KeyedContext;
67
import dev.restate.sdk.Awaitable;
78
import dev.restate.sdk.common.syscalls.Syscalls;
89
import java.time.Duration;
@@ -15,24 +16,17 @@ public class {{className}} {
1516
private {{className}}() {}
1617

1718
/**
18-
* Create a new client.
19-
*/
20-
public static {{serviceName}}RestateClient newClient() {
21-
return newClient(RestateContext.fromSyscalls(Syscalls.current()));
22-
}
23-
24-
/**
25-
* Create a new client from the given {@link RestateContext}.
19+
* Create a new client from the given {@link KeyedContext}.
2620
*/
27-
public static {{serviceName}}RestateClient newClient(RestateContext ctx) {
21+
public static {{serviceName}}RestateClient newClient(UnkeyedContext ctx) {
2822
return new {{serviceName}}RestateClient(ctx);
2923
}
3024

3125
{{{apidoc}}}
3226
public static final class {{serviceName}}RestateClient {
33-
private final RestateContext ctx;
27+
private final UnkeyedContext ctx;
3428
35-
{{serviceName}}RestateClient(RestateContext ctx) {
29+
{{serviceName}}RestateClient(UnkeyedContext ctx) {
3630
this.ctx = ctx;
3731
}
3832

@@ -57,9 +51,9 @@ public class {{className}} {
5751
}
5852

5953
public static final class {{serviceName}}RestateOneWayClient {
60-
private final RestateContext ctx;
54+
private final UnkeyedContext ctx;
6155
62-
{{serviceName}}RestateOneWayClient(RestateContext ctx) {
56+
{{serviceName}}RestateOneWayClient(UnkeyedContext ctx) {
6357
this.ctx = ctx;
6458
}
6559

@@ -74,10 +68,10 @@ public class {{className}} {
7468
}
7569

7670
public static final class {{serviceName}}RestateDelayedClient {
77-
private final RestateContext ctx;
71+
private final UnkeyedContext ctx;
7872
private final Duration delay;
7973
80-
{{serviceName}}RestateDelayedClient(RestateContext ctx, Duration delay) {
74+
{{serviceName}}RestateDelayedClient(UnkeyedContext ctx, Duration delay) {
8175
this.ctx = ctx;
8276
this.delay = delay;
8377
}
@@ -100,7 +94,7 @@ public class {{className}} {
10094
@java.lang.Deprecated
10195
{{/deprecated}}
10296
{{{apidoc}}}
103-
public {{#isOutputEmpty}}void{{/isOutputEmpty}}{{^isOutputEmpty}}{{outputType}}{{/isOutputEmpty}} {{methodName}}(RestateContext context{{^isInputEmpty}}, {{inputType}} request{{/isInputEmpty}}) throws dev.restate.sdk.common.TerminalException {
97+
public {{#isOutputEmpty}}void{{/isOutputEmpty}}{{^isOutputEmpty}}{{outputType}}{{/isOutputEmpty}} {{methodName}}({{contextType}} context{{^isInputEmpty}}, {{inputType}} request{{/isInputEmpty}}) throws dev.restate.sdk.common.TerminalException {
10498
throw new dev.restate.sdk.common.TerminalException(dev.restate.sdk.common.TerminalException.Code.UNIMPLEMENTED);
10599
}
106100

@@ -120,34 +114,34 @@ public class {{className}} {
120114
private static final class HandlerAdapter<Req, Resp> implements
121115
io.grpc.stub.ServerCalls.UnaryMethod<Req, Resp> {
122116
123-
private final java.util.function.BiFunction<RestateContext, Req, Resp> handler;
117+
private final java.util.function.BiFunction<KeyedContext, Req, Resp> handler;
124118
125-
private HandlerAdapter(java.util.function.BiFunction<RestateContext, Req, Resp> handler) {
119+
private HandlerAdapter(java.util.function.BiFunction<KeyedContext, Req, Resp> handler) {
126120
this.handler = handler;
127121
}
128122

129123
@Override
130124
public void invoke(Req request, io.grpc.stub.StreamObserver<Resp> responseObserver) {
131-
responseObserver.onNext(handler.apply(RestateContext.fromSyscalls(Syscalls.current()), request));
125+
responseObserver.onNext(handler.apply(KeyedContext.fromSyscalls(Syscalls.current()), request));
132126
responseObserver.onCompleted();
133127
}
134128

135-
private static <Req, Resp> HandlerAdapter<Req, Resp> of(java.util.function.BiFunction<RestateContext, Req, Resp> handler) {
129+
private static <Req, Resp> HandlerAdapter<Req, Resp> of(java.util.function.BiFunction<KeyedContext, Req, Resp> handler) {
136130
return new HandlerAdapter<>(handler);
137131
}
138132

139-
private static <Resp> HandlerAdapter<com.google.protobuf.Empty, Resp> of(java.util.function.Function<RestateContext, Resp> handler) {
133+
private static <Resp> HandlerAdapter<com.google.protobuf.Empty, Resp> of(java.util.function.Function<KeyedContext, Resp> handler) {
140134
return new HandlerAdapter<>((ctx, e) -> handler.apply(ctx));
141135
}
142136

143-
private static <Req> HandlerAdapter<Req, com.google.protobuf.Empty> of(java.util.function.BiConsumer<RestateContext, Req> handler) {
137+
private static <Req> HandlerAdapter<Req, com.google.protobuf.Empty> of(java.util.function.BiConsumer<KeyedContext, Req> handler) {
144138
return new HandlerAdapter<>((ctx, req) -> {
145139
handler.accept(ctx, req);
146140
return com.google.protobuf.Empty.getDefaultInstance();
147141
});
148142
}
149143

150-
private static HandlerAdapter<com.google.protobuf.Empty, com.google.protobuf.Empty> of(java.util.function.Consumer<RestateContext> handler) {
144+
private static HandlerAdapter<com.google.protobuf.Empty, com.google.protobuf.Empty> of(java.util.function.Consumer<KeyedContext> handler) {
151145
return new HandlerAdapter<>((ctx, req) -> {
152146
handler.accept(ctx);
153147
return com.google.protobuf.Empty.getDefaultInstance();

‎protoc-gen-restate/src/main/resources/ktStub.mustache

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,14 @@
22
package {{packageName}};
33
{{/packageName}}
44

5-
import dev.restate.sdk.kotlin.RestateContext;
5+
import dev.restate.sdk.kotlin.UnkeyedContext;
6+
import dev.restate.sdk.kotlin.KeyedContext;
67
import dev.restate.sdk.kotlin.Awaitable;
78
import dev.restate.sdk.kotlin.RestateKtService;
89
import dev.restate.sdk.common.syscalls.Syscalls;
910
import io.grpc.kotlin.ClientCalls.unaryRpc
1011
import io.grpc.kotlin.ServerCalls.unaryServerMethodDefinition
1112
import {{packageName}}.{{serviceName}}Grpc.getServiceDescriptor;
12-
import dev.restate.sdk.kotlin.restateContextFromSyscalls;
1313
import kotlin.time.Duration
1414

1515
{{#deprecated}}
@@ -18,14 +18,14 @@ import kotlin.time.Duration
1818
public object {{className}} {
1919
2020
/**
21-
* Create a new client from the given [RestateContext].
21+
* Create a new client from the given [UnkeyedContext].
2222
*/
23-
fun newClient(ctx: RestateContext): {{serviceName}}RestateKtClient {
23+
fun newClient(ctx: UnkeyedContext): {{serviceName}}RestateKtClient {
2424
return {{serviceName}}RestateKtClient(ctx);
2525
}
2626

2727
{{{javadoc}}}
28-
public class {{serviceName}}RestateKtClient(private val ctx: RestateContext) {
28+
public class {{serviceName}}RestateKtClient(private val ctx: UnkeyedContext) {
2929
// Create a variant of this client to execute oneWay calls.
3030
public fun oneWay(): {{serviceName}}RestateKtOneWayClient {
3131
return {{serviceName}}RestateKtOneWayClient(ctx);
@@ -46,7 +46,7 @@ public object {{className}} {
4646
{{/methods}}
4747
}
4848

49-
public class {{serviceName}}RestateKtOneWayClient(private val ctx: RestateContext) {
49+
public class {{serviceName}}RestateKtOneWayClient(private val ctx: UnkeyedContext) {
5050
{{#methods}}
5151
{{#deprecated}}@Deprecated{{/deprecated}}
5252
{{{javadoc}}}
@@ -57,7 +57,7 @@ public object {{className}} {
5757
{{/methods}}
5858
}
5959

60-
public class {{serviceName}}RestateKtDelayedClient(private val ctx: RestateContext, private val delay: Duration) {
60+
public class {{serviceName}}RestateKtDelayedClient(private val ctx: UnkeyedContext, private val delay: Duration) {
6161
{{#methods}}
6262
{{#deprecated}}@Deprecated{{/deprecated}}
6363
{{{javadoc}}}
@@ -78,7 +78,7 @@ public object {{className}} {
7878
@Deprecated
7979
{{/deprecated}}
8080
{{{javadoc}}}
81-
public open suspend fun {{methodName}}(context: RestateContext{{^isInputEmpty}}, request: {{inputType}} {{/isInputEmpty}}){{^isOutputEmpty}}: {{outputType}}{{/isOutputEmpty}} {
81+
public open suspend fun {{methodName}}(context: {{contextType}}{{^isInputEmpty}}, request: {{inputType}} {{/isInputEmpty}}){{^isOutputEmpty}}: {{outputType}}{{/isOutputEmpty}} {
8282
throw dev.restate.sdk.common.TerminalException(dev.restate.sdk.common.TerminalException.Code.UNIMPLEMENTED);
8383
}
8484

@@ -91,18 +91,18 @@ public object {{className}} {
9191
descriptor = {{packageName}}.{{serviceName}}Grpc.{{methodDescriptorGetter}}(),
9292
implementation = {
9393
{{#isInputEmpty}}{{#isOutputEmpty}}
94-
{{methodName}}(restateContextFromSyscalls(Syscalls.current()))
94+
{{methodName}}(KeyedContext.fromSyscalls(Syscalls.current()))
9595
return@unaryServerMethodDefinition com.google.protobuf.Empty.getDefaultInstance()
9696
{{/isOutputEmpty}}{{/isInputEmpty}}
9797
{{#isInputEmpty}}{{^isOutputEmpty}}
98-
return@unaryServerMethodDefinition {{methodName}}(restateContextFromSyscalls(Syscalls.current()))
98+
return@unaryServerMethodDefinition {{methodName}}(KeyedContext.fromSyscalls(Syscalls.current()))
9999
{{/isOutputEmpty}}{{/isInputEmpty}}
100100
{{^isInputEmpty}}{{#isOutputEmpty}}
101-
{{methodName}}(restateContextFromSyscalls(Syscalls.current()), it)
101+
{{methodName}}(KeyedContext.fromSyscalls(Syscalls.current()), it)
102102
return@unaryServerMethodDefinition com.google.protobuf.Empty.getDefaultInstance()
103103
{{/isOutputEmpty}}{{/isInputEmpty}}
104104
{{^isInputEmpty}}{{^isOutputEmpty}}
105-
return@unaryServerMethodDefinition {{methodName}}(restateContextFromSyscalls(Syscalls.current()), it)
105+
return@unaryServerMethodDefinition {{methodName}}(KeyedContext.fromSyscalls(Syscalls.current()), it)
106106
{{/isOutputEmpty}}{{/isInputEmpty}}
107107
}
108108
))

‎sdk-api-kotlin/src/main/kotlin/dev/restate/sdk/kotlin/RestateContextImpl.kt renamed to ‎sdk-api-kotlin/src/main/kotlin/dev/restate/sdk/kotlin/ContextImpl.kt

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,7 @@ import kotlin.time.Duration
2424
import kotlin.time.toJavaDuration
2525
import kotlinx.coroutines.*
2626

27-
internal class RestateContextImpl internal constructor(private val syscalls: Syscalls) :
28-
RestateContext {
27+
internal class ContextImpl internal constructor(private val syscalls: Syscalls) : KeyedContext {
2928
override suspend fun <T : Any> get(key: StateKey<T>): T? {
3029
val deferred: Deferred<ByteString> =
3130
suspendCancellableCoroutine { cont: CancellableContinuation<Deferred<ByteString>> ->

0 commit comments

Comments
 (0)