Tests for scope and concurrency limit

e2e-tests/src/main/kotlin/dev/restate/sdktesting/junit/TestSuites.kt

@@ -11,6 +11,7 @@ package dev.restate.sdktesting.junit
 import dev.restate.sdktesting.tests.AwakeableIngressEndpointTest
 import dev.restate.sdktesting.tests.AwakeableLeaderTransferTest
 import dev.restate.sdktesting.tests.BackwardCompatibilityTest
+import dev.restate.sdktesting.tests.ConcurrencyLimitTest
 import dev.restate.sdktesting.tests.ForwardCompatibilityTest
 import dev.restate.sdktesting.tests.IngressTest
 import dev.restate.sdktesting.tests.InvokerMemoryTest
@@ -20,6 +21,7 @@ import dev.restate.sdktesting.tests.OpenAPITest
 import dev.restate.sdktesting.tests.PauseResumeChangingDeploymentTest
 import dev.restate.sdktesting.tests.PauseResumeTest
 import dev.restate.sdktesting.tests.RestartAsNewInvocationTest
+import dev.restate.sdktesting.tests.ScopeIsolationTest
 import dev.restate.sdktesting.tests.StatePatchingTest
 import dev.restate.sdktesting.tests.TracingTest
 import dev.restate.sdktesting.tests.UpgradeWithInFlightInvocation
