Merge pull request quarkusio#40183 from mkouba/websockets-next-optimi…

…ze-01 WebSockets Next: avoid unnecessary bean lookups
zakkak · Apr 22, 2024 · 63170de · 63170de
2 parents 923e83d + 1acdb64
commit 63170de
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 10 deletions.
diff --git a/...loyment/src/main/java/io/quarkus/websockets/next/deployment/WebSocketServerProcessor.java b/...loyment/src/main/java/io/quarkus/websockets/next/deployment/WebSocketServerProcessor.java
@@ -413,7 +413,7 @@ private void validateOnClose(Callback callback) {
      *     }
      *
      *     public Uni doOnTextMessage(String message) {
-     *         Uni uni = ((Echo) super.beanInstance("MTd91f3oxHtG8gnznR7XcZBCLdE")).echo((String) message);
+     *         Uni uni = ((Echo) super.beanInstance().echo((String) message);
      *         if (uni != null) {
      *             // The lambda is implemented as a generated function: Echo_WebSocketEndpoint$$function$$1
      *             return uni.chain(m -> sendText(m, false));
@@ -423,7 +423,7 @@ private void validateOnClose(Callback callback) {
      *     }
      *
      *     public Uni doOnTextMessage(Object message) {
-     *         Object bean = super.beanInstance("egBJQ7_QAFkQlYXSTKE0XlN3wow");
+     *         Object bean = super.beanInstance();
      *         try {
      *             String ret = ((EchoEndpoint) bean).echo((String) message);
      *             return ret != null ? super.sendText(ret, false) : Uni.createFrom().voidItem();
@@ -445,6 +445,10 @@ private void validateOnClose(Callback callback) {
      *     public WebSocketEndpoint.ExecutionModel onTextMessageExecutionModel() {
      *         return ExecutionModel.EVENT_LOOP;
      *     }
+     *
+     *     public String beanIdentifier() {
+     *        return "egBJQ7_QAFkQlYXSTKE0XlN3wow";
+     *     }
      * }
      * </pre>
      *
@@ -485,13 +489,15 @@ private String generateEndpoint(WebSocketEndpointBuildItem endpoint,
         MethodCreator executionMode = endpointCreator.getMethodCreator("executionMode", WebSocket.ExecutionMode.class);
         executionMode.returnValue(executionMode.load(endpoint.executionMode));
 
+        MethodCreator beanIdentifier = endpointCreator.getMethodCreator("beanIdentifier", String.class);
+        beanIdentifier.returnValue(beanIdentifier.load(endpoint.bean.getIdentifier()));
+
         if (endpoint.onOpen != null) {
             Callback callback = endpoint.onOpen;
             MethodCreator doOnOpen = endpointCreator.getMethodCreator("doOnOpen", Uni.class, Object.class);
-            // Foo foo = beanInstance("foo");
+            // Foo foo = beanInstance();
             ResultHandle beanInstance = doOnOpen.invokeVirtualMethod(
-                    MethodDescriptor.ofMethod(WebSocketEndpointBase.class, "beanInstance", Object.class, String.class),
-                    doOnOpen.getThis(), doOnOpen.load(endpoint.bean.getIdentifier()));
+                    MethodDescriptor.ofMethod(WebSocketEndpointBase.class, "beanInstance", Object.class), doOnOpen.getThis());
             // Call the business method
             TryBlock tryBlock = onErrorTryBlock(doOnOpen, doOnOpen.getThis());
             ResultHandle[] args = callback.generateArguments(tryBlock.getThis(), tryBlock, transformedAnnotations, index);
@@ -515,8 +521,7 @@ private String generateEndpoint(WebSocketEndpointBuildItem endpoint,
             MethodCreator doOnClose = endpointCreator.getMethodCreator("doOnClose", Uni.class, Object.class);
             // Foo foo = beanInstance("foo");
             ResultHandle beanInstance = doOnClose.invokeVirtualMethod(
-                    MethodDescriptor.ofMethod(WebSocketEndpointBase.class, "beanInstance", Object.class, String.class),
-                    doOnClose.getThis(), doOnClose.load(endpoint.bean.getIdentifier()));
+                    MethodDescriptor.ofMethod(WebSocketEndpointBase.class, "beanInstance", Object.class), doOnClose.getThis());
             // Call the business method
             TryBlock tryBlock = onErrorTryBlock(doOnClose, doOnClose.getThis());
             ResultHandle[] args = callback.generateArguments(tryBlock.getThis(), tryBlock, transformedAnnotations, index);
@@ -663,10 +668,9 @@ private void generateOnMessage(ClassCreator endpointCreator, WebSocketEndpointBu
                 methodParameterType);
 
         TryBlock tryBlock = onErrorTryBlock(doOnMessage, doOnMessage.getThis());
-        // Foo foo = beanInstance("foo");
+        // Foo foo = beanInstance();
         ResultHandle beanInstance = tryBlock.invokeVirtualMethod(
-                MethodDescriptor.ofMethod(WebSocketEndpointBase.class, "beanInstance", Object.class, String.class),
-                tryBlock.getThis(), tryBlock.load(endpoint.bean.getIdentifier()));
+                MethodDescriptor.ofMethod(WebSocketEndpointBase.class, "beanInstance", Object.class), tryBlock.getThis());
         ResultHandle[] args = callback.generateArguments(tryBlock.getThis(), tryBlock, transformedAnnotations, index);
         // Call the business method
         ResultHandle ret = tryBlock.invokeVirtualMethod(MethodDescriptor.of(callback.method), beanInstance,

diff --git a/...xt/server/runtime/src/main/java/io/quarkus/websockets/next/runtime/WebSocketEndpoint.java b/...xt/server/runtime/src/main/java/io/quarkus/websockets/next/runtime/WebSocketEndpoint.java
@@ -79,6 +79,13 @@ default ExecutionModel onCloseExecutionModel() {
 
     Uni<Void> doOnError(Throwable t);
 
+    /**
+     *
+     * @return the identifier of the bean with callbacks
+     * @see io.quarkus.arc.InjectableBean#getIdentifier()
+     */
+    String beanIdentifier();
+
     enum ExecutionModel {
         WORKER_THREAD,
         VIRTUAL_THREAD,

diff --git a/...erver/runtime/src/main/java/io/quarkus/websockets/next/runtime/WebSocketEndpointBase.java b/...erver/runtime/src/main/java/io/quarkus/websockets/next/runtime/WebSocketEndpointBase.java
@@ -5,10 +5,14 @@
 import java.util.function.Consumer;
 import java.util.function.Function;
 
+import jakarta.enterprise.context.ApplicationScoped;
+import jakarta.inject.Singleton;
+
 import org.jboss.logging.Logger;
 
 import io.quarkus.arc.Arc;
 import io.quarkus.arc.ArcContainer;
+import io.quarkus.arc.InjectableBean;
 import io.quarkus.arc.InjectableContext.ContextState;
 import io.quarkus.virtual.threads.VirtualThreadsRecorder;
 import io.quarkus.websockets.next.WebSocket.ExecutionMode;
@@ -43,6 +47,9 @@ public abstract class WebSocketEndpointBase implements WebSocketEndpoint {
 
     private final ContextSupport contextSupport;
 
+    private final InjectableBean<?> bean;
+    private final Object beanInstance;
+
     public WebSocketEndpointBase(WebSocketConnection connection, Codecs codecs,
             WebSocketsRuntimeConfig config, ContextSupport contextSupport) {
         this.connection = connection;
@@ -51,6 +58,16 @@ public WebSocketEndpointBase(WebSocketConnection connection, Codecs codecs,
         this.config = config;
         this.container = Arc.container();
         this.contextSupport = contextSupport;
+        InjectableBean<?> bean = container.bean(beanIdentifier());
+        if (bean.getScope().equals(ApplicationScoped.class)
+                || bean.getScope().equals(Singleton.class)) {
+            // For certain scopes, we can optimize and obtain the contextual reference immediately
+            this.bean = null;
+            this.beanInstance = container.instance(bean).get();
+        } else {
+            this.bean = bean;
+            this.beanInstance = null;
+        }
     }
 
     @Override
@@ -238,6 +255,10 @@ public void handle(Void event) {
         return UniHelper.toUni(promise.future());
     }
 
+    public Object beanInstance() {
+        return beanInstance != null ? beanInstance : container.instance(bean).get();
+    }
+
     public Object beanInstance(String identifier) {
         return container.instance(container.bean(identifier)).get();
     }