killbill-memoizeit

diff --git a/profiles/killbill/src/main/java/org/killbill/billing/server/modules/KillBillShiroWebModule.java b/profiles/killbill/src/main/java/org/killbill/billing/server/modules/KillBillShiroWebModule.java
index 4d8f5c1..49057fb 100644
--- a/profiles/killbill/src/main/java/org/killbill/billing/server/modules/KillBillShiroWebModule.java
+++ b/profiles/killbill/src/main/java/org/killbill/billing/server/modules/KillBillShiroWebModule.java
@@ -31,10 +31,10 @@ import org.apache.shiro.session.mgt.SessionManager;
 import org.apache.shiro.web.filter.authc.BasicHttpAuthenticationFilter;
 import org.apache.shiro.web.mgt.DefaultWebSecurityManager;
 import org.apache.shiro.web.mgt.WebSecurityManager;
-import org.apache.shiro.web.session.mgt.DefaultWebSessionManager;
 import org.apache.shiro.web.util.WebUtils;
 import org.killbill.billing.jaxrs.resources.JaxrsResource;
 import org.killbill.billing.server.security.FirstSuccessfulStrategyWith540;
+import org.killbill.billing.server.security.KillBillWebSessionManager;
 import org.killbill.billing.server.security.KillbillJdbcTenantRealm;
 import org.killbill.billing.util.config.definition.RbacConfig;
 import org.killbill.billing.util.glue.EhcacheShiroManagerProvider;
