2 files changed, 58 insertions, 4 deletions

diff --git a/server/src/main/java/com/vaadin/server/VaadinService.java b/server/src/main/java/com/vaadin/server/VaadinService.java
index 888b82a3cb..25354074e1 100644
--- a/server/src/main/java/com/vaadin/server/VaadinService.java
+++ b/server/src/main/java/com/vaadin/server/VaadinService.java
@@ -1156,8 +1156,13 @@ public abstract class VaadinService implements Serializable {
             if (value instanceof VaadinSession) {
                 // set flag to avoid cleanup
                 VaadinSession serviceSession = (VaadinSession) value;
-                serviceSession.setAttribute(PRESERVE_UNBOUND_SESSION_ATTRIBUTE,
-                        Boolean.TRUE);
+                serviceSession.lock();
+                try {
+                    serviceSession.setAttribute(
+                            PRESERVE_UNBOUND_SESSION_ATTRIBUTE, Boolean.TRUE);
+                } finally {
+                    serviceSession.unlock();
+                }
             }
             attrs.put(name, value);
         }
@@ -1183,8 +1188,13 @@ public abstract class VaadinService implements Serializable {
                         serviceSession.getLockInstance());
 
                 service.storeSession(serviceSession, newSession);
-                serviceSession.setAttribute(PRESERVE_UNBOUND_SESSION_ATTRIBUTE,
-                        null);
+                serviceSession.lock();
+                try {
+                    serviceSession.setAttribute(
+                            PRESERVE_UNBOUND_SESSION_ATTRIBUTE, null);
+                } finally {
+                    serviceSession.unlock();
+                }
             }
         }
 
diff --git a/server/src/test/java/com/vaadin/server/VaadinServiceTest.java b/server/src/test/java/com/vaadin/server/VaadinServiceTest.java
index f10f55db36..f179ad1d58 100644
--- a/server/src/test/java/com/vaadin/server/VaadinServiceTest.java
+++ b/server/src/test/java/com/vaadin/server/VaadinServiceTest.java
@@ -8,8 +8,11 @@ import javax.servlet.ServletConfig;
 import javax.servlet.ServletException;
 import javax.servlet.http.HttpSessionBindingEvent;
 
+import java.util.Collections;
+
 import org.easymock.EasyMock;
 import org.junit.Test;
+import org.mockito.Mockito;
 
 import com.vaadin.shared.Registration;
 import com.vaadin.util.CurrentInstance;
@@ -179,4 +182,45 @@ public class VaadinServiceTest {
         assertEquals(1, listener.callCount);
         assertEquals(1, listener2.callCount);
     }
+
+    @Test
+    public void reinitializeSession_setVaadinSessionAttriuteWithLock() {
+        VaadinRequest request = Mockito.mock(VaadinRequest.class);
+
+        VaadinSession vaadinSession = Mockito.mock(VaadinSession.class);
+        VaadinSession newVaadinSession = Mockito.mock(VaadinSession.class);
+
+        WrappedSession session = mockSession(request, vaadinSession, "foo");
+
+        Mockito.doAnswer(invocation -> {
+            mockSession(request, newVaadinSession, "bar");
+            return null;
+        }).when(session).invalidate();
+
+        VaadinService.reinitializeSession(request);
+
+        Mockito.verify(vaadinSession, Mockito.times(2)).lock();
+        Mockito.verify(vaadinSession).setAttribute(
+                VaadinService.PRESERVE_UNBOUND_SESSION_ATTRIBUTE, Boolean.TRUE);
+        Mockito.verify(vaadinSession).setAttribute(
+                VaadinService.PRESERVE_UNBOUND_SESSION_ATTRIBUTE, null);
+        Mockito.verify(vaadinSession, Mockito.times(2)).unlock();
+    }
+
+    private WrappedSession mockSession(VaadinRequest request,
+            VaadinSession vaadinSession, String attributeName) {
+        WrappedSession session = Mockito.mock(WrappedSession.class);
+        Mockito.when(request.getWrappedSession()).thenReturn(session);
+
+        Mockito.when(session.getAttributeNames())
+                .thenReturn(Collections.singleton(attributeName));
+
+        Mockito.when(session.getAttribute(attributeName))
+                .thenReturn(vaadinSession);
+
+        VaadinService service = Mockito.mock(VaadinService.class);
+
+        Mockito.when(vaadinSession.getService()).thenReturn(service);
+        return session;
+    }
 }