fix: fix inspection and injection of Vaadin scoped beans (#152)

Vaadin scoped beans requires VaadinSession and UI thread locals to be set
in order to lookup beans.
This change tracks Vaadin scoped beans during serialization and makes sure
that required thread locals are set during deserialization.

The serialVersionUID added to TransientDescriptor should keep it compatible
with older version, preventing deserialization of previous data to fail.

Fixes #140
Requires vaadin/flow#20394

Co-authored-by: Tamas Mak <tamas@vaadin.com>

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

@@ -321,6 +321,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,"
@@ -418,7 +423,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);
+    }
 }