Add ctx.stateKeys() (#219)

sdk-api-kotlin/src/main/kotlin/dev/restate/sdk/kotlin/ContextImpl.kt

@@ -47,6 +47,25 @@ internal class ContextImpl internal constructor(private val syscalls: Syscalls)
     return key.serde().deserializeWrappingException(syscalls, readyResult.value!!)!!
   }
 
+  override suspend fun stateKeys(): Collection<String> {
+    val deferred: Deferred<Collection<String>> =
+        suspendCancellableCoroutine { cont: CancellableContinuation<Deferred<Collection<String>>> ->
+          syscalls.getKeys(completingContinuation(cont))
+        }
+
+    if (!deferred.isCompleted) {
+      suspendCancellableCoroutine { cont: CancellableContinuation<Unit> ->
+        syscalls.resolveDeferred(deferred, completingUnitContinuation(cont))
+      }
+    }
+
+    val readyResult = deferred.toResult()!!
+    if (!readyResult.isSuccess) {
+      throw readyResult.failure!!
+    }
+    return readyResult.value!!
+  }
+
   override suspend fun <T : Any> set(key: StateKey<T>, value: T) {
     val serializedValue = key.serde().serializeWrappingException(syscalls, value)!!
     return suspendCancellableCoroutine { cont: CancellableContinuation<Unit> ->

sdk-api-kotlin/src/main/kotlin/dev/restate/sdk/kotlin/api.kt

@@ -222,6 +222,13 @@ sealed interface KeyedContext : UnkeyedContext {
    */
   suspend fun <T : Any> get(key: StateKey<T>): T?
 
+  /**
+   * Gets all the known state keys for this service instance.
+   *
+   * @return the immutable collection of known state keys.
+   */
+  suspend fun stateKeys(): Collection<String>
+
   /**
    * Sets the given value under the given key, serializing the value using the registered
    * [dev.restate.sdk.core.serde.Serde] in the interceptor.

sdk-api-kotlin/src/test/kotlin/dev/restate/sdk/kotlin/EagerStateTest.kt

@@ -90,4 +90,14 @@ class EagerStateTest : EagerStateTestSuite() {
   override fun getClearAllAndGet(): BindableService {
     return GetClearAllAndGet()
   }
+
+  private class ListKeys : GreeterRestateKt.GreeterRestateKtImplBase() {
+    override suspend fun greet(context: KeyedContext, request: GreetingRequest): GreetingResponse {
+      return greetingResponse { message = context.stateKeys().joinToString(separator = ",") }
+    }
+  }
+
+  override fun listKeys(): BindableService {
+    return ListKeys()
+  }
 }

sdk-api/src/main/java/dev/restate/sdk/ContextImpl.java

@@ -17,6 +17,7 @@
 import dev.restate.sdk.common.syscalls.Syscalls;
 import io.grpc.MethodDescriptor;
 import java.time.Duration;
+import java.util.Collection;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
@@ -43,6 +44,17 @@ public <T> Optional<T> get(StateKey<T> key) {
         .map(bs -> Util.deserializeWrappingException(syscalls, key.serde(), bs));
   }
 
+  @Override
+  public Collection<String> stateKeys() {
+    Deferred<Collection<String>> deferred = Util.blockOnSyscall(syscalls::getKeys);
+
+    if (!deferred.isCompleted()) {
+      Util.<Void>blockOnSyscall(cb -> syscalls.resolveDeferred(deferred, cb));
+    }
+
+    return Util.unwrapResult(deferred.toResult());
+  }
+
   @Override
   public void clear(StateKey<?> key) {
     Util.<Void>blockOnSyscall(cb -> syscalls.clear(key.name(), cb));

sdk-api/src/main/java/dev/restate/sdk/KeyedContext.java

@@ -10,6 +10,7 @@
 
 import dev.restate.sdk.common.*;
 import dev.restate.sdk.common.syscalls.Syscalls;
+import java.util.Collection;
 import java.util.Optional;
 import javax.annotation.Nonnull;
 import javax.annotation.concurrent.NotThreadSafe;
@@ -34,6 +35,13 @@ public interface KeyedContext extends UnkeyedContext {
    */
   <T> Optional<T> get(StateKey<T> key);
 
+  /**
+   * Gets all the known state keys for this service instance.
+   *
+   * @return the immutable collection of known state keys.
+   */
+  Collection<String> stateKeys();
+
   /**
    * Clears the state stored under key.
    *

sdk-api/src/test/java/dev/restate/sdk/EagerStateTest.java

@@ -12,8 +12,10 @@
 
 import dev.restate.sdk.common.CoreSerdes;
 import dev.restate.sdk.common.StateKey;
+import dev.restate.sdk.common.TerminalException;
 import dev.restate.sdk.core.EagerStateTestSuite;
 import dev.restate.sdk.core.testservices.GreeterGrpc;
+import dev.restate.sdk.core.testservices.GreeterRestate;
 import dev.restate.sdk.core.testservices.GreetingRequest;
 import dev.restate.sdk.core.testservices.GreetingResponse;
 import io.grpc.BindableService;
@@ -119,4 +121,19 @@ public void greet(GreetingRequest request, StreamObserver<GreetingResponse> resp
   protected BindableService getClearAllAndGet() {
     return new GetClearAllAndGet();
   }
+
+  private static class ListKeys extends GreeterRestate.GreeterRestateImplBase {
+    @Override
+    public GreetingResponse greet(KeyedContext context, GreetingRequest request)
+        throws TerminalException {
+      return GreetingResponse.newBuilder()
+          .setMessage(String.join(",", context.stateKeys()))
+          .build();
+    }
+  }
+
+  @Override
+  protected BindableService listKeys() {
+    return new ListKeys();
+  }
 }

sdk-common/src/main/java/dev/restate/sdk/common/syscalls/Syscalls.java

@@ -14,6 +14,7 @@
 import io.grpc.Context;
 import io.grpc.MethodDescriptor;
 import java.time.Duration;
+import java.util.*;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -61,6 +62,8 @@ static Syscalls current() {
 
   void get(String name, SyscallCallback<Deferred<ByteString>> callback);
 
+  void getKeys(SyscallCallback<Deferred<Collection<String>>> callback);
+
   void clear(String name, SyscallCallback<Void> callback);
 
   void clearAll(SyscallCallback<Void> callback);

sdk-core/src/main/java/dev/restate/sdk/core/Entries.java

@@ -10,13 +10,16 @@
 
 import com.google.protobuf.ByteString;
 import com.google.protobuf.Empty;
+import com.google.protobuf.InvalidProtocolBufferException;
 import com.google.protobuf.MessageLite;
 import dev.restate.generated.service.protocol.Protocol;
 import dev.restate.generated.service.protocol.Protocol.*;
 import dev.restate.sdk.common.syscalls.Result;
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.trace.Span;
+import java.util.Collection;
 import java.util.function.Function;
+import java.util.stream.Collectors;
 
 final class Entries {
   static final String AWAKEABLE_IDENTIFIER_PREFIX = "prom_1";
@@ -183,6 +186,78 @@ void updateUserStateStorageWithCompletion(
     }
   }
 
+  static final class GetStateKeysEntry
+      extends CompletableJournalEntry<GetStateKeysEntryMessage, Collection<String>> {
+
+    static final GetStateKeysEntry INSTANCE = new GetStateKeysEntry();
+
+    private GetStateKeysEntry() {}
+
+    @Override
+    void trace(GetStateKeysEntryMessage expected, Span span) {
+      span.addEvent("GetStateKeys");
+    }
+
+    @Override
+    public boolean hasResult(GetStateKeysEntryMessage actual) {
+      return actual.getResultCase() != GetStateKeysEntryMessage.ResultCase.RESULT_NOT_SET;
+    }
+
+    @Override
+    void checkEntryHeader(GetStateKeysEntryMessage expected, MessageLite actual)
+        throws ProtocolException {
+      if (!(actual instanceof GetStateKeysEntryMessage)) {
+        throw ProtocolException.entryDoesNotMatch(expected, actual);
+      }
+    }
+
+    @Override
+    public Result<Collection<String>> parseEntryResult(GetStateKeysEntryMessage actual) {
+      if (actual.getResultCase() == GetStateKeysEntryMessage.ResultCase.VALUE) {
+        return Result.success(
+            actual.getValue().getKeysList().stream()
+                .map(ByteString::toStringUtf8)
+                .collect(Collectors.toUnmodifiableList()));
+      } else if (actual.getResultCase() == GetStateKeysEntryMessage.ResultCase.FAILURE) {
+        return Result.failure(Util.toRestateException(actual.getFailure()));
+      } else {
+        throw new IllegalStateException("GetStateKeysEntryMessage has not been completed.");
+      }
+    }
+
+    @Override
+    public Result<Collection<String>> parseCompletionResult(CompletionMessage actual) {
+      if (actual.getResultCase() == CompletionMessage.ResultCase.VALUE) {
+        GetStateKeysEntryMessage.StateKeys stateKeys;
+        try {
+          stateKeys = GetStateKeysEntryMessage.StateKeys.parseFrom(actual.getValue());
+        } catch (InvalidProtocolBufferException e) {
+          throw new ProtocolException(
+              "Cannot parse get state keys completion", e, ProtocolException.PROTOCOL_VIOLATION);
+        }
+        return Result.success(
+            stateKeys.getKeysList().stream()
+                .map(ByteString::toStringUtf8)
+                .collect(Collectors.toUnmodifiableList()));
+      } else if (actual.getResultCase() == CompletionMessage.ResultCase.FAILURE) {
+        return Result.failure(Util.toRestateException(actual.getFailure()));
+      }
+      return super.parseCompletionResult(actual);
+    }
+
+    @Override
+    GetStateKeysEntryMessage tryCompleteWithUserStateStorage(
+        GetStateKeysEntryMessage expected, UserStateStore userStateStore) {
+      if (userStateStore.isComplete()) {
+        return expected.toBuilder()
+            .setValue(
+                GetStateKeysEntryMessage.StateKeys.newBuilder().addAllKeys(userStateStore.keys()))
+            .build();
+      }
+      return expected;
+    }
+  }
+
   static final class ClearStateEntry extends JournalEntry<ClearStateEntryMessage> {
 
     static final ClearStateEntry INSTANCE = new ClearStateEntry();

sdk-core/src/main/java/dev/restate/sdk/core/ExecutorSwitchingSyscalls.java

@@ -17,6 +17,7 @@
 import dev.restate.sdk.common.syscalls.SyscallCallback;
 import io.grpc.MethodDescriptor;
 import java.time.Duration;
+import java.util.Collection;
 import java.util.Map;
 import java.util.concurrent.Executor;
 
@@ -50,6 +51,11 @@ public void get(String name, SyscallCallback<Deferred<ByteString>> callback) {
     syscallsExecutor.execute(() -> syscalls.get(name, callback));
   }
 
+  @Override
+  public void getKeys(SyscallCallback<Deferred<Collection<String>>> callback) {
+    syscallsExecutor.execute(() -> syscalls.getKeys(callback));
+  }
+
   @Override
   public void clear(String name, SyscallCallback<Void> callback) {
     syscallsExecutor.execute(() -> syscalls.clear(name, callback));

sdk-core/src/main/java/dev/restate/sdk/core/MessageHeader.java

@@ -87,6 +87,14 @@ public static MessageHeader fromMessage(MessageLite msg) {
       return new MessageHeader(MessageType.ClearStateEntryMessage, 0, msg.getSerializedSize());
     } else if (msg instanceof Protocol.ClearAllStateEntryMessage) {
       return new MessageHeader(MessageType.ClearAllStateEntryMessage, 0, msg.getSerializedSize());
+    } else if (msg instanceof Protocol.GetStateKeysEntryMessage) {
+      return new MessageHeader(
+          MessageType.GetStateKeysEntryMessage,
+          ((Protocol.GetStateKeysEntryMessage) msg).getResultCase()
+                  != Protocol.GetStateKeysEntryMessage.ResultCase.RESULT_NOT_SET
+              ? DONE_FLAG
+              : 0,
+          msg.getSerializedSize());
     } else if (msg instanceof Protocol.SleepEntryMessage) {
       return new MessageHeader(
           MessageType.SleepEntryMessage,

sdk-core/src/main/java/dev/restate/sdk/core/MessageType.java

@@ -30,6 +30,7 @@ public enum MessageType {
   SetStateEntryMessage,
   ClearStateEntryMessage,
   ClearAllStateEntryMessage,
+  GetStateKeysEntryMessage,
 
   // Syscalls
   SleepEntryMessage,
@@ -54,6 +55,7 @@ public enum MessageType {
   public static final short SET_STATE_ENTRY_MESSAGE_TYPE = 0x0801;
   public static final short CLEAR_STATE_ENTRY_MESSAGE_TYPE = 0x0802;
   public static final short CLEAR_ALL_STATE_ENTRY_MESSAGE_TYPE = 0x0803;
+  public static final short GET_STATE_KEYS_ENTRY_MESSAGE_TYPE = 0x0804;
   public static final short SLEEP_ENTRY_MESSAGE_TYPE = 0x0C00;
   public static final short INVOKE_ENTRY_MESSAGE_TYPE = 0x0C01;
   public static final short BACKGROUND_INVOKE_ENTRY_MESSAGE_TYPE = 0x0C02;
@@ -88,6 +90,8 @@ public Parser<? extends MessageLite> messageParser() {
         return Protocol.ClearStateEntryMessage.parser();
       case ClearAllStateEntryMessage:
         return Protocol.ClearAllStateEntryMessage.parser();
+      case GetStateKeysEntryMessage:
+        return Protocol.GetStateKeysEntryMessage.parser();
       case SleepEntryMessage:
         return Protocol.SleepEntryMessage.parser();
       case InvokeEntryMessage:
@@ -132,6 +136,8 @@ public short encode() {
         return CLEAR_STATE_ENTRY_MESSAGE_TYPE;
       case ClearAllStateEntryMessage:
         return CLEAR_ALL_STATE_ENTRY_MESSAGE_TYPE;
+      case GetStateKeysEntryMessage:
+        return GET_STATE_KEYS_ENTRY_MESSAGE_TYPE;
       case SleepEntryMessage:
         return SLEEP_ENTRY_MESSAGE_TYPE;
       case InvokeEntryMessage:
@@ -176,6 +182,8 @@ public static MessageType decode(short value) throws ProtocolException {
         return ClearStateEntryMessage;
       case CLEAR_ALL_STATE_ENTRY_MESSAGE_TYPE:
         return ClearAllStateEntryMessage;
+      case GET_STATE_KEYS_ENTRY_MESSAGE_TYPE:
+        return GetStateKeysEntryMessage;
       case SLEEP_ENTRY_MESSAGE_TYPE:
         return SleepEntryMessage;
       case INVOKE_ENTRY_MESSAGE_TYPE:

sdk-core/src/main/java/dev/restate/sdk/core/SyscallsImpl.java

@@ -23,9 +23,7 @@
 import java.nio.ByteBuffer;
 import java.time.Duration;
 import java.time.Instant;
-import java.util.AbstractMap;
-import java.util.Base64;
-import java.util.Map;
+import java.util.*;
 import java.util.function.Function;
 import javax.annotation.Nullable;
 import org.apache.logging.log4j.LogManager;
@@ -104,6 +102,19 @@ public void get(String name, SyscallCallback<Deferred<ByteString>> callback) {
         callback);
   }
 
+  @Override
+  public void getKeys(SyscallCallback<Deferred<Collection<String>>> callback) {
+    wrapAndPropagateExceptions(
+        () -> {
+          LOG.trace("get keys");
+          this.stateMachine.processCompletableJournalEntry(
+              Protocol.GetStateKeysEntryMessage.newBuilder().build(),
+              GetStateKeysEntry.INSTANCE,
+              callback);
+        },
+        callback);
+  }
+
   @Override
   public void clear(String name, SyscallCallback<Void> callback) {
     wrapAndPropagateExceptions(

sdk-core/src/main/java/dev/restate/sdk/core/UserStateStore.java

@@ -11,6 +11,7 @@
 import com.google.protobuf.ByteString;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Set;
 import java.util.stream.Collectors;
 
 final class UserStateStore {
@@ -68,4 +69,12 @@ public void clearAll() {
     this.map.clear();
     this.isPartial = false;
   }
+
+  public boolean isComplete() {
+    return !isPartial;
+  }
+
+  public Set<ByteString> keys() {
+    return this.map.keySet();
+  }
 }

sdk-core/src/main/java/dev/restate/sdk/core/Util.java

@@ -90,7 +90,7 @@ static boolean isTerminalException(Throwable throwable) {
 
   static void assertIsEntry(MessageLite msg) {
     if (!isEntry(msg)) {
-      throw new IllegalStateException("Expected input to be entry");
+      throw new IllegalStateException("Expected input to be entry: " + msg);
     }
   }
 
@@ -110,6 +110,7 @@ static boolean isEntry(MessageLite msg) {
     return msg instanceof Protocol.PollInputStreamEntryMessage
         || msg instanceof Protocol.OutputStreamEntryMessage
         || msg instanceof Protocol.GetStateEntryMessage
+        || msg instanceof Protocol.GetStateKeysEntryMessage
         || msg instanceof Protocol.SetStateEntryMessage
         || msg instanceof Protocol.ClearStateEntryMessage
         || msg instanceof Protocol.ClearAllStateEntryMessage