added interface based calling pattern and entity re-entrancy

Author: bachuv

client/src/main/java/com/microsoft/durabletask/EntityProxy.java

@@ -0,0 +1,182 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.microsoft.durabletask;
+
+import javax.annotation.Nonnull;
+import java.lang.reflect.InvocationHandler;
+import java.lang.reflect.Method;
+import java.lang.reflect.ParameterizedType;
+import java.lang.reflect.Proxy;
+import java.lang.reflect.Type;
+
+/**
+ * Creates type-safe proxies for interacting with durable entities from orchestrations.
+ * <p>
+ * A typed entity proxy is a JDK dynamic proxy that implements a user-defined interface.
+ * Method calls on the proxy are translated into entity operations:
+ * <ul>
+ *   <li>{@code void} methods become fire-and-forget signals via
+ *       {@link TaskOrchestrationContext#signalEntity}</li>
+ *   <li>Methods returning {@link Task}{@code <V>} become two-way calls via
+ *       {@link TaskOrchestrationContext#callEntity}</li>
+ * </ul>
+ * <p>
+ * The method name is used as the entity operation name (case-insensitive matching on the
+ * entity side). Methods must accept 0 or 1 parameters; the single parameter is passed as
+ * the operation input.
+ *
+ * <p>Example:
+ * <pre>{@code
+ * // Define entity operations as an interface
+ * public interface ICounter {
+ *     void add(int amount);            // fire-and-forget signal
+ *     void reset();                    // fire-and-forget signal
+ *     Task<Integer> get();             // two-way call returning a result
+ * }
+ *
+ * // Use in an orchestration
+ * ICounter counter = ctx.createEntityProxy(entityId, ICounter.class);
+ * counter.add(5);
+ * counter.reset();
+ * int value = counter.get().await();
+ * }</pre>
+ *
+ * @see TaskOrchestrationContext#createEntityProxy(EntityInstanceId, Class)
+ * @see TaskOrchestrationEntityFeature#createProxy(EntityInstanceId, Class)
+ */
+public final class EntityProxy {
+
+    private EntityProxy() {
+        // Utility class — not instantiable
+    }
+
+    /**
+     * Creates a typed entity proxy for the given entity instance.
+     *
+     * @param ctx            the orchestration context (used to send signals and calls)
+     * @param entityId       the target entity's instance ID
+     * @param proxyInterface the interface whose methods map to entity operations
+     * @param <T>            the proxy interface type
+     * @return a proxy instance that implements {@code proxyInterface}
+     * @throws IllegalArgumentException if {@code proxyInterface} is not an interface
+     */
+    @SuppressWarnings("unchecked")
+    public static <T> T create(
+            @Nonnull TaskOrchestrationContext ctx,
+            @Nonnull EntityInstanceId entityId,
+            @Nonnull Class<T> proxyInterface) {
+        if (ctx == null) {
+            throw new IllegalArgumentException("ctx must not be null");
+        }
+        if (entityId == null) {
+            throw new IllegalArgumentException("entityId must not be null");
+        }
+        if (proxyInterface == null) {
+            throw new IllegalArgumentException("proxyInterface must not be null");
+        }
+        if (!proxyInterface.isInterface()) {
+            throw new IllegalArgumentException(
+                    "proxyInterface must be an interface, got: " + proxyInterface.getName());
+        }
+
+        return (T) Proxy.newProxyInstance(
+                proxyInterface.getClassLoader(),
+                new Class<?>[]{ proxyInterface },
+                new EntityInvocationHandler(ctx, entityId));
+    }
+
+    /**
+     * Invocation handler that translates interface method calls into entity operations.
+     */
+    private static final class EntityInvocationHandler implements InvocationHandler {
+        private final TaskOrchestrationContext ctx;
+        private final EntityInstanceId entityId;
+
+        EntityInvocationHandler(TaskOrchestrationContext ctx, EntityInstanceId entityId) {
+            this.ctx = ctx;
+            this.entityId = entityId;
+        }
+
+        @Override
+        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
+            // Handle java.lang.Object methods
+            if (method.getDeclaringClass() == Object.class) {
+                switch (method.getName()) {
+                    case "toString":
+                        return "EntityProxy[" + entityId + "]";
+                    case "hashCode":
+                        return entityId.hashCode();
+                    case "equals":
+                        if (args[0] == proxy) {
+                            return true;
+                        }
+                        if (args[0] == null || !Proxy.isProxyClass(args[0].getClass())) {
+                            return false;
+                        }
+                        InvocationHandler otherHandler = Proxy.getInvocationHandler(args[0]);
+                        if (otherHandler instanceof EntityInvocationHandler) {
+                            return entityId.equals(((EntityInvocationHandler) otherHandler).entityId);
+                        }
+                        return false;
+                    default:
+                        return method.invoke(this, args);
+                }
+            }
+
+            String operationName = method.getName();
+            Object input = (args != null && args.length == 1) ? args[0] : null;
+
+            if (args != null && args.length > 1) {
+                throw new UnsupportedOperationException(
+                        "Entity proxy methods must have 0 or 1 parameters. " +
+                        "Method '" + operationName + "' has " + args.length + " parameters. " +
+                        "Use a single wrapper object to pass multiple values.");
+            }
+
+            Class<?> returnType = method.getReturnType();
+
+            if (returnType == void.class) {
+                // Fire-and-forget signal
+                ctx.signalEntity(entityId, operationName, input);
+                return null;
+            } else if (Task.class.isAssignableFrom(returnType)) {
+                // Two-way entity call — extract the Task<V> type parameter
+                Class<?> resultType = extractTaskTypeParameter(method);
+                return ctx.callEntity(entityId, operationName, input, resultType);
+            } else {
+                throw new UnsupportedOperationException(
+                        "Entity proxy methods must return void (for signals) or Task<V> (for calls). " +
+                        "Method '" + operationName + "' returns " + returnType.getName() + ".");
+            }
+        }
+
+        /**
+         * Extracts the generic type parameter from a method returning {@code Task<V>}.
+         * Falls back to {@code Void.class} if the type cannot be determined.
+         */
+        private static Class<?> extractTaskTypeParameter(Method method) {
+            Type genericReturnType = method.getGenericReturnType();
+            if (genericReturnType instanceof ParameterizedType) {
+                ParameterizedType pt = (ParameterizedType) genericReturnType;
+                Type[] typeArgs = pt.getActualTypeArguments();
+                if (typeArgs.length > 0) {
+                    return getRawClass(typeArgs[0]);
+                }
+            }
+            return Void.class;
+        }
+
+        /**
+         * Resolves a {@link Type} to its raw {@link Class}, handling parameterized types
+         * and wildcard types.
+         */
+        private static Class<?> getRawClass(Type type) {
+            if (type instanceof Class) {
+                return (Class<?>) type;
+            } else if (type instanceof ParameterizedType) {
+                return getRawClass(((ParameterizedType) type).getRawType());
+            }
+            return Object.class;
+        }
+    }
+}

