Examples with CancellationFlag org.jboss.as.ejb3.component.interceptors.CancellationFlag used on opensource projects

Example 6 with CancellationFlag

use of org.jboss.as.ejb3.component.interceptors.CancellationFlag in project wildfly by wildfly.

the class AssociationImpl method receiveInvocationRequest.

@Override
public CancelHandle receiveInvocationRequest(@NotNull final InvocationRequest invocationRequest) {
    final EJBIdentifier ejbIdentifier = invocationRequest.getEJBIdentifier();
    final String appName = ejbIdentifier.getAppName();
    final String moduleName = ejbIdentifier.getModuleName();
    final String distinctName = ejbIdentifier.getDistinctName();
    final String beanName = ejbIdentifier.getBeanName();
    final EjbDeploymentInformation ejbDeploymentInformation = findEJB(appName, moduleName, distinctName, beanName);
    if (ejbDeploymentInformation == null) {
        invocationRequest.writeNoSuchEJB();
        return CancelHandle.NULL;
    }
    final ClassLoader classLoader = ejbDeploymentInformation.getDeploymentClassLoader();
    ClassLoader originalTccl = WildFlySecurityManager.getCurrentContextClassLoaderPrivileged();
    WildFlySecurityManager.setCurrentContextClassLoaderPrivileged(classLoader);
    final InvocationRequest.Resolved requestContent;
    try {
        requestContent = invocationRequest.getRequestContent(classLoader);
    } catch (IOException | ClassNotFoundException e) {
        invocationRequest.writeException(new EJBException(e));
        return CancelHandle.NULL;
    } finally {
        WildFlySecurityManager.setCurrentContextClassLoaderPrivileged(originalTccl);
    }
    final Map<String, Object> attachments = requestContent.getAttachments();
    final EJBLocator<?> ejbLocator = requestContent.getEJBLocator();
    final String viewClassName = ejbLocator.getViewType().getName();
    if (!ejbDeploymentInformation.isRemoteView(viewClassName)) {
        invocationRequest.writeWrongViewType();
        return CancelHandle.NULL;
    }
    final ComponentView componentView = ejbDeploymentInformation.getView(viewClassName);
    final Method invokedMethod = findMethod(componentView, invocationRequest.getMethodLocator());
    if (invokedMethod == null) {
        invocationRequest.writeNoSuchMethod();
        return CancelHandle.NULL;
    }
    final Component component = componentView.getComponent();
    try {
        component.waitForComponentStart();
    } catch (RuntimeException e) {
        invocationRequest.writeException(new EJBException(e));
        return CancelHandle.NULL;
    }
    final EJBLocator<?> actualLocator;
    if (component instanceof StatefulSessionComponent) {
        if (ejbLocator.isStateless()) {
            final SessionID sessionID = ((StatefulSessionComponent) component).createSessionRemote();
            try {
                invocationRequest.convertToStateful(sessionID);
            } catch (IllegalArgumentException e) {
                // cannot convert (old protocol)
                invocationRequest.writeNotStateful();
                return CancelHandle.NULL;
            }
            actualLocator = ejbLocator.withSession(sessionID);
        } else {
            actualLocator = ejbLocator;
        }
    } else {
        if (ejbLocator.isStateful()) {
            invocationRequest.writeNotStateful();
            return CancelHandle.NULL;
        } else {
            actualLocator = ejbLocator;
        }
    }
    final boolean isAsync = componentView.isAsynchronous(invokedMethod);
    final boolean oneWay = isAsync && invokedMethod.getReturnType() == void.class;
    if (oneWay) {
        // send immediate response
        updateAffinities(invocationRequest, attachments, ejbLocator, componentView);
        requestContent.writeInvocationResult(null);
    } else if (isAsync) {
        invocationRequest.writeProceedAsync();
    }
    final CancellationFlag cancellationFlag = new CancellationFlag();
    Runnable runnable = () -> {
        if (!cancellationFlag.runIfNotCancelled()) {
            if (!oneWay)
                invocationRequest.writeCancelResponse();
            return;
        }
        // invoke the method
        final Object result;
        try {
            final Map<String, Object> contextDataHolder = new HashMap<>();
            result = invokeMethod(componentView, invokedMethod, invocationRequest, requestContent, cancellationFlag, actualLocator, contextDataHolder);
            attachments.putAll(contextDataHolder);
        } catch (EJBComponentUnavailableException ex) {
            // if the Jakarta Enterprise Beans are shutting down when the invocation was done, then it's as good as the Jakarta Enterprise Beans not being available. The client has to know about this as
            // a "no such EJB" failure so that it can retry the invocation on a different node if possible.
            EjbLogger.EJB3_INVOCATION_LOGGER.debugf("Cannot handle method invocation: %s on bean: %s due to Jakarta Enterprise Beans component unavailability exception. Returning a no such Jakarta Enterprise Beans available message back to client", invokedMethod, beanName);
            if (!oneWay)
                invocationRequest.writeNoSuchEJB();
            return;
        } catch (ComponentIsStoppedException ex) {
            EjbLogger.EJB3_INVOCATION_LOGGER.debugf("Cannot handle method invocation: %s on bean: %s due to Jakarta Enterprise Beans component stopped exception. Returning a no such Jakarta Enterprise Beans available message back to client", invokedMethod, beanName);
            if (!oneWay)
                invocationRequest.writeNoSuchEJB();
            return;
        // TODO should we write a specifc response with a specific protocol letting client know that server is suspending?
        } catch (CancellationException ex) {
            if (!oneWay)
                invocationRequest.writeCancelResponse();
            return;
        } catch (Exception exception) {
            if (oneWay)
                return;
            // write out the failure
            final Exception exceptionToWrite;
            final Throwable cause = exception.getCause();
            if (componentView.getComponent() instanceof StatefulSessionComponent && exception instanceof EJBException && cause != null) {
                if (!(componentView.getComponent().isRemotable(cause))) {
                    // Avoid serializing the cause of the exception in case it is not remotable
                    // Client might not be able to deserialize and throw ClassNotFoundException
                    exceptionToWrite = new EJBException(exception.getLocalizedMessage());
                } else {
                    exceptionToWrite = exception;
                }
            } else {
                exceptionToWrite = exception;
            }
            invocationRequest.writeException(exceptionToWrite);
            return;
        }
        // invocation was successful
        if (!oneWay)
            try {
                updateAffinities(invocationRequest, attachments, actualLocator, componentView);
                requestContent.writeInvocationResult(result);
            } catch (Throwable ioe) {
                EjbLogger.REMOTE_LOGGER.couldNotWriteMethodInvocation(ioe, invokedMethod, beanName, appName, moduleName, distinctName);
            }
    };
    // invoke the method and write out the response, possibly on a separate thread
    execute(invocationRequest, runnable, isAsync, false);
    return cancellationFlag::cancel;
}