Refactor session stores by providing a with_lock helper and by moving the responsitility to handle :renew and :drop up to Abstract::ID.

Signed-off-by: raggi <jftucker@gmail.com>

5 files changed, 72 insertions, 101 deletions

diff --git a/lib/rack/session/abstract/id.rb b/lib/rack/session/abstract/id.rb
index 4d211c4d..1f6824ad 100644
--- a/lib/rack/session/abstract/id.rb
+++ b/lib/rack/session/abstract/id.rb
@@ -115,6 +115,11 @@ module Rack
           options = env['rack.session.options']
           session_id = options[:id]
 
+          if options[:drop] || options[:renew]
+            session_id = destroy_session(env, session_id, options)
+            return [status, headers, body] unless session_id
+          end
+
           if not data = set_session(env, session_id, session, options)
             env["rack.errors"].puts("Warning! #{self.class.name} failed to save session. Content dropped.")
           elsif options[:defer] and not options[:renew]
@@ -151,9 +156,17 @@ module Rack
         # All thread safety and session storage proceedures should occur here.
         # Should return true or false dependant on whether or not the session
         # was saved or not.
+
         def set_session(env, sid, session, options)
           raise '#set_session not implemented.'
         end
+
+        # All thread safety and session destroy proceedures should occur here.
+        # Should return a new session id or nil if options[:drop]
+
+        def destroy_session(env, sid, options)
+          raise '#destroy_session not implemented'
+        end
       end
     end
   end
diff --git a/lib/rack/session/cookie.rb b/lib/rack/session/cookie.rb
index 99d39e9d..d908ab5d 100644
--- a/lib/rack/session/cookie.rb
+++ b/lib/rack/session/cookie.rb
@@ -50,6 +50,12 @@ module Rack
         env["rack.session.options"] = @default_options.dup
       end
 
+      # Overwrite set cookie to bypass content equality and always stream the cookie.
+
+      def set_cookie(env, headers, cookie)
+        Utils.set_cookie_header!(headers, @key, cookie)
+      end
+
       def set_session(env, session_id, session, options)
         session_data = [Marshal.dump(session)].pack("m*")
 
@@ -65,10 +71,9 @@ module Rack
         end
       end
 
-      # Overwrite set cookie to bypass content equality and always stream the cookie.
-
-      def set_cookie(env, headers, cookie)
-        Utils.set_cookie_header!(headers, @key, cookie)
+      def destroy_session(env, session_id, options)
+        # Nothing to do here, data is in the client
+        generate_sid unless options[:drop]
       end
 
       def generate_hmac(data)
diff --git a/lib/rack/session/memcache.rb b/lib/rack/session/memcache.rb
index 5e77774a..0f83dcbd 100644
--- a/lib/rack/session/memcache.rb
+++ b/lib/rack/session/memcache.rb
@@ -21,6 +21,7 @@ module Rack
 
     class Memcache < Abstract::ID
       attr_reader :mutex, :pool
+
       DEFAULT_OPTIONS = Abstract::ID::DEFAULT_OPTIONS.merge \
         :namespace => 'rack:session',
         :memcache_server => 'localhost:11211'
@@ -45,75 +46,48 @@ module Rack
         end
       end
 
-      def get_session(env, session_id)
-        @mutex.lock if env['rack.multithread']
-        unless session_id and session = @pool.get(session_id)
-          session_id, session = generate_sid, {}
-          unless /^STORED/ =~ @pool.add(session_id, session)
-            raise "Session collision on '#{session_id.inspect}'"
+      def get_session(env, sid)
+        with_lock(env, [nil, {}]) do
+          unless sid and session = @pool.get(sid)
+            sid, session = generate_sid, {}
+            unless /^STORED/ =~ @pool.add(sid, session)
+              raise "Session collision on '#{sid.inspect}'"
+            end
           end
+          [sid, session]
         end
-        session.instance_variable_set '@old', @pool.get(session_id, true)
-        return [session_id, session]
-      rescue MemCache::MemCacheError, Errno::ECONNREFUSED
-        # MemCache server cannot be contacted
-        warn "#{self} is unable to find memcached server."
-        warn $!.inspect
-        return [ nil, {} ]
-      ensure
-        @mutex.unlock if @mutex.locked?
       end
 
       def set_session(env, session_id, new_session, options)
         expiry = options[:expire_after]
         expiry = expiry.nil? ? 0 : expiry + 1
 
-        @mutex.lock if env['rack.multithread']
-        if options[:renew] or options[:drop]
-          @pool.delete session_id
-          return false if options[:drop]
-          session_id = generate_sid
-          @pool.add session_id, {} # so we don't worry about cache miss on #set
+        with_lock(env, false) do
+          @pool.set session_id, new_session, expiry
+          session_id
         end
+      end
 