@@ -35,6 +37,7 @@ object TestSuites : SuiteProvider {
           emptyMap(),
           listOf(
               clazz<AwakeableIngressEndpointTest>(),
+              clazz<ConcurrencyLimitTest>(),
               clazz<IngressTest>(),
               clazz<InvokerMemoryTest>(),
               clazz<JournalRetentionTest>(),
@@ -43,6 +46,7 @@ object TestSuites : SuiteProvider {
               clazz<PauseResumeChangingDeploymentTest>(),
               clazz<PauseResumeTest>(),
               clazz<RestartAsNewInvocationTest>(),
+              clazz<ScopeIsolationTest>(),
               clazz<StatePatchingTest>(),
               clazz<TracingTest>(),
               clazz<UpgradeWithNewInvocation>(),
@@ -56,6 +60,8 @@ object TestSuites : SuiteProvider {
               "RESTATE_DEFAULT_NUM_PARTITIONS" to "4",
           ),
           listOf(
+              clazz<ConcurrencyLimitTest>(),
+              clazz<ScopeIsolationTest>(),
               clazz<AwakeableIngressEndpointTest>(),
               clazz<AwakeableLeaderTransferTest>(),
               clazz<IngressTest>(),

e2e-tests/src/main/kotlin/dev/restate/sdktesting/tests/ConcurrencyLimitTest.kt

@@ -0,0 +1,171 @@
+// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
+//
+// This file is part of the Restate SDK Test suite tool,
+// 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-test-suite/blob/main/LICENSE
+package dev.restate.sdktesting.tests
+
+import dev.restate.client.Client
+import dev.restate.client.kotlin.attachSuspend
+import dev.restate.client.kotlin.response
+import dev.restate.client.kotlin.toVirtualObject
+import dev.restate.sdk.annotation.Handler
+import dev.restate.sdk.annotation.Name
+import dev.restate.sdk.annotation.Service
+import dev.restate.sdk.annotation.Shared
+import dev.restate.sdk.annotation.VirtualObject
+import dev.restate.sdk.common.TerminalException
+import dev.restate.sdk.endpoint.Endpoint
+import dev.restate.sdk.kotlin.awakeable
+import dev.restate.sdk.kotlin.awakeableHandle
+import dev.restate.sdk.kotlin.call
+import dev.restate.sdk.kotlin.get
+import dev.restate.sdk.kotlin.resolve
+import dev.restate.sdk.kotlin.set
+import dev.restate.sdk.kotlin.state
+import dev.restate.sdk.kotlin.toVirtualObject
+import dev.restate.sdktesting.infra.*
+import dev.restate.serde.TypeTag
+import java.net.URI
+import java.util.UUID
+import kotlin.time.Duration.Companion.seconds
+import kotlinx.serialization.json.Json
+import org.assertj.core.api.Assertions.assertThat
+import org.awaitility.kotlin.await
+import org.awaitility.kotlin.withAlias
+import org.junit.jupiter.api.DisplayName
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.RegisterExtension
+
+/**
+ * Verifies the runtime enforces the rule-book action-concurrency limit on a scope. Scoped
+ * invocations are limited to N in flight, and the held excess progresses as the running ones
+ * complete.
+ */
+class ConcurrencyLimitTest {
+
+  @Service
+  @Name("BlockingProxy")
+  class BlockingProxy {
+    @Handler
+    suspend fun block(key: String): String =
+        toVirtualObject<Blocker>(key).request { run() }.call().await()
+  }
+
+  @VirtualObject
+  @Name("Blocker")
+  class Blocker {
+    @Handler
+    suspend fun run(): String {
+      val awk = awakeable<String>()
+      state().set("awk", awk.id)
+      return awk.await()
+    }
+
+    @Shared suspend fun getAwakeable(): String = state().get<String>("awk") ?: ""
+
+    @Shared
+    suspend fun resolveAwakeable(value: String) {
+      val id = state().get<String>("awk") ?: ""
+      if (id.isEmpty()) throw TerminalException("Awakeable not registered yet")
+      awakeableHandle(id).resolve(value)
+    }
+  }
+
+  companion object {
+    @RegisterExtension
+    @JvmField
+    val deployerExt: RestateDeployerExtension = RestateDeployerExtension {
+      withEndpoint(Endpoint.bind(BlockingProxy()).bind(Blocker()))
+      // Scoped invocations and rule-book require the vqueues experimental feature.
+      // TODO: drop this once the minimum supported Restate version is v1.8, where vqueues is on by
+      //   default.
+      withEnv("RESTATE_EXPERIMENTAL_ENABLE_VQUEUES", "true")
+      // Reduce rule-book activation latency so the test isn't gated on the default 30s poll.
+      withEnv("RESTATE_WORKER__RULE_BOOK_POLL_INTERVAL", "1s")
+    }
+  }
+
+  @Test
+  @DisplayName(
+      "Action concurrency limit on a scope holds excess invocations and releases on completion")
+  fun actionConcurrencyLimitIsRespected(
+      @InjectIngressURI ingressURI: URI,
+      @InjectAdminURI adminURI: URI,
+      @InjectContainerHandle(hostName = RESTATE_RUNTIME) runtimeHandle: ContainerHandle,
+      @InjectClient ingressClient: Client,
+  ) =
+      runTest(timeout = 120.seconds) {
+        val runId = UUID.randomUUID().toString().take(8)
+        val scope = "myscope-$runId"
+        val limit = 2
+        val invocationCount = 4
+        val blockerKeys = (0 until invocationCount).map { "block-key-$runId-$it" }
+
+        val ruleVersion =
+            upsertActionConcurrencyRule(adminURI, pattern = scope, actionConcurrency = limit)
+                .version
+        awaitRuleBookApplied(runtimeHandle, ruleVersion)
+
+        val outerIds =
+            blockerKeys.map { key ->
+              sendInvocationWithScope(
+                  ingressURI, scope, "BlockingProxy", "block", Json.encodeToString(key))
+            }
+
+        val blockerTargetFilter = "target LIKE 'Blocker/%/run'"
+
+        // Wait until exactly `limit` Blocker invocations exist; the rest are held by the rule.
+        await withAlias
+            "exactly $limit Blocker invocations are in flight" untilAsserted
+            {
+              assertThat(getAllInvocations(adminURI, blockerTargetFilter)).hasSize(limit)
+            }
+
+        // Resolve one awakeable at a time. After each resolve, a held outer becomes running and
+        // spawns its Blocker.
+        val unresolvedKeys = blockerKeys.toMutableSet()
+        repeat(invocationCount) {
+          var activeKey: String? = null
+          await withAlias
+              "find a Blocker (among unresolved keys) that has registered its awakeable" untilAsserted
+              {
+                val found =
+                    unresolvedKeys.firstNotNullOfOrNull { key ->
+                      val awkId =
+                          ingressClient
+                              .toVirtualObject<Blocker>(key)
+                              .request { getAwakeable() }
+                              .options(idempotentCallOptions)
+                              .call()
+                              .response
+                      if (awkId.isNotEmpty()) key else null
+                    }
+                assertThat(found).isNotNull
+                activeKey = found
+              }
+
+          ingressClient
+              .toVirtualObject<Blocker>(activeKey!!)
+              .request { resolveAwakeable("done") }
+              .options(idempotentCallOptions)
+              .call()
+          unresolvedKeys.remove(activeKey!!)
+        }
+
+        // All four outer BlockingProxy invocations must complete successfully.
+        outerIds.forEach { outerId ->
+          val response =
+              ingressClient
+                  .invocationHandle(outerId, TypeTag.of(String::class.java))
+                  .attachSuspend()
+                  .response
+          assertThat(response).isEqualTo("done")
+        }
+
+        bulkDeleteRules(adminURI, listOf(scope))
+      }
+}

e2e-tests/src/main/kotlin/dev/restate/sdktesting/tests/ScopeIsolationTest.kt

@@ -0,0 +1,93 @@
+// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
+//
+// This file is part of the Restate SDK Test suite tool,
+// 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-test-suite/blob/main/LICENSE
+package dev.restate.sdktesting.tests
+
+import dev.restate.client.Client
+import dev.restate.client.kotlin.attachSuspend
+import dev.restate.client.kotlin.response
+import dev.restate.sdk.annotation.Handler
+import dev.restate.sdk.annotation.Name
+import dev.restate.sdk.annotation.Service
+import dev.restate.sdk.endpoint.Endpoint
+import dev.restate.sdk.kotlin.runBlock
+import dev.restate.sdktesting.infra.*
+import dev.restate.serde.TypeTag
+import java.net.URI
+import java.util.UUID
+import kotlinx.serialization.json.Json
+import org.assertj.core.api.Assertions.assertThat
+import org.junit.jupiter.api.DisplayName
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.RegisterExtension
+
+/**
+ * Verifies that scope is part of an invocation's identity: the same idempotency key sent in two
+ * different scopes produces two distinct invocations rather than colliding on dedup. This is a
+ * runtime guarantee independent of any rule-book rule.
+ */
+class ScopeIsolationTest {
+
+  @Service
+  @Name("Random")
+  class Random {
+    @Handler suspend fun genRandomUUID(input: String) = runBlock { UUID.randomUUID().toString() }
+  }
+
+  companion object {
+    @RegisterExtension
+    @JvmField
+    val deployerExt: RestateDeployerExtension = RestateDeployerExtension {
+      withEndpoint(Endpoint.bind(Random()))
+      // Scoped invocations require the vqueues experimental feature.
+      // TODO: drop this once the minimum supported Restate version is v1.8,
+      //  where vqueues are enabled by default.
+      withEnv("RESTATE_EXPERIMENTAL_ENABLE_VQUEUES", "true")
+    }
+  }
+
+  @Test
+  @DisplayName("Same idempotency key in two different scopes produces two distinct invocations")
+  fun sameIdempotencyKeyAcrossScopesIsolates(
+      @InjectIngressURI ingressURI: URI,
+      @InjectClient ingressClient: Client,
+  ) = runTest {
+    val sharedIdempotencyKey = "shared-idempotency"
+    val scopeA = UUID.randomUUID().toString()
+    val scopeB = UUID.randomUUID().toString()
+    val body = Json.encodeToString("foo")
+
+    val idA =
+        sendInvocationWithScope(
+            ingressURI,
+            scopeA,
+            "Random",
+            "genRandomUUID",
+            body,
+            idempotencyKey = sharedIdempotencyKey)
+    val idB =
+        sendInvocationWithScope(
+            ingressURI,
+            scopeB,
+            "Random",
+            "genRandomUUID",
+            body,
+            idempotencyKey = sharedIdempotencyKey)
+
+    assertThat(idA).isNotEqualTo(idB)
+
+    // Each invocation generates its own random value via `runBlock`. If the runtime had
+    // collapsed the two calls under shared idempotency, both responses would be identical;
+    // distinct values prove the scope isolated them.
+    val respA =
+        ingressClient.invocationHandle(idA, TypeTag.of(String::class.java)).attachSuspend().response
+    val respB =
+        ingressClient.invocationHandle(idB, TypeTag.of(String::class.java)).attachSuspend().response
+    assertThat(respA).isNotEqualTo(respB)
+  }
+}