StickySessionRepository.java

package de.classyfi.session;

import java.time.Duration;
import java.time.Instant;
import java.util.Set;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.springframework.context.ApplicationEvent;
import org.springframework.context.ApplicationEventPublisher;
import org.springframework.session.MapSession;
import org.springframework.session.MapSessionRepository;
import org.springframework.session.Session;
import org.springframework.session.SessionRepository;
import org.springframework.session.events.AbstractSessionEvent;
import org.springframework.session.events.SessionCreatedEvent;
import org.springframework.session.events.SessionDeletedEvent;
import org.springframework.session.events.SessionDestroyedEvent;
import org.springframework.session.events.SessionExpiredEvent;

import de.classyfi.session.StickySessionRepository.StickySession;

/**
 * {@link SessionRepository} implementation that delegates to a (usually remote) session repository, but keeps
 * a local copy of the session in the configured cache.
 *
 * If the remote repository has a session with a newer {@link Session#getLastAccessedTime() lastAccessedTime}
 * than the local copy, the local copy is ignored.
 *
 * When used with a remote repository like a {@link org.springframework.session.data.redis.RedisIndexedSessionRepository},
 * this implementation provides no advantages in terms of latency, transferred data, or CPU cycles, because
 * the remote session is always fetched (and deserialized) and updated synchronously. However, it allows access
 * to transient variables stored within session attributes.
 *
 * @author Bernhard Frauendienst <[email protected]>
 */
public final class StickySessionRepository<S extends Session>
    implements SessionRepository<StickySession<S>> {

  private final static Logger log = LogManager.getLogger(StickySessionRepository.class);

  private final SessionRepository<S> delegate;

  private final MapSessionRepository sessionCache;

  private ApplicationEventPublisher eventPublisher = event -> {};

  private EventPublisher stickyEventPublisher = null;

  public StickySessionRepository(SessionRepository<S> delegate, MapSessionRepository sessionCache) {
    this.delegate = delegate;
    this.sessionCache = sessionCache;
  }

  public void setApplicationEventPublisher(ApplicationEventPublisher eventPublisher) {
    this.eventPublisher = eventPublisher;
  }

  /**
   * Returns an ApplicationEventPublisher that should be supplied to the delegate session repository. It will re-publish the
   * session events from the delegate repository with StickySession objects and with *this* as source.
   *
   */
  public ApplicationEventPublisher getStickyEventPublisher() {
    if (stickyEventPublisher == null) {
      stickyEventPublisher = new EventPublisher();
    }
    return stickyEventPublisher;
  }

  @Override
  public StickySession<S> createSession() {
    var delegate = this.delegate.createSession();
    return new StickySession<>(delegate);
  }

  @Override
  public void save(StickySession<S> session) {
    delegate.save(session.delegate);
    session.cached.setLastAccessedTime(session.delegate.getLastAccessedTime());
    final MapSession existingCached = sessionCache.findById(session.getId());
    if (existingCached == null || !existingCached.getLastAccessedTime().isAfter(session.getLastAccessedTime())) {
      sessionCache.save(session.cached);
    }
  }

  @Override
  public StickySession<S> findById(String id) {
    var delegate = this.delegate.findById(id);
    if (delegate == null || delegate.isExpired()) {
      sessionCache.deleteById(id);
      return null;
    }
    var cached = sessionCache.findById(id);
    if (cached == null) {
      return new StickySession<>(delegate);
    }
    // if the delegate session is newer than our cache, we need to evict it
    if (delegate.getLastAccessedTime().isAfter(cached.getLastAccessedTime())) {
      sessionCache.deleteById(id);
      return new StickySession<>(delegate);
    }
    return new StickySession<>(delegate, cached);
  }

  @Override
  public void deleteById(String id) {
    sessionCache.deleteById(id);
    delegate.deleteById(id);
  }

  final static class StickySession<S extends Session> implements Session {

    private S delegate;

    private MapSession cached;

    public StickySession(S delegate) {
      this(delegate, new MapSession(delegate));
    }

    public StickySession(S delegate, MapSession cached) {
      this.delegate = delegate;
      this.cached = cached;
    }

    @Override
    public String getId() {
      return delegate.getId();
    }

    @Override
    public String changeSessionId() {
      var sessionId = delegate.changeSessionId();
      cached.setId(sessionId);
      return sessionId;
    }

    @Override
    public <T> T getAttribute(String attributeName) {
      return cached.getAttribute(attributeName);
    }

    @Override
    public Set<String> getAttributeNames() {
      return cached.getAttributeNames();
    }

    @Override
    public void setAttribute(String attributeName, Object attributeValue) {
      cached.setAttribute(attributeName, attributeValue);
      delegate.setAttribute(attributeName, attributeValue);
    }

    @Override
    public void removeAttribute(String attributeName) {
      cached.removeAttribute(attributeName);
      delegate.removeAttribute(attributeName);
    }

    @Override
    public Instant getCreationTime() {
      return cached.getCreationTime();
    }

    @Override
    public Instant getLastAccessedTime() {
      return cached.getLastAccessedTime();
    }

    @Override
    public void setLastAccessedTime(Instant lastAccessedTime) {
      cached.setLastAccessedTime(lastAccessedTime);
      delegate.setLastAccessedTime(lastAccessedTime);
    }

    @Override
    public Duration getMaxInactiveInterval() {
      return cached.getMaxInactiveInterval();
    }

    @Override
    public void setMaxInactiveInterval(Duration interval) {
      cached.setMaxInactiveInterval(interval);
      delegate.setMaxInactiveInterval(interval);
    }

    @Override
    public boolean isExpired() {
      return delegate.isExpired();
    }
  }

  class EventPublisher implements ApplicationEventPublisher {
    @Override
    public void publishEvent(ApplicationEvent event) {
      if (event instanceof AbstractSessionEvent) {
        publishEvent((AbstractSessionEvent) event);
      }
    }

    @Override
    public void publishEvent(Object event) {
      if (event instanceof AbstractSessionEvent) {
        publishEvent((AbstractSessionEvent) event);
      }
    }

    private void publishEvent(AbstractSessionEvent event) {
      if (event.getSource() != delegate) {
        log.warn("Will not publish " + event.getClass().getSimpleName() + " not originating from " + delegate);
        return;
      }

      S delegateSession = event.getSession();
      if (delegateSession == null) {
        // AbstractSessionEvent javadocs claims this can happen. AFAICT, the source code says otherwise.
        log.warn("Cannot publish " + event.getClass().getSimpleName() + " for session " + event.getSessionId()
            + ", no cached session found.");
        return;
      }
      var cached = sessionCache.findById(event.getSessionId());
      if (cached == null) {
        log.warn("Cannot publish " + event.getClass().getSimpleName() + " for session " + event.getSessionId()
            + ", no cached session found.");
        return;
      }
      var session = new StickySession<>(delegateSession, cached);
      if (event instanceof SessionCreatedEvent) {
        eventPublisher.publishEvent(new SessionCreatedEvent(StickySessionRepository.this, session));
      } else if (event instanceof SessionDestroyedEvent) {
        if (event instanceof SessionDeletedEvent) {
          eventPublisher.publishEvent(new SessionDeletedEvent(StickySessionRepository.this, session));
        } else if (event instanceof SessionExpiredEvent) {
          eventPublisher.publishEvent(new SessionExpiredEvent(StickySessionRepository.this, session));
        }
        sessionCache.deleteById(event.getSessionId());
      } else {
        log.warn("Unknown event type " + event.getClass());
      }
    }
  }
}