-        session = @pool.get(session_id) || {}
-        old_session = new_session.instance_variable_get '@old'
-        old_session = old_session ? Marshal.load(old_session) : {}
-
-        unless Hash === old_session and Hash === new_session
-          env['rack.errors'].
-            puts 'Bad old_session or new_session sessions provided.'
-        else # merge sessions
-          # alterations are either update or delete, making as few changes as
-          # possible to prevent possible issues.
-
-          # removed keys
-          delete = old_session.keys - new_session.keys
-          if $VERBOSE and not delete.empty?
-            env['rack.errors'].
-              puts "//@#{session_id}: delete #{delete*','}"
-          end
-          delete.each{|k| session.delete k }
-
-          # added or altered keys
-          update = new_session.keys.
-            select{|k| new_session[k] != old_session[k] }
-          if $VERBOSE and not update.empty?
-            env['rack.errors'].puts "//@#{session_id}: update #{update*','}"
-          end
-          update.each{|k| session[k] = new_session[k] }
+      def destroy_session(env, session_id, options)
+        with_lock(env) do
+          @pool.delete(session_id)
+          generate_sid unless options[:drop]
         end
+      end
 
-        @pool.set session_id, session, expiry
-        session_id
+      def with_lock(env, default=nil)
+        @mutex.lock if env['rack.multithread']
+        yield
       rescue MemCache::MemCacheError, Errno::ECONNREFUSED
-        # MemCache server cannot be contacted
-        warn "#{self} is unable to find memcached server."
-        warn $!.inspect
-        return false
+        if $VERBOSE
+          warn "#{self} is unable to find memcached server."
+          warn $!.inspect
+        end
+        default
       ensure
         @mutex.unlock if @mutex.locked?
       end
+
     end
   end
 end
diff --git a/lib/rack/session/pool.rb b/lib/rack/session/pool.rb
index b3f8bd72..d4774fed 100644
--- a/lib/rack/session/pool.rb
+++ b/lib/rack/session/pool.rb
@@ -42,59 +42,38 @@ module Rack
       end
 
       def get_session(env, sid)
-        session = @pool[sid] if sid
-        @mutex.lock if env['rack.multithread']
-        unless sid and session
-          env['rack.errors'].puts("Session '#{sid.inspect}' not found, initializing...") if $VERBOSE and not sid.nil?
-          session = {}
-          sid = generate_sid
-          @pool.store sid, session
+        with_lock(env, [nil, {}]) do
+          unless sid and session = @pool[sid]
+            sid, session = generate_sid, {}
+            @pool.store sid, session
+          end
+          [sid, session]
         end
-        session.instance_variable_set('@old', {}.merge(session))
-        return [sid, session]
-      ensure
-        @mutex.unlock if env['rack.multithread']
       end
 
       def set_session(env, session_id, new_session, options)
-        @mutex.lock if env['rack.multithread']
-        session = @pool[session_id]
-        if options[:renew] or options[:drop]
-          @pool.delete session_id
-          return false if options[:drop]
-          session_id = generate_sid
-          @pool.store session_id, 0
+        with_lock(env, false) do
+          @pool.store session_id, new_session
+          session_id
         end
-        old_session = new_session.instance_variable_get('@old') || {}
-        session = merge_sessions session_id, old_session, new_session, session
-        @pool.store session_id, session
-        return session_id
-      rescue
-        warn "#{new_session.inspect} has been lost."
-        warn $!.inspect
-      ensure
-        @mutex.unlock if env['rack.multithread']
       end
 
-      private
-
-      def merge_sessions sid, old, new, cur=nil
-        cur ||= {}
-        unless Hash === old and Hash === new
-          warn 'Bad old or new sessions provided.'
-          return cur
+      def destroy_session(env, session_id, options)
+        with_lock(env) do
+          @pool.delete(session_id)
+          generate_sid unless options[:drop]
         end
+      end
 
-        delete = old.keys - new.keys
-        warn "//@#{sid}: dropping #{delete*','}" if $DEBUG and not delete.empty?
-        delete.each{|k| cur.delete k }
-
-        update = new.keys.select{|k| new[k] != old[k] }
-        warn "//@#{sid}: updating #{update*','}" if $DEBUG and not update.empty?
-        update.each{|k| cur[k] = new[k] }
-
-        cur
+      def with_lock(env, default=nil)
+        @mutex.lock if env['rack.multithread']
+        yield
+      rescue
+        default
+      ensure
+        @mutex.unlock if @mutex.locked?
       end
+
     end
   end
 end
diff --git a/test/spec_session_pool.rb b/test/spec_session_pool.rb
index 60eb23a4..58f193c3 100644
--- a/test/spec_session_pool.rb
+++ b/test/spec_session_pool.rb
@@ -83,7 +83,7 @@ describe Rack::Session::Pool do
     pool.pool.size.should.equal 1
 
     res2 = dreq.get("/", "HTTP_COOKIE" => cookie)
-    res2["Set-Cookie"].should.equal nil
+    res2["Set-Cookie"].should.be.nil
     res2.body.should.equal '{"counter"=>2}'
     pool.pool.size.should.equal 0