fix: fix inspection and injection of Vaadin scoped beans

kubernetes-kit-starter/src/main/java/com/vaadin/kubernetes/starter/sessiontracker/PessimisticSerializationRequiredException.java

@@ -0,0 +1,43 @@
+/*-
+ * Copyright (C) 2022 Vaadin Ltd
+ *
+ * This program is available under Vaadin Commercial License and Service Terms.
+ *
+ *
+ * See <https://vaadin.com/commercial-license-and-service-terms> for the full
+ * license.
+ */
+package com.vaadin.kubernetes.starter.sessiontracker;
+
+/**
+ * Exception raise during session serialization to indicate that VaadinSession
+ * lock is required to complete the operation.
+ */
+public class PessimisticSerializationRequiredException
+        extends RuntimeException {
+
+    /**
+     * Constructs a new exception with the specified detail message.
+     *
+     * @param message
+     *            the detail message. The detail message is saved for later
+     *            retrieval by the {@link #getMessage()} method.
+     */
+    public PessimisticSerializationRequiredException(String message) {
+        super(message);
+    }
+
+    /**
+     * Constructs a new exception with the specified detail message and cause.
+     *
+     * @param message
+     *            the detail message.
+     * @param cause
+     *            the cause. (A {@code null} value is permitted, and indicates
+     *            that the cause is nonexistent or unknown.)
+     */
+    public PessimisticSerializationRequiredException(String message,
+            Throwable cause) {
+        super(message, cause);
+    }
+}

kubernetes-kit-starter/src/main/java/com/vaadin/kubernetes/starter/sessiontracker/SessionSerializer.java

@@ -313,6 +313,11 @@ private void handleSessionSerialization(String sessionId,
                     return;
                 }
             }
