diff --git a/lib/redis/semaphore.rb b/lib/redis/semaphore.rb
index 461fa72..ea91315 100644
--- a/lib/redis/semaphore.rb
+++ b/lib/redis/semaphore.rb
@@ -7,7 +7,7 @@ class Semaphore
 
     # stale_client_timeout is the threshold of time before we assume
     # that something has gone terribly wrong with a client and we
-    # invalidate it's lock.
+    # invalidate its lock.
     # Default is nil for which we don't check for stale clients
     # Redis::Semaphore.new(:my_semaphore, :stale_client_timeout => 30, :redis => myRedis)
     # Redis::Semaphore.new(:my_semaphore, :redis => myRedis)
@@ -81,6 +81,7 @@ def lock(timeout = nil)
         end
       end
 
+      set_expiration_if_necessary
       return_value
     end
     alias_method :wait, :lock
@@ -199,7 +200,7 @@ def create!
 
     def set_expiration_if_necessary
       if @expiration
-        [available_key, exists_key, version_key].each do |key|
+        [available_key, exists_key, version_key, @stale_client_timeout.nil?? grabbed_key : nil ].each do |key|
           @redis.expire(key, @expiration)
         end
       end
diff --git a/redis-semaphore.gemspec b/redis-semaphore.gemspec
index 4c7d2ce..689df36 100644
--- a/redis-semaphore.gemspec
+++ b/redis-semaphore.gemspec
@@ -1,6 +1,6 @@
 Gem::Specification.new do |s|
   s.name        = 'redis-semaphore'
-  s.version     = '0.3.1'
+  s.version     = '0.4.0'
   s.summary     = 'Implements a distributed semaphore or mutex using Redis.'
   s.authors     = ['David Verhasselt']
   s.email       = 'david@crowdway.com'