Scope and limit key (#597)

Author: slinkydeveloper

.github/workflows/integration-ci.yaml

@@ -107,9 +107,12 @@ jobs:
 
       - name: Run test tool
         continue-on-error: ${{ inputs.continueOnError == 'true' }}
-        uses: restatedev/e2e/sdk-tests@v2.1
+        uses: restatedev/e2e/sdk-tests@v2.2
         with:
           envVars: ${{ inputs.envVars }}
+          # Scope/limit key need the FFM state machine (protocol v7, JRE >= 23). On the legacy
+          # state machine (JRE < 23) those tests can't pass, so exclude them there.
+          exclusionsFile: ${{ matrix.jreVersion < 23 && '.tools/exclusions-legacy-statemachine.yaml' || '' }}
           testArtifactOutput: ${{ inputs.testArtifactOutput != '' && format('{0}-jre{1}', inputs.testArtifactOutput, matrix.jreVersion) || format('sdk-java-jre{0}-integration-test-report', matrix.jreVersion) }}
           restateContainerImage: ${{ inputs.restateCommit != '' && 'localhost/restatedev/restate-commit-download:latest' || (inputs.restateImage != '' && inputs.restateImage || 'ghcr.io/restatedev/restate:main') }}
           serviceContainerImage: ${{ inputs.serviceImage != '' && inputs.serviceImage || 'restatedev/test-services-java' }}

.github/workflows/integration.yaml

@@ -117,9 +117,12 @@ jobs:
 
       - name: Run test tool
         continue-on-error: ${{ inputs.continueOnError == 'true' }}
-        uses: restatedev/e2e/sdk-tests@v2.1
+        uses: restatedev/e2e/sdk-tests@v2.2
         with:
           envVars: ${{ inputs.envVars }}
+          # Scope/limit key need the FFM state machine (protocol v7, JRE >= 23). On the legacy
+          # state machine (JRE < 23) those tests can't pass, so exclude them there.
+          exclusionsFile: ${{ matrix.jreVersion < 23 && '.tools/exclusions-legacy-statemachine.yaml' || '' }}
           testArtifactOutput: ${{ inputs.testArtifactOutput != '' && format('{0}-jre{1}', inputs.testArtifactOutput, matrix.jreVersion) || format('sdk-java-jre{0}-integration-test-report', matrix.jreVersion) }}
           restateContainerImage: ${{ inputs.restateCommit != '' && 'localhost/restatedev/restate-commit-download:latest' || (inputs.restateImage != '' && inputs.restateImage || 'ghcr.io/restatedev/restate:main') }}
           serviceContainerImage: ${{ inputs.serviceImage != '' && inputs.serviceImage || 'restatedev/test-services-java' }}

.tools/exclusions-legacy-statemachine.yaml

@@ -0,0 +1,9 @@
+# Tests excluded when the SDK runs on the legacy (pure-Java) state machine, i.e. on JRE < 23,
+# where the Panama/FFM state machine is unavailable and only service protocol v6 is negotiated.
+#
+# Scope and limit key require service protocol v7 (native/FFM state machine, JRE >= 23), so these
+# tests must be skipped on the legacy state machine. This file is passed via `--exclusions-file`
+# only for the JRE 17 and 21 CI legs; the JRE 25 leg runs them.
+exclusions:
+  "default":
+  - "dev.restate.sdktesting.tests.ServiceToServiceScopeConcurrency.scopeAndLimitKeyArePropagatedOnServiceToServiceCalls"

client-kotlin/src/main/kotlin/dev/restate/client/kotlin/ingress.kt

@@ -271,6 +271,108 @@ val <Res> Response<Res>.response: Res
 val <Res> SendResponse<Res>.sendStatus: SendResponse.SendStatus
   get() = this.sendStatus()
 
+/**
+ * **PREVIEW:** Returns a [ScopedKotlinClient] that routes all outgoing calls within the given
+ * scope.
+ *
+ * **NOTE:** This API is in preview and is not enabled by default. To use it in restate-server 1.7,
+ * enable the flow control and protocol v7 experimental features, via
+ * `RESTATE_EXPERIMENTAL_ENABLE_PROTOCOL_V7=true` and `RESTATE_EXPERIMENTAL_ENABLE_VQUEUES=true`.
+ * These can be enabled only on **new clusters**, for more info check out
+ * https://docs.restate.dev/services/flow-control#enabling-flow-control. If these experimental
+ * features aren't enabled, the call fails with a retryable error and keeps retrying until they are.
+ *
+ * A scope is a sub-grouping of resources (invocations, virtual object instances, workflow
+ * instances, concurrency limits) within the Restate cluster. It becomes part of the target identity
+ * tuple:
+ * - `scope, service, handler, idempotencyKey?`
+ * - `scope, virtualObject, objectKey, handler, idempotencyKey?`
+ * - `scope, workflow, workflowKey, handler`
+ *
+ * Under the hood, the scope contributes to the partition key, so all resources in a scope get
+ * co-located by the restate-server.
+ *
+ * Omitting the scope (i.e. using the regular `service` / `workflow` methods) is equivalent to
+ * calling with no scope, which is the existing behavior.
+ *
+ * The scope key must consist only of `[a-zA-Z0-9_.-]` characters, with `1 <= length <= 36` chars.
+ *
+ * Example usage:
+ * ```kotlin
+ * // Route a call into a named scope
+ * client.scope("tenant-123").service<MyService>().process(payload)
+ * ```
+ *
+ * @param scopeKey the scope identifier
+ * @see <a
+ *   href="https://docs.restate.dev/services/flow-control">https://docs.restate.dev/services/flow-control</a>
+ */
+@org.jetbrains.annotations.ApiStatus.Experimental
+fun Client.scope(scopeKey: String): ScopedKotlinClient = ScopedKotlinClient(this, scopeKey)
+
+/**
+ * **PREVIEW:** A client for making RPC calls within a specific scope.
+ *
+ * Obtain an instance via [Client.scope].
+ *
+ * @see Client.scope
+ */
+@org.jetbrains.annotations.ApiStatus.Experimental
+class ScopedKotlinClient
+@PublishedApi
+internal constructor(
+    @PublishedApi internal val client: Client,
+    @PublishedApi internal val scopeKey: String,
+) {
+  /** @see Client.service */
+  @org.jetbrains.annotations.ApiStatus.Experimental
+  inline fun <reified SVC : Any> service(): SVC {
+    return service(client, SVC::class.java, scopeKey)
+  }
+
+  /** @see Client.virtualObject */
+  @org.jetbrains.annotations.ApiStatus.Experimental
+  inline fun <reified SVC : Any> virtualObject(key: String): SVC {
+    return virtualObject(client, SVC::class.java, key, scopeKey)
+  }
+
+  /** @see Client.workflow */
+  @org.jetbrains.annotations.ApiStatus.Experimental
+  inline fun <reified SVC : Any> workflow(key: String): SVC {
+    return workflow(client, SVC::class.java, key, scopeKey)
+  }
+
+  /** @see Client.toService */
+  @org.jetbrains.annotations.ApiStatus.Experimental
+  inline fun <reified SVC : Any> toService(): KClientRequestBuilder<SVC> {
+    ReflectionUtils.mustHaveServiceAnnotation(SVC::class.java)
+    require(ReflectionUtils.isKotlinClass(SVC::class.java)) {
+      "Using Java classes with Kotlin's API is not supported"
+    }
+    return KClientRequestBuilder(client, SVC::class.java, null, scopeKey)
+  }
+
+  /** @see Client.toVirtualObject */
+  @org.jetbrains.annotations.ApiStatus.Experimental
+  inline fun <reified SVC : Any> toVirtualObject(key: String): KClientRequestBuilder<SVC> {
+    ReflectionUtils.mustHaveVirtualObjectAnnotation(SVC::class.java)
+    require(ReflectionUtils.isKotlinClass(SVC::class.java)) {
+      "Using Java classes with Kotlin's API is not supported"
+    }
+    return KClientRequestBuilder(client, SVC::class.java, key, scopeKey)
+  }
+
+  /** @see Client.toWorkflow */
+  @org.jetbrains.annotations.ApiStatus.Experimental
+  inline fun <reified SVC : Any> toWorkflow(key: String): KClientRequestBuilder<SVC> {
+    ReflectionUtils.mustHaveWorkflowAnnotation(SVC::class.java)
+    require(ReflectionUtils.isKotlinClass(SVC::class.java)) {
+      "Using Java classes with Kotlin's API is not supported"
+    }
+    return KClientRequestBuilder(client, SVC::class.java, key, scopeKey)
+  }
+}
+
 /**
  * Create a proxy client for a Restate service.
  *
@@ -329,15 +431,15 @@ inline fun <reified SVC : Any> Client.workflow(key: String): SVC {
  * @return a proxy that intercepts method calls and executes them via the client
  */
 @PublishedApi
-internal fun <SVC : Any> service(client: Client, clazz: Class<SVC>): SVC {
+internal fun <SVC : Any> service(client: Client, clazz: Class<SVC>, scope: String? = null): SVC {
   ReflectionUtils.mustHaveServiceAnnotation(clazz)
   require(ReflectionUtils.isKotlinClass(clazz)) {
     "Using Java classes with Kotlin's API is not supported"
   }
 
   val serviceName = ReflectionUtils.extractServiceName(clazz)
   return ProxySupport.createProxy(clazz) { invocation ->
-    val request = invocation.captureInvocation(serviceName, null).toRequest()
+    val request = invocation.captureInvocation(serviceName, null, scope).toRequest()
     @Suppress("UNCHECKED_CAST") val continuation = invocation.arguments.last() as Continuation<Any?>
 
     // Start a coroutine that calls the client and resumes the continuation
@@ -356,15 +458,20 @@ internal fun <SVC : Any> service(client: Client, clazz: Class<SVC>): SVC {
  * @return a proxy that intercepts method calls and executes them via the client
  */
 @PublishedApi
-internal fun <SVC : Any> virtualObject(client: Client, clazz: Class<SVC>, key: String): SVC {
+internal fun <SVC : Any> virtualObject(
+    client: Client,
+    clazz: Class<SVC>,
+    key: String,
+    scope: String? = null,
+): SVC {
   ReflectionUtils.mustHaveVirtualObjectAnnotation(clazz)
   require(ReflectionUtils.isKotlinClass(clazz)) {
     "Using Java classes with Kotlin's API is not supported"
   }
 
   val serviceName = ReflectionUtils.extractServiceName(clazz)
   return ProxySupport.createProxy(clazz) { invocation ->
-    val request = invocation.captureInvocation(serviceName, key).toRequest()
+    val request = invocation.captureInvocation(serviceName, key, scope).toRequest()
     @Suppress("UNCHECKED_CAST") val continuation = invocation.arguments.last() as Continuation<Any?>
 
     // Start a coroutine that calls the client and resumes the continuation
@@ -383,15 +490,20 @@ internal fun <SVC : Any> virtualObject(client: Client, clazz: Class<SVC>, key: S
  * @return a proxy that intercepts method calls and executes them via the client
  */
 @PublishedApi
-internal fun <SVC : Any> workflow(client: Client, clazz: Class<SVC>, key: String): SVC {
+internal fun <SVC : Any> workflow(
+    client: Client,
+    clazz: Class<SVC>,
+    key: String,
+    scope: String? = null,
+): SVC {
   ReflectionUtils.mustHaveWorkflowAnnotation(clazz)
   require(ReflectionUtils.isKotlinClass(clazz)) {
     "Using Java classes with Kotlin's API is not supported"
   }
 
   val serviceName = ReflectionUtils.extractServiceName(clazz)
   return ProxySupport.createProxy(clazz) { invocation ->
-    val request = invocation.captureInvocation(serviceName, key).toRequest()
+    val request = invocation.captureInvocation(serviceName, key, scope).toRequest()
     @Suppress("UNCHECKED_CAST") val continuation = invocation.arguments.last() as Continuation<Any?>
 
     // Start a coroutine that calls the client and resumes the continuation
@@ -414,6 +526,7 @@ internal constructor(
     private val client: Client,
     private val clazz: Class<SVC>,
     private val key: String?,
+    private val scope: String? = null,
 ) {
   /**
    * Create a request by invoking a method on the target.
@@ -428,7 +541,7 @@ internal constructor(
   suspend fun <Res> request(block: suspend SVC.() -> Res): KClientRequest<Any?, Res> {
     return KClientRequestImpl(
         client,
-        RequestCaptureProxy(clazz, key).capture(block as suspend SVC.() -> Any?).toRequest(),
+        RequestCaptureProxy(clazz, key, scope).capture(block as suspend SVC.() -> Any?).toRequest(),
     )
         as KClientRequest<Any?, Res>
   }

client/src/main/java/dev/restate/client/Client.java

@@ -531,6 +531,62 @@ default Response<Output<Res>> getOutput() throws IngressException {
     }
   }
 
+  /**
+   * <b>PREVIEW:</b> Returns a {@link ScopedClient} that routes all outgoing calls within the given
+   * scope.
+   *
+   * <p><b>NOTE:</b> This API is in preview and is not enabled by default. To use it in
+   * restate-server 1.7, enable the flow control and protocol v7 experimental features, via {@code
+   * RESTATE_EXPERIMENTAL_ENABLE_PROTOCOL_V7=true} and {@code
+   * RESTATE_EXPERIMENTAL_ENABLE_VQUEUES=true}. These can be enabled only on <b>new clusters</b>,
+   * for more info check out https://docs.restate.dev/services/flow-control#enabling-flow-control.
+   * If these experimental features aren't enabled, the call fails with a retryable error and keeps
+   * retrying until they are.
+   *
+   * <p>A scope is a sub-grouping of resources (invocations, virtual object instances, workflow
+   * instances, concurrency limits) within the Restate cluster. It becomes part of the target
+   * identity tuple:
+   *
+   * <ul>
+   *   <li>{@code scope, service, handler, idempotencyKey?}
+   *   <li>{@code scope, virtualObject, objectKey, handler, idempotencyKey?}
+   *   <li>{@code scope, workflow, workflowKey, handler}
+   * </ul>
+   *
+   * <p>Under the hood, the scope contributes to the partition key, so all resources in a scope get
+   * co-located by the restate-server.
+   *
+   * <p>Omitting the scope (i.e. using the regular {@link #service(Class)} / {@link #workflow(Class,
+   * String)} methods) is equivalent to calling with no scope, which is the existing behavior.
+   *
+   * <p>The scope key must consist only of {@code [a-zA-Z0-9_.-]} characters, with {@code 1 <=
+   * length <= 36} chars.
+   *
+   * <pre>{@code
+   * Client client = Client.connect("http://localhost:8080");
+   *
+   * // Route a call into a named scope
+   * client.scope("tenant-123").service(MyService.class).process(payload);
+   *
+   * // Idempotency keys are scoped — "req-1" in "tenant-123" is distinct from "req-1" in "tenant-456"
+   * client.scope("tenant-123").serviceHandle(MyService.class)
+   *     .call(MyService::process, payload, InvocationOptions.idempotencyKey("req-1").build());
+   *
+   * // Combine with a limit key to enforce per-scope concurrency limits
+   * client.scope("tenant-123").workflowHandle(MyWorkflow.class, "wf-key")
+   *     .call(MyWorkflow::run, input, InvocationOptions.limitKey("api-key/user42").build());
+   * }</pre>
+   *
+   * @param scopeKey the scope identifier
+   * @return a {@link ScopedClient}
+   * @see <a
+   *     href="https://docs.restate.dev/services/flow-control">https://docs.restate.dev/services/flow-control</a>
+   */
+  @org.jetbrains.annotations.ApiStatus.Experimental
+  default ScopedClient scope(String scopeKey) {
+    return new ScopedClient(this, scopeKey);
+  }
+
   /**
    * Simple API to invoke a Restate service from the ingress.
    *

client/src/main/java/dev/restate/client/ClientServiceHandleImpl.java

@@ -29,14 +29,21 @@ final class ClientServiceHandleImpl<SVC> implements ClientServiceHandle<SVC> {
   private final Class<SVC> clazz;
   private final String serviceName;
   private final @Nullable String key;
+  private final @Nullable String scope;
 
   private MethodInfoCollector<SVC> methodInfoCollector;
 
   ClientServiceHandleImpl(Client innerClient, Class<SVC> clazz, @Nullable String key) {
+    this(innerClient, clazz, key, null);
+  }
+
+  ClientServiceHandleImpl(
+      Client innerClient, Class<SVC> clazz, @Nullable String key, @Nullable String scope) {
     this.innerClient = innerClient;
     this.clazz = clazz;
     this.serviceName = ReflectionUtils.extractServiceName(clazz);
     this.key = key;
+    this.scope = scope;
   }
 
   @SuppressWarnings("unchecked")
@@ -46,6 +53,7 @@ public <I, O> CompletableFuture<Response<O>> callAsync(
     MethodInfo methodInfo = getMethodInfoCollector().resolve(s, input);
     return innerClient.callAsync(
         toRequest(
+            scope,
             serviceName,
             key,
             methodInfo.getHandlerName(),
@@ -62,6 +70,7 @@ public <I> CompletableFuture<Response<Void>> callAsync(
     MethodInfo methodInfo = getMethodInfoCollector().resolve(s, input);
     return innerClient.callAsync(
         toRequest(
+            scope,
             serviceName,
             key,
             methodInfo.getHandlerName(),
@@ -78,6 +87,7 @@ public <O> CompletableFuture<Response<O>> callAsync(
     MethodInfo methodInfo = getMethodInfoCollector().resolve(s);
     return innerClient.callAsync(
         toRequest(
+            scope,
             serviceName,
             key,
             methodInfo.getHandlerName(),
@@ -93,6 +103,7 @@ public CompletableFuture<Response<Void>> callAsync(
     MethodInfo methodInfo = getMethodInfoCollector().resolve(s);
     return innerClient.callAsync(
         toRequest(
+            scope,
             serviceName,
             key,
             methodInfo.getHandlerName(),
@@ -109,6 +120,7 @@ public <I, O> CompletableFuture<SendResponse<O>> sendAsync(
     MethodInfo methodInfo = getMethodInfoCollector().resolve(s, input);
     return innerClient.sendAsync(
         toRequest(
+            scope,
             serviceName,
             key,
             methodInfo.getHandlerName(),
@@ -126,6 +138,7 @@ public <I> CompletableFuture<SendResponse<Void>> sendAsync(
     MethodInfo methodInfo = getMethodInfoCollector().resolve(s, input);
     return innerClient.sendAsync(
         toRequest(
+            scope,
             serviceName,
             key,
             methodInfo.getHandlerName(),
@@ -143,6 +156,7 @@ public <O> CompletableFuture<SendResponse<O>> sendAsync(
     MethodInfo methodInfo = getMethodInfoCollector().resolve(s);
     return innerClient.sendAsync(
         toRequest(
+            scope,
             serviceName,
             key,
             methodInfo.getHandlerName(),
@@ -159,6 +173,7 @@ public CompletableFuture<SendResponse<Void>> sendAsync(
     MethodInfo methodInfo = getMethodInfoCollector().resolve(s);
     return innerClient.sendAsync(
         toRequest(
+            scope,
             serviceName,
             key,
             methodInfo.getHandlerName(),