+        } catch (PessimisticSerializationRequiredException e) {
+            getLogger().warn(
+                    "Optimistic serialization of session {} with distributed key {} cannot be completed "
+                            + " because VaadinSession lock is required. Switching to pessimistic locking.",
+                    sessionId, clusterKey, e);
         } catch (NotSerializableException e) {
             getLogger().error(
                     "Optimistic serialization of session {} with distributed key {} failed,"
@@ -410,7 +415,8 @@ private SessionInfo serializeOptimisticLocking(String sessionId,
             logSessionDebugInfo("Serialized session " + sessionId
                     + " with distributed key " + clusterKey, attributes);
             return info;
-        } catch (NotSerializableException e) {
+        } catch (NotSerializableException
+                | PessimisticSerializationRequiredException e) {
             throw e;
         } catch (Exception e) {
             getLogger().trace(

kubernetes-kit-starter/src/main/java/com/vaadin/kubernetes/starter/sessiontracker/serialization/SpringTransientHandler.java

@@ -13,9 +13,13 @@
 import java.lang.reflect.InaccessibleObjectException;
 import java.lang.reflect.Modifier;
 import java.util.ArrayList;
+import java.util.Collections;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -25,6 +29,11 @@
 import org.springframework.context.ApplicationContext;
 
 import com.vaadin.flow.internal.ReflectTools;
+import com.vaadin.flow.server.VaadinSession;
+import com.vaadin.flow.spring.annotation.RouteScope;
+import com.vaadin.flow.spring.annotation.UIScope;
+import com.vaadin.flow.spring.annotation.VaadinSessionScope;
+import com.vaadin.kubernetes.starter.sessiontracker.PessimisticSerializationRequiredException;
 
 /**
  * Spring specific implementation of {@link TransientHandler}, capable to
@@ -58,50 +67,137 @@ public void inject(Object obj, List<TransientDescriptor> transients) {
     }
 
     private void injectField(Object obj, TransientDescriptor descriptor) {
-        getLogger().debug("Injecting '{}' into transient field {} of type {}",
+        getLogger().debug(
+                "Injecting '{}' into transient field '{}' of type '{}'",
                 descriptor.getInstanceReference(), descriptor.getName(),
                 obj.getClass());
-        ReflectTools.setJavaFieldValue(obj, descriptor.getField(),
-                appCtx.getBean(descriptor.getInstanceReference()));
+        try {
+            ReflectTools.setJavaFieldValue(obj, descriptor.getField(),
+                    appCtx.getBean(descriptor.getInstanceReference()));
+        } catch (RuntimeException ex) {
+            getLogger().error(
+                    "Failed injecting '{}' into transient field '{}' of type '{}'",
+                    descriptor.getInstanceReference(), descriptor.getName(),
+                    obj.getClass());
+            throw ex;
+        }
     }
 
     public List<TransientDescriptor> inspect(Object target) {
-        return findTransientFields(target.getClass(), f -> true).stream()
-                .map(field -> detectBean(target, field))
-                .filter(Objects::nonNull).collect(Collectors.toList());
+        List<Injectable> injectables = findTransientFields(target.getClass(),
+                f -> true).stream().map(field -> detectBean(target, field))
+                .filter(Objects::nonNull).toList();
+        return createDescriptors(target, injectables);
     }
 
-    private TransientDescriptor detectBean(Object target, Field field) {
+    private Injectable detectBean(Object target, Field field) {
         Object value = getFieldValue(target, field);
-
         if (value != null) {
             Class<?> valueType = value.getClass();
             getLogger().trace(
                     "Inspecting field {} of class {} for injected beans",
                     field.getName(), target.getClass());
-            TransientDescriptor transientDescriptor = appCtx
-                    .getBeansOfType(valueType).entrySet().stream()
-                    .filter(e -> e.getValue() == value || matchesPrototype(
-                            e.getKey(), e.getValue(), valueType))
-                    .map(Map.Entry::getKey).findFirst()
-                    .map(beanName -> new TransientDescriptor(field, beanName))
-                    .orElse(null);
-            if (transientDescriptor != null) {
-                getLogger().trace("Bean {} found for field {} of class {}",
-                        transientDescriptor.getInstanceReference(),
-                        field.getName(), target.getClass());
-            } else {
-                getLogger().trace("No bean detected for field {} of class {}",
-                        field.getName(), target.getClass());
+            Set<String> beanNames = new LinkedHashSet<>(List
+                    .of(appCtx.getBeanNamesForType(valueType, true, false)));
+            List<String> vaadinScopedBeanNames = new ArrayList<>();
+            Collections.addAll(vaadinScopedBeanNames,
+                    appCtx.getBeanNamesForAnnotation(VaadinSessionScope.class));
+            Collections.addAll(vaadinScopedBeanNames,
+                    appCtx.getBeanNamesForAnnotation(UIScope.class));
+            Collections.addAll(vaadinScopedBeanNames,
+                    appCtx.getBeanNamesForAnnotation(RouteScope.class));
+
+            boolean vaadinScoped = beanNames.stream()
+                    .anyMatch(vaadinScopedBeanNames::contains);
+            if (vaadinScoped && VaadinSession.getCurrent() == null) {
+                getLogger().warn(
+                        "VaadinSession is not available when trying to inspect Vaadin scoped bean: {}."
+                                + "Transient fields might not be registered for deserialization.",
+                        beanNames);
+                beanNames.removeIf(vaadinScopedBeanNames::contains);
             }
-            return transientDescriptor;
+            return new Injectable(field, value, beanNames, vaadinScoped);
         }
         getLogger().trace(
                 "No bean detected for field {} of class {}, field value is null",
                 field.getName(), target.getClass());
         return null;
     }
 
+    private record Injectable(Field field, Object value, Set<String> beanNames,
+            boolean vaadinScoped) {
+    }
+
+    private TransientDescriptor createDescriptor(Object target,
+            Injectable injectable) {
+        Field field = injectable.field;
+        Object value = injectable.value;
+        Class<?> valueType = value.getClass();
+        TransientDescriptor transientDescriptor;
+        transientDescriptor = injectable.beanNames.stream()
+                .map(beanName -> Map.entry(beanName, appCtx.getBean(beanName)))
+                .filter(e -> e.getValue() == value || matchesPrototype(
+                        e.getKey(), e.getValue(), valueType))
+                .map(Map.Entry::getKey).findFirst()
+                .map(beanName -> new TransientDescriptor(field, beanName,
+                        injectable.vaadinScoped))
+                .orElse(null);
+        if (transientDescriptor != null) {
+            getLogger().trace("Bean {} found for field {} of class {}",
+                    transientDescriptor.getInstanceReference(), field.getName(),
+                    target.getClass());
+        } else {
+            getLogger().trace("No bean detected for field {} of class {}",
+                    field.getName(), target.getClass());
+        }
+        return transientDescriptor;
+    }
+
+    private List<TransientDescriptor> createDescriptors(Object target,
+            List<Injectable> injectables) {
+        boolean sessionLocked = false;
+        if (injectables.stream().anyMatch(Injectable::vaadinScoped)) {
+            // Bean has Vaadin scope, lookup needs VaadinSession lock
+            VaadinSession vaadinSession = VaadinSession.getCurrent();
+            if (vaadinSession != null) {
+                try {
+                    sessionLocked = vaadinSession.getLockInstance().tryLock(1,
+                            TimeUnit.SECONDS);
+                    if (!sessionLocked) {
+                        throw new PessimisticSerializationRequiredException(
+                                "Unable to acquire VaadinSession lock to lookup Vaadin scoped beans. "
+                                        + collectVaadinScopedCandidates(
+                                                injectables));
+                    }
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                    throw new PessimisticSerializationRequiredException(
+                            "Unable to acquire VaadinSession lock to lookup Vaadin scoped beans. "
+                                    + collectVaadinScopedCandidates(
+                                            injectables),
+                            e);
+                }
+            }
+        }
+        try {
+            return injectables.stream()
+                    .map(injectable -> createDescriptor(target, injectable))
+                    .filter(Objects::nonNull).toList();
+        } finally {
+            if (sessionLocked) {
+                VaadinSession.getCurrent().getLockInstance().unlock();
+            }
+        }
+    }
+
+    private String collectVaadinScopedCandidates(List<Injectable> injectables) {
+        return injectables.stream().filter(Injectable::vaadinScoped)
+                .map(injectable -> String.format(
+                        "[Field: %s, bean candidates: %s]",
+                        injectable.field.getName(), injectable.beanNames))
+                .collect(Collectors.joining(", "));
+    }
+
     private boolean matchesPrototype(String beanName, Object beanDefinition,
             Class<?> fieldValueType) {
         return appCtx.containsBeanDefinition(beanName)

kubernetes-kit-starter/src/main/java/com/vaadin/kubernetes/starter/sessiontracker/serialization/TransientAwareHolder.java

@@ -10,9 +10,21 @@
 package com.vaadin.kubernetes.starter.sessiontracker.serialization;
 
 import java.io.Serializable;
+import java.lang.reflect.Field;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import com.vaadin.flow.component.UI;
+import com.vaadin.flow.internal.CurrentInstance;
+import com.vaadin.flow.internal.ReflectTools;
+import com.vaadin.flow.server.VaadinSession;
 
 /**
  * A serializable class that holds information about an object to be
@@ -28,10 +40,21 @@ final class TransientAwareHolder implements Serializable {
 
     private final List<TransientDescriptor> transientDescriptors;
     private final Object source; // NOSONAR
+    private final UI ui;
+    private final VaadinSession session;
 
     TransientAwareHolder(Object source, List<TransientDescriptor> descriptors) {
         this.source = source;
         this.transientDescriptors = new ArrayList<>(descriptors);
+        if (descriptors.stream()
+                .anyMatch(TransientDescriptor::isVaadinScoped)) {
+            this.ui = UI.getCurrent();
+            this.session = ui != null ? ui.getSession()
+                    : VaadinSession.getCurrent();
+        } else {
+            this.ui = null;
+            this.session = null;
+        }
     }
 
     /**
@@ -53,4 +76,59 @@ Object source() {
         return source;
     }
 
+    /**
+     * Executes the given runnable making sure that Vaadin thread locals are
+     * set, when they are available.
+     * 
+     * @param runnable
+     *            the action to execute.
+     */
+    void inVaadinScope(Runnable runnable) {
+        Map<Class<?>, CurrentInstance> instanceMap = null;
+        if (ui != null) {
+            instanceMap = CurrentInstance.setCurrent(ui);
+        } else if (session != null) {
+            instanceMap = CurrentInstance.setCurrent(session);
+        }
+        Runnable cleaner = injectLock(session);
+        try {
+            runnable.run();
+        } finally {
+            if (instanceMap != null) {
+                CurrentInstance.restoreInstances(instanceMap);
+                cleaner.run();
+            }
+        }
+    }
+
+    // VaadinSession lock is usually set by calling
+    // VaadinSession.refreshTransients(WrappedSession,VaadinService), but during
+    // deserialization none of the required objects are available.
+    // This method injects a temporary lock instance into the provided
+    // VaadinSession and returns a runnable that will remove it when executed.
+    private static Runnable injectLock(VaadinSession session) {
+        if (session != null) {
+            try {
+                Field field = VaadinSession.class.getDeclaredField("lock");
+                Lock lock = new ReentrantLock();
+                lock.lock();
+                ReflectTools.setJavaFieldValue(session, field, lock);
+                return () -> removeLock(session, field);
+            } catch (NoSuchFieldException e) {
+                getLogger().debug("Cannot access lock field on VaadinSession",
+                        e);
+            }
+        }
+        return () -> {
+        };
+    }
+
+    private static void removeLock(VaadinSession session, Field field) {
+        session.getLockInstance().unlock();
+        ReflectTools.setJavaFieldValue(session, field, null);
+    }
+
+    private static Logger getLogger() {
+        return LoggerFactory.getLogger(TransientAwareHolder.class);
+    }
 }