Fix case where we pass ServiceDefinition to the builder already (#402)

restatedev · Nov 18, 2024 · fa728dd · fa728dd
1 parent c86ae92
commit fa728dd
Showing 1 changed file with 18 additions and 0 deletions.
diff --git a/sdk-core/src/main/java/dev/restate/sdk/core/RestateEndpoint.java b/sdk-core/src/main/java/dev/restate/sdk/core/RestateEndpoint.java
@@ -238,6 +238,24 @@ private ServiceDefinitionFactoryDiscovery() {
   @SuppressWarnings("unchecked")
   public static ServiceDefinitionFactory<Object, Object> discoverServiceDefinitionFactory(
       Object service) {
+    if (service instanceof ServiceDefinitionFactory<?, ?>) {
+      // We got this already
+      return (ServiceDefinitionFactory<Object, Object>) service;
+    }
+    if (service instanceof ServiceDefinition<?>) {
+      // We got this already
+      return new ServiceDefinitionFactory<>() {
+        @Override
+        public ServiceDefinition<Object> create(Object serviceObject) {
+          return (ServiceDefinition<Object>) serviceObject;
+        }
+
+        @Override
+        public boolean supports(Object serviceObject) {
+          return serviceObject == service;
+        }
+      };
+    }
     return Objects.requireNonNull(
         ServiceDefinitionFactorySingleton.INSTANCE.discoverFactory(service),
         () ->