Add liveness health check support

When MailRoom is run in Kubernetes, we have found occasions where
MailRoom appears to have attempted to stop running, but `Net::IMAP` is
stuck waiting for threads (ruby/net-imap#14).

This commit adds an HTTP liveness checker to enable detection of a
terminated MailRoom pod.

README.md

@@ -42,6 +42,9 @@ You will also need to install `faraday` or `letter_opener` if you use the `postb
 
 ```yaml
 ---
+:health_check:
+  :address: "127.0.0.1"
+  :port: 8080
 :mailboxes:
   -
     :email: "[email protected]"
@@ -133,6 +136,7 @@ You will also need to install `faraday` or `letter_opener` if you use the `postb
 **Note:** If using `delete_after_delivery`, you also probably want to use
 `expunge_deleted` unless you really know what you're doing.
 
+<<<<<<< HEAD
 ## inbox_method
 
 By default, IMAP mode is assumed for reading a mailbox.
@@ -218,6 +222,16 @@ for Microsoft Cloud for US Government:
       :graph_endpoint: https://graph.microsoft.us
 ```
 
+## health_check ##
+
+Requires `webrick` gem to be installed.
+
+This option enables an HTTP server that listens to a bind address
+defined by `address` and `port`. The following endpoints are supported:
+
+* `/liveness`: This returns a 200 status code with `OK` as the body if
+the server is running. Otherwise, it returns a 500 status code.
+
 ## delivery_method ##
 
 ### postback ###

lib/mail_room.rb

@@ -7,6 +7,7 @@ module MailRoom
 
 require "mail_room/version"
 require "mail_room/configuration"
+require "mail_room/health_check"
 require "mail_room/mailbox"
 require "mail_room/mailbox_watcher"
 require "mail_room/message"

lib/mail_room/cli.rb

@@ -40,7 +40,7 @@ def initialize(args)
       end.parse!(args)
 
       self.configuration = Configuration.new(options)
-      self.coordinator = Coordinator.new(configuration.mailboxes)
+      self.coordinator = Coordinator.new(configuration.mailboxes, configuration.health_check)
     end
 
     # Start the coordinator running, sets up signal traps

lib/mail_room/configuration.rb

@@ -4,7 +4,7 @@ module MailRoom
   # Wraps configuration for a set of individual mailboxes with global config
   # @author Tony Pitale
   class Configuration
-    attr_accessor :mailboxes, :log_path, :quiet
+    attr_accessor :mailboxes, :log_path, :quiet, :health_check
 
     # Initialize a new configuration of mailboxes
     def initialize(options={})
@@ -18,6 +18,7 @@ def initialize(options={})
           config_file = YAML.load(erb.result)
 
           set_mailboxes(config_file[:mailboxes])
+          set_health_check(config_file[:health_check])
         rescue => e
           raise e unless quiet
         end
@@ -32,5 +33,14 @@ def set_mailboxes(mailboxes_config)
         self.mailboxes << Mailbox.new(attributes)
       end
     end
+
+    # Builds the health checker from YAML configuration
+    #
+    # @param health_check_config nil or a Hash containing :address and :port
+    def set_health_check(health_check_config)
+      return unless health_check_config
+
+      self.health_check = HealthCheck.new(health_check_config)
+    end
   end
 end

lib/mail_room/coordinator.rb

@@ -2,31 +2,35 @@ module MailRoom
   # Coordinate the mailbox watchers
   # @author Tony Pitale
   class Coordinator
-    attr_accessor :watchers, :running
+    attr_accessor :watchers, :running, :health_check
 
     # build watchers for a set of mailboxes
     # @params mailboxes [Array<MailRoom::Mailbox>] mailboxes to be watched
-    def initialize(mailboxes)
+    # @params health_check <MailRoom::HealthCheck> health checker to run
+    def initialize(mailboxes, health_check = nil)
       self.watchers = []
 
+      @health_check = health_check
       mailboxes.each {|box| self.watchers << MailboxWatcher.new(box)}
     end
 
     alias :running? :running
 
     # start each of the watchers to running
     def run
+      health_check&.run
       watchers.each(&:run)
-      
+
       self.running = true
-      
+
       sleep_while_running
     ensure
       quit
     end
 
     # quit each of the watchers when we're done running
     def quit
+      health_check&.quit
       watchers.each(&:quit)
     end
 

lib/mail_room/health_check.rb

@@ -0,0 +1,60 @@
+# frozen_string_literal: true
+
+module MailRoom
+  class HealthCheck
+    attr_reader :address, :port, :running
+
+    def initialize(attributes = {})
+      @address = attributes[:address]
+      @port = attributes[:port]
+
+      validate!
+    end
+
+    def run
+      @server = create_server
+
+      @thread = Thread.new do
+        @server.start
+      end
+
+      @thread.abort_on_exception = true
+      @running = true
+    end
+
+    def quit
+      @running = false
+      @server&.shutdown
+      @thread&.join(60)
+    end
+
+    private
+
+    def validate!
+      raise 'No health check address specified' unless address
+      raise "Health check port #{@port.to_i} is invalid" unless port.to_i.positive?
+    end
+
+    def create_server
+      require 'webrick'
+
+      server = ::WEBrick::HTTPServer.new(Port: port, BindAddress: address, AccessLog: [])
+
+      server.mount_proc '/liveness' do |_req, res|
+        handle_liveness(res)
+      end
+
+      server
+    end
+
+    def handle_liveness(res)
+      if @running
+        res.status = 200
+        res.body = "OK\n"
+      else
+        res.status = 500
+        res.body = "Not running\n"
+      end
+    end
+  end
+end

mail_room.gemspec

@@ -26,6 +26,7 @@ Gem::Specification.new do |gem|
   gem.add_development_dependency "rubocop", "~> 1.11"
   gem.add_development_dependency "mocha", "~> 1.11"
   gem.add_development_dependency "simplecov"
+  gem.add_development_dependency "webrick", "~> 1.6"
 
   # for testing delivery methods
   gem.add_development_dependency "faraday"

spec/fixtures/test_config.yml

@@ -1,4 +1,7 @@
 ---
+:health_check:
+  :address: "127.0.0.1"
+  :port: 8080
 :mailboxes:
   -
     :email: "[email protected]"

spec/lib/cli_spec.rb

@@ -5,14 +5,14 @@
   let!(:configuration) {MailRoom::Configuration.new({config_path: config_path})}
   let(:coordinator) {stub(run: true, quit: true)}
   let(:configuration_args) { anything }
-  let(:coordinator_args) { anything }
+  let(:coordinator_args) { [anything, anything] }
 
   describe '.new' do
     let(:args) {["-c", "a path"]}
 
     before :each do
       MailRoom::Configuration.expects(:new).with(configuration_args).returns(configuration)
-      MailRoom::Coordinator.stubs(:new).with(coordinator_args).returns(coordinator)
+      MailRoom::Coordinator.stubs(:new).with(*coordinator_args).returns(coordinator)
     end
 
     context 'with configuration args' do
@@ -27,7 +27,7 @@
 
     context 'with coordinator args' do
       let(:coordinator_args) do
-        configuration.mailboxes
+        [configuration.mailboxes, anything]
       end
 
       it 'creates a new coordinator with configuration' do

spec/lib/configuration_spec.rb

@@ -3,7 +3,7 @@
 describe MailRoom::Configuration do
   let(:config_path) {File.expand_path('../fixtures/test_config.yml', File.dirname(__FILE__))}
 
-  describe 'set_mailboxes' do
+  describe '#initalize' do
     context 'with config_path' do
       let(:configuration) { MailRoom::Configuration.new(config_path: config_path) }
 
@@ -12,6 +12,10 @@
 
         expect(configuration.mailboxes).to eq(['mailbox1', 'mailbox2'])
       end
+
+      it 'parses health check' do
+        expect(configuration.health_check).to be_a(MailRoom::HealthCheck)
+      end
     end
 
     context 'without config_path' do
@@ -23,6 +27,10 @@
 
         expect(configuration.mailboxes).to eq([])
       end
+
+      it 'sets the health check to nil' do
+        expect(configuration.health_check).to be_nil
+      end
     end
   end
 end

spec/lib/coordinator_spec.rb

@@ -15,22 +15,36 @@
       coordinator = MailRoom::Coordinator.new([])
       expect(coordinator.watchers).to eq([])
     end
+
+    it 'sets the health check' do
+      health_check = MailRoom::HealthCheck.new({ address: '127.0.0.1', port: 8080})
+      coordinator = MailRoom::Coordinator.new([], health_check)
+
+      expect(coordinator.health_check).to eq(health_check)
+    end
   end
 
   describe '#run' do
     it 'runs each watcher' do
       watcher = stub
       watcher.stubs(:run)
       watcher.stubs(:quit)
+
+      health_check = stub
+      health_check.stubs(:run)
+      health_check.stubs(:quit)
+
       MailRoom::MailboxWatcher.stubs(:new).returns(watcher)
-      coordinator = MailRoom::Coordinator.new(['mailbox1'])
+      coordinator = MailRoom::Coordinator.new(['mailbox1'], health_check)
       coordinator.stubs(:sleep_while_running)
       watcher.expects(:run)
       watcher.expects(:quit)
+      health_check.expects(:run)
+      health_check.expects(:quit)
 
       coordinator.run
     end
-    
+
     it 'should go to sleep after running watchers' do
       coordinator = MailRoom::Coordinator.new([])
       coordinator.stubs(:running=)

spec/lib/health_check_spec.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe MailRoom::HealthCheck do
+  let(:address) { '127.0.0.1' }
+  let(:port) { 8000 }
+  let(:params) { { address: address, port: port } }
+  subject { described_class.new(params) }
+
+  describe '#initialize' do
+    context 'with valid parameters' do
+      it 'validates successfully' do
+        expect(subject).to be_a(described_class)
+      end
+    end
+
+    context 'with invalid address' do
+      let(:address) { nil }
+
+      it 'raises an error' do
+        expect { subject }.to raise_error('No health check address specified')
+      end
+    end
+
+    context 'with invalid port' do
+      let(:port) { nil }
+
+      it 'raises an error' do
+        expect { subject }.to raise_error('Health check port 0 is invalid')
+      end
+    end
+  end
+
+  describe '#run' do
+    it 'sets running to true' do
+      server = stub(start: true)
+      subject.stubs(:create_server).returns(server)
+
+      subject.run
+
+      expect(subject.running).to be true
+    end
+  end
+
+  describe '#quit' do
+    it 'sets running to false' do
+      server = stub(start: true, shutdown: true)
+      subject.stubs(:create_server).returns(server)
+
+      subject.run
+      subject.quit
+
+      expect(subject.running).to be false
+    end
+  end
+end