From 72d761e7b13842e758193848b12967b8c2ebe270 Mon Sep 17 00:00:00 2001
From: Jon Leighton <j@jonathanleighton.com>
Date: Tue, 9 Dec 2014 23:42:06 +0000
Subject: [PATCH] Prevent the user having to run `spring stop` after adding
 commands

Too many people are getting confused by this
---
 CHANGELOG.md                       |  4 +++
 README.md                          |  3 +-
 lib/spring/client/run.rb           | 51 ++++++++++++++++++++++++------
 lib/spring/client/stop.rb          | 27 ++++------------
 lib/spring/env.rb                  | 25 +++++++++++++++
 lib/spring/test/acceptance_test.rb |  3 ++
 6 files changed, 81 insertions(+), 32 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 52072618..65f93464 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,9 @@
 ## Next version
 
+* Automatically restart spring after new commands are added. This means
+  that you can add spring-commands-rspec to your Gemfile and then
+  immediately start using it, without having to run `spring stop`.
+  (Spring will effectively run `spring stop` for you.)
 * Make app reloading work in apps which spew out lots of output on
   startup (previously a buffer would fill up and cause the process to
   hang). Issue #332.
diff --git a/README.md b/README.md
index 2ccbed02..ed9bfb6a 100644
--- a/README.md
+++ b/README.md
@@ -218,8 +218,7 @@ speed-up).
 
 ### Additional commands
 
-You can add these to your Gemfile for additional commands (run `spring stop` afterwards
-to pick up the changes):
+You can add these to your Gemfile for additional commands:
 
 * [spring-commands-rspec](https://github.com/jonleighton/spring-commands-rspec)
 * [spring-commands-cucumber](https://github.com/jonleighton/spring-commands-cucumber)
diff --git a/lib/spring/client/run.rb b/lib/spring/client/run.rb
index a03c7d02..ac98bbe8 100644
--- a/lib/spring/client/run.rb
+++ b/lib/spring/client/run.rb
@@ -21,7 +21,37 @@ def server
       end
 
       def call
-        boot_server unless env.server_running?
+        if env.server_running?
+          warm_run
+        else
+          cold_run
+        end
+      rescue Errno::ECONNRESET
+        exit 1
+      ensure
+        server.close if @server
+      end
+
+      def warm_run
+        run
+      rescue CommandNotFound
+        require "spring/commands"
+
+        if Spring.command?(args.first)
+          # Command installed since spring started
+          stop_server
+          cold_run
+        else
+          raise
+        end
+      end
+
+      def cold_run
+        boot_server
+        run
+      end
+
+      def run
         verify_server_version
 
         application, client = UNIXSocket.pair
@@ -29,19 +59,16 @@ def call
         queue_signals
         connect_to_application(client)
         run_command(client, application)
-      rescue Errno::ECONNRESET
-        exit 1
-      ensure
-        server.close if @server
       end
 
       def boot_server
         env.socket_path.unlink if env.socket_path.exist?
 
-        pid = fork {
-          require "spring/server"
-          Spring::Server.boot
-        }
+        pid = Process.spawn(
+          "ruby",
+          "-r", "spring/server",
+          "-e", "Spring::Server.boot"
+        )
 
         until env.socket_path.exist?
           _, status = Process.waitpid2(pid, Process::WNOHANG)
@@ -50,6 +77,12 @@ def boot_server
         end
       end
 
+      def stop_server
+        server.close
+        @server = nil
+        env.stop
+      end
+
       def verify_server_version
         server_version = server.gets.chomp
         if server_version != env.version
diff --git a/lib/spring/client/stop.rb b/lib/spring/client/stop.rb
index a448de29..af4da509 100644
--- a/lib/spring/client/stop.rb
+++ b/lib/spring/client/stop.rb
@@ -3,35 +3,20 @@
 module Spring
   module Client
     class Stop < Command
-      TIMEOUT = 2 # seconds
-
       def self.description
         "Stop all spring processes for this project."
       end
 
       def call
-        if env.server_running?
-          timeout = Time.now + TIMEOUT
-          kill 'TERM'
-          sleep 0.1 until !env.server_running? || Time.now >= timeout
-
-          if env.server_running?
-            $stderr.puts "Spring did not stop; killing forcibly."
-            kill 'KILL'
-          else
-            puts "Spring stopped."
-          end
-        else
+        case env.stop
+        when :stopped
+          puts "Spring stopped."
+        when :killed
+          $stderr.puts "Spring did not stop; killing forcibly."
+        when :not_running
           puts "Spring is not running"
         end
       end
-
-      def kill(sig)
-        pid = env.pid
-        Process.kill(sig, pid) if pid
-      rescue Errno::ESRCH
-        # already dead
-      end
     end
   end
 end
diff --git a/lib/spring/env.rb b/lib/spring/env.rb
index 9ab6347c..e1e9656f 100644
--- a/lib/spring/env.rb
+++ b/lib/spring/env.rb
@@ -9,6 +9,7 @@
 
 module Spring
   IGNORE_SIGNALS = %w(INT QUIT)
+  STOP_TIMEOUT = 2 # seconds
 
   class Env
     attr_reader :log_file
@@ -80,5 +81,29 @@ def log(message)
       log_file.puts "[#{Time.now}] [#{Process.pid}] #{message}"
       log_file.flush
     end
+
+    def stop
+      if server_running?
+        timeout = Time.now + STOP_TIMEOUT
+        kill 'TERM'
+        sleep 0.1 until !server_running? || Time.now >= timeout
+
+        if server_running?
+          kill 'KILL'
+          :killed
+        else
+          :stopped
+        end
+      else
+        :not_running
+      end
+    end
+
+    def kill(sig)
+      pid = self.pid
+      Process.kill(sig, pid) if pid
+    rescue Errno::ESRCH
+      # already dead
+    end
   end
 end
diff --git a/lib/spring/test/acceptance_test.rb b/lib/spring/test/acceptance_test.rb
index 8260310f..3ebca0fb 100644
--- a/lib/spring/test/acceptance_test.rb
+++ b/lib/spring/test/acceptance_test.rb
@@ -147,6 +147,9 @@ def self.omg
       end
 
       test "custom commands" do
+        # Start spring before setting up the command, to test that it gracefully upgrades itself
+        assert_success "bin/rails runner ''"
+
         File.write(app.spring_config, <<-CODE)
           class CustomCommand
             def call