@@ -111,7 +111,7 @@ public class KillBillShiroWebModule extends ShiroWebModuleWith435 {
     protected void bindSessionManager(final AnnotatedBindingBuilder<SessionManager> bind) {
         // Bypass the servlet container completely for session management and delegate it to Shiro.
         // The default session timeout is 30 minutes.
-        bind.to(DefaultWebSessionManager.class).asEagerSingleton();
+        bind.to(KillBillWebSessionManager.class).asEagerSingleton();
 
         // Magic provider to configure the session DAO
         bind(JDBCSessionDao.class).toProvider(JDBCSessionDaoProvider.class).asEagerSingleton();

diff --git a/profiles/killbill/src/main/java/org/killbill/billing/server/security/KillBillWebSessionManager.java b/profiles/killbill/src/main/java/org/killbill/billing/server/security/KillBillWebSessionManager.java
new file mode 100644
index 0000000..e67ec0e
--- /dev/null
+++ b/profiles/killbill/src/main/java/org/killbill/billing/server/security/KillBillWebSessionManager.java
@@ -0,0 +1,48 @@
+/*
+ * Copyright 2017 Groupon, Inc
+ * Copyright 2017 The Billing Project, LLC
+ *
+ * The Billing Project licenses this file to you under the Apache License, version 2.0
+ * (the "License"); you may not use this file except in compliance with the
+ * License.  You may obtain a copy of the License at:
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.killbill.billing.server.security;
+
+import org.apache.shiro.session.Session;
+import org.apache.shiro.session.mgt.SessionContext;
+import org.apache.shiro.web.session.mgt.DefaultWebSessionManager;
+
+public class KillBillWebSessionManager extends DefaultWebSessionManager {
+
+    @Override
+    protected Session newSessionInstance(final SessionContext context) {
+        final Session session = super.newSessionInstance(context);
+
+        // DefaultWebSessionManager will call applyGlobalSessionTimeout() in
+        // start() below instead, which in turn calls onChange() and triggers a DAO UPDATE call
+        session.setTimeout(getGlobalSessionTimeout());
+
+        return session;
+    }
+
+    @Override
+    public Session start(final SessionContext context) {
+        final Session session = createSession(context);
+
+        // See above
+        //applyGlobalSessionTimeout(session);
+
+        onStart(session, context);
+        notifyStart(session);
+        return createExposedSession(session, context);
+    }
+}

diff --git a/util/src/main/java/org/killbill/billing/util/glue/EhcacheShiroManagerProvider.java b/util/src/main/java/org/killbill/billing/util/glue/EhcacheShiroManagerProvider.java
index 1af8a98..943450c 100644
--- a/util/src/main/java/org/killbill/billing/util/glue/EhcacheShiroManagerProvider.java
+++ b/util/src/main/java/org/killbill/billing/util/glue/EhcacheShiroManagerProvider.java
@@ -61,7 +61,9 @@ public class EhcacheShiroManagerProvider extends CacheProviderBase implements Pr
 
         if (securityManager instanceof DefaultSecurityManager) {
             // For RBAC only (see also KillbillJdbcTenantRealmProvider)
-            ((DefaultSecurityManager) securityManager).setCacheManager(shiroEhCacheManager);
+            final DefaultSecurityManager securityManager = (DefaultSecurityManager) this.securityManager;
+            securityManager.setCacheManager(shiroEhCacheManager);
+            securityManager.setSubjectDAO(new KillBillSubjectDAO());
         }
 
         return shiroEhCacheManager;

diff --git a/util/src/main/java/org/killbill/billing/util/glue/KillBillSubjectDAO.java b/util/src/main/java/org/killbill/billing/util/glue/KillBillSubjectDAO.java
new file mode 100644
index 0000000..176956f
--- /dev/null
+++ b/util/src/main/java/org/killbill/billing/util/glue/KillBillSubjectDAO.java
@@ -0,0 +1,83 @@
+/*
+ * Copyright 2017 Groupon, Inc
+ * Copyright 2017 The Billing Project, LLC
+ *
+ * The Billing Project licenses this file to you under the Apache License, version 2.0
+ * (the "License"); you may not use this file except in compliance with the
+ * License.  You may obtain a copy of the License at:
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.killbill.billing.util.glue;
+
+import org.apache.shiro.mgt.DefaultSubjectDAO;
+import org.apache.shiro.mgt.SessionsSecurityManager;
+import org.apache.shiro.session.Session;
+import org.apache.shiro.session.mgt.DefaultSessionManager;
+import org.apache.shiro.subject.Subject;
+import org.apache.shiro.subject.support.DelegatingSubject;
+import org.apache.shiro.util.CollectionUtils;
+import org.killbill.billing.util.security.shiro.dao.JDBCSessionDao;
+
+public class KillBillSubjectDAO extends DefaultSubjectDAO {
+
+    @Override
+    protected void saveToSession(final Subject subject) {
+        boolean updatesDisabled = false;
+
+        Session session = subject.getSession(false);
+        if (session == null && !CollectionUtils.isEmpty(subject.getPrincipals())) {
+            // Force the creation of the session here to get the id
+            session = subject.getSession();
+            // Optimize the session creation path: the default saveToSession implementation
+            // will call setAttribute() several times in a row, causing unnecessary DAO UPDATE queries
+            updatesDisabled = disableUpdatesForSession(subject, session);
+        }
+
+        super.saveToSession(subject);
+
+        if (updatesDisabled) {
+            enableUpdatesForSession(subject, session);
+        }
+    }
+
+    private boolean disableUpdatesForSession(final Subject subject, final Session session) {
+        final JDBCSessionDao sessionDAO = getJDBCSessionDao(subject);
+        if (sessionDAO != null) {
+            sessionDAO.disableUpdatesForSession(session);
+            return true;
+        }
+        return false;
+    }
+
+    private void enableUpdatesForSession(final Subject subject, final Session session) {
+        final JDBCSessionDao sessionDAO = getJDBCSessionDao(subject);
+        if (sessionDAO != null) {
+            sessionDAO.enableUpdatesForSession(session);
+        }
+    }
+
+    private JDBCSessionDao getJDBCSessionDao(final Subject subject) {
+        if (subject instanceof DelegatingSubject) {
+            final DelegatingSubject delegatingSubject = (DelegatingSubject) subject;
+            if (delegatingSubject.getSecurityManager() instanceof SessionsSecurityManager) {
+                final SessionsSecurityManager securityManager = (SessionsSecurityManager) delegatingSubject.getSecurityManager();
+                if (securityManager.getSessionManager() instanceof DefaultSessionManager) {
+                    final DefaultSessionManager sessionManager = (DefaultSessionManager) securityManager.getSessionManager();
+                    if (sessionManager.getSessionDAO() instanceof JDBCSessionDao) {
+                        return (JDBCSessionDao) sessionManager.getSessionDAO();
+                    }
+                }
+            }
+        }
+
+        return null;
+    }
+}

diff --git a/util/src/main/java/org/killbill/billing/util/security/shiro/dao/JDBCSessionDao.java b/util/src/main/java/org/killbill/billing/util/security/shiro/dao/JDBCSessionDao.java
index 41a05eb..a6a5a34 100644
--- a/util/src/main/java/org/killbill/billing/util/security/shiro/dao/JDBCSessionDao.java
+++ b/util/src/main/java/org/killbill/billing/util/security/shiro/dao/JDBCSessionDao.java
@@ -1,7 +1,7 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
- * Copyright 2014-2016 Groupon, Inc
- * Copyright 2014-2016 The Billing Project, LLC
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 The Billing Project, LLC
  *
  * The Billing Project licenses this file to you under the Apache License, version 2.0
  * (the "License"); you may not use this file except in compliance with the
@@ -21,6 +21,7 @@ package org.killbill.billing.util.security.shiro.dao;
 import java.io.IOException;
 import java.io.Serializable;
 import java.util.UUID;
+import java.util.concurrent.TimeUnit;
 
 import javax.inject.Inject;
 
@@ -31,12 +32,17 @@ import org.skife.jdbi.v2.IDBI;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
+
 public class JDBCSessionDao extends CachingSessionDAO {
 
     private static final Logger log = LoggerFactory.getLogger(JDBCSessionDao.class);
 
     private final JDBCSessionSqlDao jdbcSessionSqlDao;
 
+    private final Cache<Serializable, Boolean> noUpdateSessionsCache = CacheBuilder.<Serializable, Boolean>newBuilder().expireAfterWrite(5, TimeUnit.SECONDS).build();
+
     @Inject
     public JDBCSessionDao(final IDBI dbi) {
         this.jdbcSessionSqlDao = dbi.onDemand(JDBCSessionSqlDao.class);
@@ -44,7 +50,9 @@ public class JDBCSessionDao extends CachingSessionDAO {
 
     @Override
     protected void doUpdate(final Session session) {
-        jdbcSessionSqlDao.update(new SessionModelDao(session));
+        if (shouldUpdateSession(session)) {
+            jdbcSessionSqlDao.update(new SessionModelDao(session));
+        }
     }
 
     @Override
@@ -56,9 +64,12 @@ public class JDBCSessionDao extends CachingSessionDAO {
     protected Serializable doCreate(final Session session) {
         final UUID sessionId = UUIDs.randomUUID();
         // See SessionModelDao#toSimpleSession for why we use toString()
-        assignSessionId(session, sessionId.toString());
+        final String sessionIdAsString = sessionId.toString();
+        assignSessionId(session, sessionIdAsString);
         jdbcSessionSqlDao.create(new SessionModelDao(session));
-        return sessionId;
+        // Make sure to return a String here as well, or Shiro will cache the Session with a UUID key
+        // while it is expecting String
+        return sessionIdAsString;
     }
 
     @Override
@@ -82,4 +93,17 @@ public class JDBCSessionDao extends CachingSessionDAO {
             return null;
         }
     }
+
+    public void disableUpdatesForSession(final Session session) {
+        noUpdateSessionsCache.put(session.getId(), Boolean.TRUE);
+    }
+
+    public void enableUpdatesForSession(final Session session) {
+        noUpdateSessionsCache.invalidate(session.getId());
+        doUpdate(session);
+    }
+
+    private boolean shouldUpdateSession(final Session session) {
+        return noUpdateSessionsCache.getIfPresent(session.getId()) == Boolean.TRUE ? Boolean.FALSE : Boolean.TRUE;
+    }
 }