client/src/main/java/com/microsoft/durabletask/TaskEntity.java

@@ -191,6 +191,136 @@ private Object handleImplicitOperations(TaskEntityOperation operation) {
                 operation.getName() + "'.");
     }
 
+    // region Re-entrant self-dispatch
+
+    /**
+     * Dispatches an operation to this entity instance synchronously (re-entrant self-call).
+     * <p>
+     * This allows one entity operation to invoke another operation on the same entity,
+     * reusing the same state and context. The dispatched operation executes inline — state
+     * mutations are visible immediately to the caller.
+     * <p>
+     * The operation is resolved using the same case-insensitive reflection dispatch as external
+     * operations: entity class methods are tried first, then state object methods (if
+     * {@linkplain #setAllowStateDispatch state dispatch} is enabled), then implicit operations.
+     *
+     * <p>Example:
+     * <pre>{@code
+     * public class BankAccount extends TaskEntity<BankAccountState> {
+     *     public void deposit(int amount) {
+     *         this.state.balance += amount;
+     *     }
+     *
+     *     public void depositWithBonus(int amount) {
+     *         dispatch("deposit", amount);          // re-entrant call
+     *         dispatch("deposit", amount / 10);     // 10% bonus
+     *     }
+     * }
+     * }</pre>
+     *
+     * @param operationName the name of the operation to dispatch (case-insensitive)
+     * @return the operation result, or {@code null} for void operations
+     * @throws IllegalStateException        if called outside of entity execution
+     * @throws UnsupportedOperationException if no matching operation is found
+     */
+    protected Object dispatch(String operationName) {
+        return dispatch(operationName, null);
+    }
+
+    /**
+     * Dispatches an operation with input to this entity instance synchronously (re-entrant self-call).
+     *
+     * @param operationName the name of the operation to dispatch (case-insensitive)
+     * @param input         the input value to pass to the operation, or {@code null}
+     * @return the operation result, or {@code null} for void operations
+     * @throws IllegalStateException        if called outside of entity execution
+     * @throws UnsupportedOperationException if no matching operation is found
+     * @see #dispatch(String)
+     */
+    protected Object dispatch(String operationName, Object input) {
+        if (this.context == null) {
+            throw new IllegalStateException(
+                    "dispatch() can only be called during entity operation execution.");
+        }
+
+        Method method = findMethod(this.getClass(), operationName);
+        Object target = this;
+
+        if (method == null && this.allowStateDispatch && this.state != null) {
+            method = findMethod(this.state.getClass(), operationName);
+            target = this.state;
+        }
+
+        if (method == null) {
+            if ("delete".equalsIgnoreCase(operationName)) {
+                this.state = null;
+                return null;
+            }
+            throw new UnsupportedOperationException(
+                    "Entity '" + this.getClass().getSimpleName() +
+                    "' does not support operation '" + operationName + "'.");
+        }
+
+        return invokeMethodDirect(method, target, input, this.context);
+    }
+
+    /**
+     * Dispatches an operation with input and a typed return value (re-entrant self-call).
+     *
+     * @param operationName the name of the operation to dispatch (case-insensitive)
+     * @param input         the input value to pass to the operation, or {@code null}
+     * @param returnType    the expected return type
+     * @param <V>           the return type
+     * @return the operation result cast to {@code V}
+     * @throws IllegalStateException        if called outside of entity execution
+     * @throws UnsupportedOperationException if no matching operation is found
+     * @throws ClassCastException           if the result cannot be cast to {@code returnType}
+     * @see #dispatch(String, Object)
+     */
+    protected <V> V dispatch(String operationName, Object input, Class<V> returnType) {
+        return returnType.cast(dispatch(operationName, input));
+    }
+
+    /**
+     * Invokes a method with a direct (already-deserialized) input value, without requiring
+     * a {@link TaskEntityOperation}. Used by {@link #dispatch} for re-entrant self-calls.
+     */
+    private static Object invokeMethodDirect(
+            Method method, Object target, Object input, TaskEntityContext context) {
+        Class<?>[] paramTypes = method.getParameterTypes();
+        Object[] args = new Object[paramTypes.length];
+
+        for (int i = 0; i < paramTypes.length; i++) {
+            if (TaskEntityContext.class.isAssignableFrom(paramTypes[i])) {
+                args[i] = context;
+            } else if (TaskEntityOperation.class.isAssignableFrom(paramTypes[i])) {
+                throw new UnsupportedOperationException(
+                        "Cannot dispatch to method '" + method.getName() +
+                        "' that accepts TaskEntityOperation. Use a simpler signature for " +
+                        "operations that support re-entrant dispatch.");
+            } else {
+                args[i] = input;
+            }
+        }
+
+        try {
+            return method.invoke(target, args);
+        } catch (InvocationTargetException e) {
+            Throwable cause = e.getTargetException();
+            if (cause instanceof RuntimeException) {
+                throw (RuntimeException) cause;
+            }
+            if (cause instanceof Error) {
+                throw (Error) cause;
+            }
+            throw new RuntimeException(cause);
+        } catch (IllegalAccessException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    // endregion
+
     /**
      * Finds a public method on the target class matching the operation name (case-insensitive).
      * Methods inherited from {@code Object} and from {@code TaskEntity} itself are excluded.

client/src/main/java/com/microsoft/durabletask/TaskOrchestrationContext.java

@@ -559,6 +559,24 @@ default <V> Task<V> waitForExternalEvent(String name, Class<V> dataType) {
         }
     }
 
+    /**
+     * Creates a typed entity proxy that maps interface method calls to entity operations.
+     * <p>
+     * This is a convenience method equivalent to calling
+     * {@code EntityProxy.create(this, entityId, proxyInterface)}.
+     *
+     * @param entityId       the target entity's instance ID
+     * @param proxyInterface the interface whose methods map to entity operations;
+     *                       {@code void} methods become signals, {@code Task<V>} methods become calls
+     * @param <T>            the proxy interface type
+     * @return a proxy instance that implements {@code proxyInterface}
+     * @throws IllegalArgumentException if {@code proxyInterface} is not an interface
+     * @see EntityProxy#create(TaskOrchestrationContext, EntityInstanceId, Class)
+     */
+    default <T> T createEntityProxy(@Nonnull EntityInstanceId entityId, @Nonnull Class<T> proxyInterface) {
+        return EntityProxy.create(this, entityId, proxyInterface);
+    }
+
     /**
      * Gets the durable entity feature for this orchestration context.
      * <p>

client/src/main/java/com/microsoft/durabletask/TaskOrchestrationEntityFeature.java

@@ -130,6 +130,22 @@ public void signalEntity(
      * @return the list of currently locked entities, or an empty list
      */
     public abstract List<EntityInstanceId> getLockedEntities();
+
+    /**
+     * Creates a typed entity proxy that maps interface method calls to entity operations.
+     * <p>
+     * This is the feature-based equivalent of
+     * {@link TaskOrchestrationContext#createEntityProxy(EntityInstanceId, Class)}.
+     *
+     * @param entityId       the target entity's instance ID
+     * @param proxyInterface the interface whose methods map to entity operations;
+     *                       {@code void} methods become signals, {@code Task<V>} methods become calls
+     * @param <T>            the proxy interface type
+     * @return a proxy instance that implements {@code proxyInterface}
+     * @throws IllegalArgumentException if {@code proxyInterface} is not an interface
+     * @see EntityProxy#create(TaskOrchestrationContext, EntityInstanceId, Class)
+     */
+    public abstract <T> T createProxy(@Nonnull EntityInstanceId entityId, @Nonnull Class<T> proxyInterface);
 }
 
 final class ContextBackedTaskOrchestrationEntityFeature extends TaskOrchestrationEntityFeature {
@@ -186,4 +202,9 @@ public boolean isInCriticalSection() {
     public List<EntityInstanceId> getLockedEntities() {
         return this.context.getLockedEntities();
     }
+
+    @Override
+    public <T> T createProxy(@Nonnull EntityInstanceId entityId, @Nonnull Class<T> proxyInterface) {
+        return EntityProxy.create(this.context, entityId, proxyInterface);
+    }
 }